From 279c4b78cc210c3edea549fe61d9bd519ff86402 Mon Sep 17 00:00:00 2001 From: junjun Date: Tue, 20 May 2025 11:55:28 +0800 Subject: [PATCH] feat(datasource): Support PostgreSQL --- backend/apps/datasource/crud/datasource.py | 43 ++++++++------- backend/apps/db/db.py | 61 ++++++++++++++++++++-- backend/pyproject.toml | 3 +- frontend/src/assets/svg/ds/pg-ds.svg | 13 +++++ frontend/src/views/ds/form.vue | 8 +-- frontend/src/views/ds/index.vue | 5 +- frontend/src/views/ds/js/ds-type.ts | 18 +++++++ 7 files changed, 119 insertions(+), 32 deletions(-) create mode 100644 frontend/src/assets/svg/ds/pg-ds.svg create mode 100644 frontend/src/views/ds/js/ds-type.ts diff --git a/backend/apps/datasource/crud/datasource.py b/backend/apps/datasource/crud/datasource.py index 70b94b3f..d164b6a5 100644 --- a/backend/apps/datasource/crud/datasource.py +++ b/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 -from sqlalchemy import text, and_ -from common.utils.utils import deepcopy_ignore_extra +import json from typing import List -from ..crud.table import delete_table_by_ds_id, update_table + +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 ..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) diff --git a/backend/apps/db/db.py b/backend/apps/db/db.py index 4f5a96d1..722ffa4b 100644 --- a/backend/apps/db/db.py +++ b/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() diff --git a/backend/pyproject.toml b/backend/pyproject.toml index d165de11..7f11a770 100644 --- a/backend/pyproject.toml +++ b/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" diff --git a/frontend/src/assets/svg/ds/pg-ds.svg b/frontend/src/assets/svg/ds/pg-ds.svg new file mode 100644 index 00000000..9a5bb6be --- /dev/null +++ b/frontend/src/assets/svg/ds/pg-ds.svg @@ -0,0 +1,13 @@ + + + + + + + + + + + + + diff --git a/frontend/src/views/ds/form.vue b/frontend/src/views/ds/form.vue index 009b485c..c5d6e1df 100644 --- a/frontend/src/views/ds/form.vue +++ b/frontend/src/views/ds/form.vue @@ -65,7 +65,7 @@ - + Get Schema @@ -101,6 +101,7 @@ import { ElMessage } from 'element-plus' import type { FormInstance, FormRules } from 'element-plus' import { Plus } from '@element-plus/icons-vue' import { useCache } from '@/utils/useCache' +import { dsType, haveSchema} from '@/views/ds/js/ds-type' const { wsCache } = useCache() const dsFormRef = ref() @@ -123,11 +124,6 @@ const rules = reactive({ }) const dialogVisible = ref(false) -const dsType = [ - {label:"Excel/CSV", value:"excel"}, - {label:"MySQL", value:"mysql"}, - {label:"SQL Server", value:"sqlServer"} -] const form = ref({ name:'', description:'', diff --git a/frontend/src/views/ds/index.vue b/frontend/src/views/ds/index.vue index e99f6e99..b05ed69e 100644 --- a/frontend/src/views/ds/index.vue +++ b/frontend/src/views/ds/index.vue @@ -28,11 +28,12 @@ +
{{ ds.name }}
-
{{ ds.type }}
+
{{ dsRelation[ds.type] }}
{{ ds.description }}
{{ datetimeFormat(ds.create_time) }}
@@ -54,6 +55,7 @@ import IconOpeEdit from '@/assets/svg/operate/ope-edit.svg' import IconOpeDelete from '@/assets/svg/operate/ope-delete.svg' import MysqlDs from '@/assets/svg/ds/mysql-ds.svg' import SQLServerDs from '@/assets/svg/ds/sqlServer-ds.svg' +import PgDs from '@/assets/svg/ds/pg-ds.svg' import { Search, List, CreditCard } from '@element-plus/icons-vue' import { ref, onMounted } from 'vue' import DsForm from './form.vue' @@ -61,6 +63,7 @@ import { datasourceApi } from '@/api/datasource' import { datetimeFormat } from '@/utils/utils' import { ElMessageBox } from 'element-plus' import { useRouter } from 'vue-router' +import { dsRelation } from '@/views/ds/js/ds-type' const searchValue = ref('') const dsForm = ref() diff --git a/frontend/src/views/ds/js/ds-type.ts b/frontend/src/views/ds/js/ds-type.ts new file mode 100644 index 00000000..abfc13ab --- /dev/null +++ b/frontend/src/views/ds/js/ds-type.ts @@ -0,0 +1,18 @@ +export const dsType = [ + {label:"MySQL", value:"mysql"}, + {label:"Microsoft SQL Server", value:"sqlServer"}, + {label:"PostgreSQL", value:"pg"}, + {label:"Excel/CSV", value:"excel"}, +] + +export const haveSchema = [ + 'sqlServer', + 'pg' +] + +export const dsRelation: any = { + "mysql": "MySQL", + "sqlServer": "Microsoft SQL Server", + "pg": "PostgreSQL", + "excel": "Excel/CSV", +} \ No newline at end of file