klinsmaya/SQLBot
1
0
Fork 0
mirror of https://gh-proxy.org/https://github.com/klinsmaya/SQLBot.git synced 2026-11-08 01:06:59 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: support kingbase datasource

Browse Source
This commit is contained in:
junjun
2025-09-28 14:20:16 +08:00
parent 3dc4b8ece3
commit 330fd2ddfd
4 changed files with 97 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/apps/datasource/crud/datasource.py
+1 -1
View File
@@ -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"""
backend/apps/db/constant.py
+1
View File
@@ -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
backend/apps/db/db.py
+63 -3
View File
@@ -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)
backend/apps/db/db_sql.py
+32 -1
View File
@@ -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