feat: full query of 1/2 type

support query way and point

no relation, also no polygon

src/yuheng/plugin/driver_db_postgresql/__main__.py

@@ -1,3 +1,5 @@
+from typing import Optional
+
 import psycopg
 import os
 import sys
@@ -24,7 +26,7 @@ def get_data(
     query_type=["line"],
     pg_schema="public",
     pg_pre_fix="planet_osm",
-) -> Waifu:
+) -> Optional[Waifu]:
     """
     # full-全量查询
     # batch-批量查询
@@ -42,25 +44,31 @@ def get_data(
         " ".join([item + "=" + config.get(item) for item in config])
     ) as connection:
         with connection.cursor() as cursor:
-            # config flag
-            flag_query_directly = False
-            # generate sql
             if query_mode == "full":
-                if len(query_type) == 1:
+                if len(query_type) == 1 and (
+                    "line" in query_type or "point" in query_type
+                ):
                     sql = f"SELECT * FROM {pg_schema}.{pg_pre_fix}_{query_type[0]}"
+                    cursor.execute(sql)
+                    for record in cursor:
+                        result.append((query_type[0], record))
                 elif (
                     len(query_type) == 2
                     and "line" in query_type
                     and "point" in query_type
                 ):
-                    if flag_query_directly:
-                        sql = f"SELECT * FROM join({pg_schema}.{pg_pre_fix}_line,{pg_schema}.{pg_pre_fix}_point)"
-            # conduct and collect
-            cursor.execute(sql)
-            for record in cursor:
-                result.append(record)
-            connection.commit()
-    print(len(result))
+                    sql_a = f"SELECT * FROM {pg_schema}.{pg_pre_fix}_{query_type[0]}"
+                    cursor.execute(sql_a)
+                    for record in cursor:
+                        result.append((query_type[0], record))
+                    sql_b = f"SELECT * FROM {pg_schema}.{pg_pre_fix}_{query_type[1]}"
+                    cursor.execute(sql_b)
+                    for record in cursor:
+                        result.append((query_type[1], record))
+                else:
+                    print("要么你就想查询line/point以外的表，要么你就输了太多项。目前无法支持。")
+                    return None
+    print("len(result)", "=", len(result))
     return result
 
 

tests/cases/plugin_driver_db_postgresql.py

@@ -9,11 +9,10 @@
 
 class TestPluginDriverDbPostgresql(unittest.TestCase):
     def setUp(self) -> None:
-        pass
-
-    def test_plugin_driver_db_postgresql_full(self):
         # NOTE: conduct this part of postgresql need you deploy a server.
+        pass
 
+    def test_plugin_driver_db_postgresql_full_1type(self):
         from yuheng.plugin.driver_db_postgresql.__main__ import get_data
 
         result = get_data(
@@ -26,7 +25,36 @@ def test_plugin_driver_db_postgresql_full(self):
             query_type=["line"],
         )
         print(len(result))
-        # assert len(result) == 0
+        # assert len(result) == 114 or 514
+
+    def test_plugin_driver_db_postgresql_full_2type(self):
+        from yuheng.plugin.driver_db_postgresql.__main__ import get_data
+
+        result = get_data(
+            connection_dbname="osm2pgsql",
+            connection_user="postgres",
+            connection_password="12345678",
+            connection_host="localhost",
+            connection_port="5432",
+            query_mode="full",
+            query_type=["line", "point"],
+        )
+        print(len(result))
+        # assert len(result) == 114 + 514
+
+    def test_plugin_driver_db_postgresql_full_invalidtype(self):
+        from yuheng.plugin.driver_db_postgresql.__main__ import get_data
+
+        result = get_data(
+            connection_dbname="osm2pgsql",
+            connection_user="postgres",
+            connection_password="12345678",
+            connection_host="localhost",
+            connection_port="5432",
+            query_mode="full",
+            query_type=["line", "point", "polygon"],
+        )
+        assert result == None
 
 
 if __name__ == "__main__":