From 330fd2ddfdefd3b664f2b61680d91a074ce2e3e5 Mon Sep 17 00:00:00 2001 From: junjun Date: Sun, 28 Sep 2025 14:20:16 +0800 Subject: [PATCH] feat: support kingbase datasource --- backend/apps/datasource/crud/datasource.py | 2 +- backend/apps/db/constant.py | 1 + backend/apps/db/db.py | 66 +++++++++++++++++++++- backend/apps/db/db_sql.py | 33 ++++++++++- 4 files changed, 97 insertions(+), 5 deletions(-) diff --git a/backend/apps/datasource/crud/datasource.py b/backend/apps/datasource/crud/datasource.py index 4380b3cf..c541a152 100644 --- a/backend/apps/datasource/crud/datasource.py +++ b/backend/apps/datasource/crud/datasource.py @@ -279,7 +279,7 @@ def preview(session: SessionDep, current_user: CurrentUser, id: int, data: Table sql = f"""SELECT TOP 100 [{"], [".join(fields)}] FROM [{conf.dbSchema}].[{data.table.table_name}] {where} """ - elif ds.type == "pg" or ds.type == "excel" or ds.type == "redshift": + elif ds.type == "pg" or ds.type == "excel" or ds.type == "redshift" or ds.type == "kingbase": sql = f"""SELECT "{'", "'.join(fields)}" FROM "{conf.dbSchema}"."{data.table.table_name}" {where} LIMIT 100""" diff --git a/backend/apps/db/constant.py b/backend/apps/db/constant.py index 4194d8b7..28d1ad0b 100644 --- a/backend/apps/db/constant.py +++ b/backend/apps/db/constant.py @@ -23,6 +23,7 @@ class DB(Enum): doris = ('doris', 'Apache Doris', '`', '`', ConnectType.py_driver) redshift = ('redshift', 'AWS Redshift', '"', '"', ConnectType.py_driver) es = ('es', 'Elasticsearch', '"', '"', ConnectType.py_driver) + kingbase = ('kingbase', 'Kingbase', '"', '"', ConnectType.py_driver) def __init__(self, type, db_name, prefix, suffix, connect_type: ConnectType): self.type = type diff --git a/backend/apps/db/db.py b/backend/apps/db/db.py index 1a22dc35..c428a576 100644 --- a/backend/apps/db/db.py +++ b/backend/apps/db/db.py @@ -5,8 +5,8 @@ import urllib.parse from decimal import Decimal from typing import Optional +import psycopg2 import pymssql - from apps.db.db_sql import get_table_sql, get_field_sql, get_version_sql from common.error import ParseSQLResultError @@ -191,6 +191,20 @@ def check_connection(trans: Optional[Trans], ds: CoreDatasource | AssistantOutDs if is_raise: raise HTTPException(status_code=500, detail=trans('i18n_ds_invalid') + f': {e.args}') return False + elif ds.type == 'kingbase': + with psycopg2.connect(host=conf.host, port=conf.port, database=conf.database, + user=conf.username, + password=conf.password, + connect_timeout=10, **extra_config_dict) as conn, conn.cursor() as cursor: + try: + cursor.execute('select 1') + SQLBotLogUtil.info("success") + return True + except Exception as e: + SQLBotLogUtil.error(f"Datasource {ds.id} connection failed: {e}") + if is_raise: + raise HTTPException(status_code=500, detail=trans('i18n_ds_invalid') + f': {e.args}') + return False elif ds.type == 'es': es_conn = get_es_connect(conf) if es_conn.ping(): @@ -269,8 +283,7 @@ def get_schema(ds: CoreDatasource): if ds.type == "sqlServer": sql = """select name from sys.schemas""" elif ds.type == "pg" or ds.type == "excel": - sql = """SELECT nspname - FROM pg_namespace""" + sql = """SELECT nspname FROM pg_namespace""" elif ds.type == "oracle": sql = """select * from all_users""" with session.execute(text(sql)) as result: @@ -294,6 +307,15 @@ def get_schema(ds: CoreDatasource): res = cursor.fetchall() res_list = [item[0] for item in res] return res_list + elif ds.type == 'kingbase': + with psycopg2.connect(host=conf.host, port=conf.port, database=conf.database, user=conf.username, + password=conf.password, + options=f"-c statement_timeout={conf.timeout * 1000}", + **extra_config_dict) as conn, conn.cursor() as cursor: + cursor.execute("""SELECT nspname FROM pg_namespace""") + res = cursor.fetchall() + res_list = [item[0] for item in res] + return res_list def get_tables(ds: CoreDatasource): @@ -331,6 +353,15 @@ def get_tables(ds: CoreDatasource): res = cursor.fetchall() res_list = [TableSchema(*item) for item in res] return res_list + elif ds.type == 'kingbase': + with psycopg2.connect(host=conf.host, port=conf.port, database=conf.database, user=conf.username, + password=conf.password, + options=f"-c statement_timeout={conf.timeout * 1000}", + **extra_config_dict) as conn, conn.cursor() as cursor: + cursor.execute(sql.format(sql_param)) + res = cursor.fetchall() + res_list = [TableSchema(*item) for item in res] + return res_list elif ds.type == 'es': res = get_es_index(conf) res_list = [TableSchema(*item) for item in res] @@ -372,6 +403,15 @@ def get_fields(ds: CoreDatasource, table_name: str = None): res = cursor.fetchall() res_list = [ColumnSchema(*item) for item in res] return res_list + elif ds.type == 'kingbase': + with psycopg2.connect(host=conf.host, port=conf.port, database=conf.database, user=conf.username, + password=conf.password, + options=f"-c statement_timeout={conf.timeout * 1000}", + **extra_config_dict) as conn, conn.cursor() as cursor: + cursor.execute(sql.format(p1, p2)) + res = cursor.fetchall() + res_list = [ColumnSchema(*item) for item in res] + return res_list elif ds.type == 'es': res = get_es_fields(conf, table_name) res_list = [ColumnSchema(*item) for item in res] @@ -457,6 +497,26 @@ def exec_sql(ds: CoreDatasource | AssistantOutDsSchema, sql: str, origin_column= "sql": bytes.decode(base64.b64encode(bytes(sql, 'utf-8')))} except Exception as ex: raise ParseSQLResultError(str(ex)) + elif ds.type == 'kingbase': + with psycopg2.connect(host=conf.host, port=conf.port, database=conf.database, user=conf.username, + password=conf.password, + options=f"-c statement_timeout={conf.timeout * 1000}", + **extra_config_dict) as conn, conn.cursor() as cursor: + try: + cursor.execute(sql) + res = cursor.fetchall() + columns = [field[0] for field in cursor.description] if origin_column else [field[0].lower() for + field in + cursor.description] + result_list = [ + {str(columns[i]): float(value) if isinstance(value, Decimal) else value for i, value in + enumerate(tuple_item)} + for tuple_item in res + ] + return {"fields": columns, "data": result_list, + "sql": bytes.decode(base64.b64encode(bytes(sql, 'utf-8')))} + except Exception as ex: + raise ParseSQLResultError(str(ex)) elif ds.type == 'es': try: res, columns = get_es_data_by_http(conf, sql) diff --git a/backend/apps/db/db_sql.py b/backend/apps/db/db_sql.py index 66616b70..bced9a22 100644 --- a/backend/apps/db/db_sql.py +++ b/backend/apps/db/db_sql.py @@ -12,7 +12,7 @@ def get_version_sql(ds: CoreDatasource, conf: DatasourceConf): return """ select SERVERPROPERTY('ProductVersion') """ - elif ds.type == "pg" or ds.type == "excel": + elif ds.type == "pg" or ds.type == "kingbase" or ds.type == "excel": return """ SELECT current_setting('server_version') """ @@ -140,6 +140,21 @@ def get_table_sql(ds: CoreDatasource, conf: DatasourceConf, db_version: str = '' WHERE TABLE_SCHEMA = %s """, conf.database + elif ds.type == "kingbase": + return """ + 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 = '{0}' + AND c.relkind IN ('r', 'v', 'p', 'm') + AND c.relname NOT LIKE 'pg_%' + AND c.relname NOT LIKE 'sql_%' + ORDER BY c.relname \ + """, conf.dbSchema elif ds.type == "es": return "", None @@ -275,5 +290,21 @@ def get_field_sql(ds: CoreDatasource, conf: DatasourceConf, table_name: str = No """ sql2 = " AND TABLE_NAME = %s" if table_name is not None and table_name != "" else "" return sql1 + sql2, conf.database, table_name + elif ds.type == "kingbase": + 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 = '{0}' + AND a.attnum > 0 + AND NOT a.attisdropped \ + """ + sql2 = " AND c.relname = '{1}'" if table_name is not None and table_name != "" else "" + return sql1 + sql2, conf.dbSchema, table_name elif ds.type == "es": return "", None, None