feat(datasource): Support PostgreSQL

Author: XiaJunjie2020

backend/apps/datasource/crud/datasource.py

@@ -1,13 +1,18 @@
-from sqlmodel import select
-from ..models.datasource import CoreDatasource, CreateDatasource, CoreTable, CoreField, ColumnSchema, EditObj
 import datetime
-from common.core.deps import SessionDep
-from apps.db.db import get_session, get_tables, get_fields, exec_sql
+import json
+from typing import List
+
 from sqlalchemy import text, and_
+from sqlmodel import select
+
+from apps.datasource.utils.utils import aes_decrypt
+from apps.db.db import get_engine, get_tables, get_fields, exec_sql
+from common.core.deps import SessionDep
 from common.utils.utils import deepcopy_ignore_extra
-from typing import List
-from ..crud.table import delete_table_by_ds_id, update_table
 from ..crud.field import delete_field_by_ds_id, update_field
+from ..crud.table import delete_table_by_ds_id, update_table
+from ..models.datasource import CoreDatasource, CreateDatasource, CoreTable, CoreField, ColumnSchema, EditObj, \
+    DatasourceConf
 
 
 def get_datasource_list(session: SessionDep) -> CoreDatasource:
@@ -17,17 +22,14 @@ def get_datasource_list(session: SessionDep) -> CoreDatasource:
 
 
 def check_status(session: SessionDep, ds: CoreDatasource):
-    conn = get_session(ds)
+    conn = get_engine(ds)
     try:
-        conn.execute(text("SELECT 1")).scalar()
-        print("success")
-        return True
+        with conn.connect() as connection:
+            print("success")
+            return True
     except Exception as e:
         print("Fail:", e)
         return False
-    finally:
-        conn.close()
-    return False
 
 
 def create_ds(session: SessionDep, create_ds: CreateDatasource):
@@ -175,13 +177,14 @@ def update_table_and_fields(session: SessionDep, data: EditObj):
 
 def preview(session: SessionDep, id: int, data: EditObj):
     ds = session.query(CoreDatasource).filter(CoreDatasource.id == id).first()
+    conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration)))
     sql: str = ""
     if ds.type == "mysql":
-        sql = f"""SELECT {", ".join([f.field_name for f in data.fields if f.checked])} FROM {data.table.table_name} LIMIT 100"""
+        sql = f"""SELECT `{"`, `".join([f.field_name for f in data.fields if f.checked])}` FROM `{data.table.table_name}` LIMIT 100"""
     elif ds.type == "sqlServer":
-        sql = f"""
-            SELECT {", ".join([f.field_name for f in data.fields if f.checked])} FROM {data.table.table_name} 
-            ORDER BY {data.fields[0].field_name}
-            OFFSET 0 ROWS FETCH NEXT 100 ROWS ONLY
-        """
+        sql = f"""SELECT [{"], [".join([f.field_name for f in data.fields if f.checked])}] FROM [{conf.dbSchema}].[{data.table.table_name}]
+            ORDER BY [{data.fields[0].field_name}]
+            OFFSET 0 ROWS FETCH NEXT 100 ROWS ONLY"""
+    elif ds.type == "pg":
+        sql = f"""SELECT "{'", "'.join([f.field_name for f in data.fields if f.checked])}" FROM "{conf.dbSchema}"."{data.table.table_name}" LIMIT 100"""
     return exec_sql(ds, sql)

backend/apps/db/db.py

@@ -1,8 +1,9 @@
+import base64
 import json
 import urllib.parse
 from typing import Any
 
-from sqlalchemy import create_engine, text, Result
+from sqlalchemy import create_engine, text, Result, Engine
 from sqlalchemy.orm import sessionmaker
 
 from apps.datasource.models.datasource import DatasourceConf, CoreDatasource, TableSchema, ColumnSchema
@@ -22,13 +23,28 @@ def get_uri(ds: CoreDatasource):
             db_url = f"mssql+pymssql://{urllib.parse.quote(conf.username)}:{urllib.parse.quote(conf.password)}@{conf.host}:{conf.port}/{urllib.parse.quote(conf.database)}?{urllib.parse.quote(conf.extraJdbc)}"
         else:
             db_url = f"mssql+pymssql://{urllib.parse.quote(conf.username)}:{urllib.parse.quote(conf.password)}@{conf.host}:{conf.port}/{urllib.parse.quote(conf.database)}"
+    elif ds.type == "pg":
+        if conf.extraJdbc is not None and conf.extraJdbc != '':
+            db_url = f"postgresql+psycopg2://{urllib.parse.quote(conf.username)}:{urllib.parse.quote(conf.password)}@{conf.host}:{conf.port}/{urllib.parse.quote(conf.database)}?{urllib.parse.quote(conf.extraJdbc)}"
+        else:
+            db_url = f"postgresql+psycopg2://{urllib.parse.quote(conf.username)}:{urllib.parse.quote(conf.password)}@{conf.host}:{conf.port}/{urllib.parse.quote(conf.database)}"
     else:
         raise 'The datasource type not support.'
     return db_url
 
 
+def get_engine(ds: CoreDatasource) -> Engine:
+    conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration)))
+    if ds.type == "pg" and (conf.dbSchema is not None and conf.dbSchema != ""):
+        engine = create_engine(get_uri(ds), connect_args={"options": f"-c search_path={conf.dbSchema}"},
+                               pool_timeout=30, pool_size=20, max_overflow=10)
+    else:
+        engine = create_engine(get_uri(ds), pool_timeout=30, pool_size=20, max_overflow=10)
+    return engine
+
+
 def get_session(ds: CoreDatasource):
-    engine = create_engine(get_uri(ds))
+    engine = get_engine(ds)
     session_maker = sessionmaker(bind=engine)
     session = session_maker()
     return session
@@ -66,6 +82,24 @@ def get_tables(ds: CoreDatasource):
                         t.TABLE_TYPE IN ('BASE TABLE', 'VIEW')
                         AND t.TABLE_SCHEMA = '{conf.dbSchema}';
                     """
+        elif ds.type == "pg":
+            sql = """
+                    SELECT 
+                        c.relname AS TABLE_NAME,
+                        COALESCE(d.description, obj_description(c.oid)) AS TABLE_COMMENT
+                    FROM 
+                        pg_class c
+                    LEFT JOIN 
+                        pg_namespace n ON n.oid = c.relnamespace
+                    LEFT JOIN 
+                        pg_description d ON d.objoid = c.oid AND d.objsubid = 0
+                    WHERE 
+                        n.nspname = current_schema() 
+                        AND c.relkind IN ('r', 'v')  
+                        AND c.relname NOT LIKE 'pg_%'
+                        AND c.relname NOT LIKE 'sql_%'
+                    ORDER BY c.relname;
+                    """
 
         result = session.execute(text(sql))
         res = result.fetchall()
@@ -113,7 +147,26 @@ def get_fields(ds: CoreDatasource, table_name: str = None):
                     WHERE 
                         C.TABLE_SCHEMA = '{conf.dbSchema}'
                     """
-            sql2 = f"AND C.TABLE_NAME = '{table_name}';" if table_name is not None and table_name != "" else ";"
+            sql2 = f" AND C.TABLE_NAME = '{table_name}';" if table_name is not None and table_name != "" else ";"
+            sql = sql1 + sql2
+        elif ds.type == "pg":
+            sql1 = """
+                    SELECT 
+                        a.attname AS COLUMN_NAME,
+                        pg_catalog.format_type(a.atttypid, a.atttypmod) AS DATA_TYPE,
+                        col_description(c.oid, a.attnum) AS COLUMN_COMMENT
+                    FROM 
+                        pg_catalog.pg_attribute a
+                    JOIN 
+                        pg_catalog.pg_class c ON a.attrelid = c.oid
+                    JOIN 
+                        pg_catalog.pg_namespace n ON n.oid = c.relnamespace  
+                    WHERE 
+                        n.nspname = current_schema()  
+                        AND a.attnum > 0              
+                        AND NOT a.attisdropped
+                    """
+            sql2 = f" AND c.relname = '{table_name}';" if table_name is not None and table_name != "" else ";"
             sql = sql1 + sql2
 
         result = session.execute(text(sql))
@@ -137,7 +190,7 @@ def exec_sql(ds: CoreDatasource, sql: str):
             {columns[i]: value for i, value in enumerate(tuple_item)}
             for tuple_item in res
         ]
-        return {"fields": columns, "data": result_list}
+        return {"fields": columns, "data": result_list, "sql": base64.b64encode(bytes(sql, 'utf-8'))}
     finally:
         if result is not None:
             result.close()

backend/pyproject.toml

@@ -31,7 +31,8 @@ dependencies = [
     "llama_index>=0.12.35",
     "pymssql (>=2.3.4,<3.0.0)",
     "pandas (>=2.2.3,<3.0.0)",
-    "openpyxl (>=3.1.5,<4.0.0)"
+    "openpyxl (>=3.1.5,<4.0.0)",
+    "psycopg2 (>=2.9.10,<3.0.0)"
 ]
 [[tool.uv.index]]
 url = "https://pypi.tuna.tsinghua.edu.cn/simple"