mirror of
https://gh-proxy.org/https://github.com/klinsmaya/SQLBot.git
synced 2026-11-08 01:06:59 +08:00
773 lines
39 KiB
Python
773 lines
39 KiB
Python
import base64
|
|
import json
|
|
import os
|
|
import platform
|
|
import re
|
|
import urllib.parse
|
|
from datetime import datetime, date, time, timedelta
|
|
from decimal import Decimal
|
|
from typing import Optional, List
|
|
|
|
import oracledb
|
|
import psycopg2
|
|
import pymssql
|
|
|
|
from apps.db.db_sql import get_table_sql, get_field_sql, get_version_sql
|
|
from common.error import ParseSQLResultError
|
|
|
|
if platform.system() != "Darwin":
|
|
import dmPython
|
|
import pymysql
|
|
import redshift_connector
|
|
from sqlalchemy import create_engine, text, Engine
|
|
from sqlalchemy.orm import sessionmaker
|
|
|
|
from apps.datasource.models.datasource import DatasourceConf, CoreDatasource, TableSchema, ColumnSchema
|
|
from apps.datasource.utils.utils import aes_decrypt
|
|
from apps.db.constant import DB, ConnectType
|
|
from apps.db.engine import get_engine_config
|
|
from apps.system.crud.assistant import get_out_ds_conf
|
|
from apps.system.schemas.system_schema import AssistantOutDsSchema
|
|
from common.core.deps import Trans
|
|
from common.utils.utils import SQLBotLogUtil, equals_ignore_case
|
|
from fastapi import HTTPException
|
|
from apps.db.es_engine import get_es_connect, get_es_index, get_es_fields, get_es_data_by_http
|
|
from common.core.config import settings
|
|
import sqlglot
|
|
from sqlglot import expressions as exp
|
|
from sqlalchemy.pool import NullPool
|
|
from pyhive import hive
|
|
|
|
try:
|
|
if os.path.exists(settings.ORACLE_CLIENT_PATH):
|
|
oracledb.init_oracle_client(
|
|
lib_dir=settings.ORACLE_CLIENT_PATH
|
|
)
|
|
SQLBotLogUtil.info("init oracle client success, use thick mode")
|
|
else:
|
|
SQLBotLogUtil.info("init oracle client failed, because not found oracle client, use thin mode")
|
|
except Exception as e:
|
|
SQLBotLogUtil.error("init oracle client failed, check your client is installed, use thin mode")
|
|
|
|
|
|
def get_uri(ds: CoreDatasource) -> str:
|
|
conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration))) if not equals_ignore_case(ds.type,
|
|
"excel") else get_engine_config()
|
|
return get_uri_from_config(ds.type, conf)
|
|
|
|
|
|
def get_uri_from_config(type: str, conf: DatasourceConf) -> str:
|
|
db_url: str
|
|
if equals_ignore_case(type, "mysql"):
|
|
checkParams(conf.extraJdbc, DB.mysql.illegalParams)
|
|
if conf.extraJdbc is not None and conf.extraJdbc != '':
|
|
db_url = f"mysql+pymysql://{urllib.parse.quote(conf.username)}:{urllib.parse.quote(conf.password)}@{conf.host}:{conf.port}/{conf.database}?{conf.extraJdbc}"
|
|
else:
|
|
db_url = f"mysql+pymysql://{urllib.parse.quote(conf.username)}:{urllib.parse.quote(conf.password)}@{conf.host}:{conf.port}/{conf.database}"
|
|
elif equals_ignore_case(type, "sqlServer"):
|
|
if conf.extraJdbc is not None and conf.extraJdbc != '':
|
|
db_url = f"mssql+pymssql://{urllib.parse.quote(conf.username)}:{urllib.parse.quote(conf.password)}@{conf.host}:{conf.port}/{conf.database}?{conf.extraJdbc}"
|
|
else:
|
|
db_url = f"mssql+pymssql://{urllib.parse.quote(conf.username)}:{urllib.parse.quote(conf.password)}@{conf.host}:{conf.port}/{conf.database}"
|
|
elif equals_ignore_case(type, "pg", "excel"):
|
|
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}/{conf.database}?{conf.extraJdbc}"
|
|
else:
|
|
db_url = f"postgresql+psycopg2://{urllib.parse.quote(conf.username)}:{urllib.parse.quote(conf.password)}@{conf.host}:{conf.port}/{conf.database}"
|
|
elif equals_ignore_case(type, "oracle"):
|
|
if equals_ignore_case(conf.mode, "service_name", "serviceName"):
|
|
if conf.extraJdbc is not None and conf.extraJdbc != '':
|
|
db_url = f"oracle+oracledb://{urllib.parse.quote(conf.username)}:{urllib.parse.quote(conf.password)}@{conf.host}:{conf.port}?service_name={conf.database}&{conf.extraJdbc}"
|
|
else:
|
|
db_url = f"oracle+oracledb://{urllib.parse.quote(conf.username)}:{urllib.parse.quote(conf.password)}@{conf.host}:{conf.port}?service_name={conf.database}"
|
|
else:
|
|
if conf.extraJdbc is not None and conf.extraJdbc != '':
|
|
db_url = f"oracle+oracledb://{urllib.parse.quote(conf.username)}:{urllib.parse.quote(conf.password)}@{conf.host}:{conf.port}/{conf.database}?{conf.extraJdbc}"
|
|
else:
|
|
db_url = f"oracle+oracledb://{urllib.parse.quote(conf.username)}:{urllib.parse.quote(conf.password)}@{conf.host}:{conf.port}/{conf.database}"
|
|
elif equals_ignore_case(type, "ck"):
|
|
if conf.extraJdbc is not None and conf.extraJdbc != '':
|
|
db_url = f"clickhouse+http://{urllib.parse.quote(conf.username)}:{urllib.parse.quote(conf.password)}@{conf.host}:{conf.port}/{conf.database}?{conf.extraJdbc}"
|
|
else:
|
|
db_url = f"clickhouse+http://{urllib.parse.quote(conf.username)}:{urllib.parse.quote(conf.password)}@{conf.host}:{conf.port}/{conf.database}"
|
|
else:
|
|
raise 'The datasource type not support.'
|
|
return db_url
|
|
|
|
|
|
def get_extra_config(conf: DatasourceConf):
|
|
config_dict = {}
|
|
if conf.extraJdbc:
|
|
config_arr = conf.extraJdbc.split("&")
|
|
for config in config_arr:
|
|
kv = config.split("=")
|
|
if len(kv) == 2 and kv[0] and kv[1]:
|
|
config_dict[kv[0]] = kv[1]
|
|
else:
|
|
raise Exception(f'param: {config} is error')
|
|
return config_dict
|
|
|
|
|
|
def get_origin_connect(type: str, conf: DatasourceConf):
|
|
extra_config_dict = get_extra_config(conf)
|
|
if equals_ignore_case(type, "sqlServer"):
|
|
# none or true, set tds_version = 7.0
|
|
if conf.lowVersion is None or conf.lowVersion:
|
|
return pymssql.connect(
|
|
server=conf.host,
|
|
port=str(conf.port),
|
|
user=conf.username,
|
|
password=conf.password,
|
|
database=conf.database,
|
|
timeout=conf.timeout,
|
|
tds_version='7.0', # options: '4.2', '7.0', '8.0' ...,
|
|
**extra_config_dict
|
|
)
|
|
else:
|
|
return pymssql.connect(
|
|
server=conf.host,
|
|
port=str(conf.port),
|
|
user=conf.username,
|
|
password=conf.password,
|
|
database=conf.database,
|
|
timeout=conf.timeout,
|
|
**extra_config_dict
|
|
)
|
|
|
|
|
|
# use sqlalchemy
|
|
def get_engine(ds: CoreDatasource, timeout: int = 0) -> Engine:
|
|
conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration))) if not equals_ignore_case(ds.type,
|
|
"excel") else get_engine_config()
|
|
if conf.timeout is None:
|
|
conf.timeout = timeout
|
|
if timeout > 0:
|
|
conf.timeout = timeout
|
|
|
|
if equals_ignore_case(ds.type, "pg"):
|
|
if conf.dbSchema is not None and conf.dbSchema != "":
|
|
engine = create_engine(get_uri(ds),
|
|
connect_args={"options": f"-c search_path={urllib.parse.quote(conf.dbSchema)}",
|
|
"connect_timeout": conf.timeout}, poolclass=NullPool)
|
|
else:
|
|
engine = create_engine(get_uri(ds), connect_args={"connect_timeout": conf.timeout}, poolclass=NullPool)
|
|
elif equals_ignore_case(ds.type, 'sqlServer'):
|
|
engine = create_engine('mssql+pymssql://', creator=lambda: get_origin_connect(ds.type, conf),
|
|
poolclass=NullPool)
|
|
elif equals_ignore_case(ds.type, 'oracle'):
|
|
engine = create_engine(get_uri(ds), poolclass=NullPool)
|
|
elif equals_ignore_case(ds.type, 'mysql'): # mysql
|
|
ssl_mode = {"require": True} if conf.ssl else None
|
|
engine = create_engine(get_uri(ds), connect_args={"connect_timeout": conf.timeout, "ssl": ssl_mode},
|
|
poolclass=NullPool)
|
|
else: # ck
|
|
engine = create_engine(get_uri(ds), connect_args={"connect_timeout": conf.timeout}, poolclass=NullPool)
|
|
return engine
|
|
|
|
|
|
def get_session(ds: CoreDatasource | AssistantOutDsSchema):
|
|
# engine = get_engine(ds) if isinstance(ds, CoreDatasource) else get_ds_engine(ds)
|
|
if isinstance(ds, AssistantOutDsSchema):
|
|
out_conf = get_out_ds_conf(ds, 30)
|
|
ds.configuration = out_conf
|
|
|
|
engine = get_engine(ds)
|
|
session_maker = sessionmaker(bind=engine)
|
|
session = session_maker()
|
|
return session
|
|
|
|
|
|
def check_connection(trans: Optional[Trans], ds: CoreDatasource | AssistantOutDsSchema, is_raise: bool = False):
|
|
if isinstance(ds, AssistantOutDsSchema):
|
|
out_conf = get_out_ds_conf(ds, 10)
|
|
ds.configuration = out_conf
|
|
|
|
db = DB.get_db(ds.type)
|
|
if db.connect_type == ConnectType.sqlalchemy:
|
|
conn = get_engine(ds, 10)
|
|
try:
|
|
with conn.connect() as connection:
|
|
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
|
|
else:
|
|
conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration)))
|
|
extra_config_dict = get_extra_config(conf)
|
|
if equals_ignore_case(ds.type, 'dm'):
|
|
with dmPython.connect(user=conf.username, password=conf.password, server=conf.host,
|
|
port=conf.port, **extra_config_dict) as conn, conn.cursor() as cursor:
|
|
try:
|
|
cursor.execute('select 1', timeout=10).fetchall()
|
|
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 equals_ignore_case(ds.type, 'doris', 'starrocks'):
|
|
ssl_args = {'ssl': {'ssl_mode': 'REQUIRE'}} if conf.ssl else {}
|
|
with pymysql.connect(user=conf.username, passwd=conf.password, host=conf.host,
|
|
port=conf.port, db=conf.database, connect_timeout=10,
|
|
read_timeout=10, **extra_config_dict, **ssl_args) 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 equals_ignore_case(ds.type, 'redshift'):
|
|
with redshift_connector.connect(host=conf.host, port=conf.port, database=conf.database,
|
|
user=conf.username,
|
|
password=conf.password,
|
|
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 equals_ignore_case(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 equals_ignore_case(ds.type, 'hive'):
|
|
with hive.connect(host=conf.host, port=conf.port, username=conf.username,
|
|
database=conf.database, **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 equals_ignore_case(ds.type, 'es'):
|
|
es_conn = get_es_connect(conf)
|
|
if es_conn.ping():
|
|
SQLBotLogUtil.info("success")
|
|
return True
|
|
else:
|
|
SQLBotLogUtil.info("failed")
|
|
return False
|
|
# else:
|
|
# conn = get_ds_engine(ds)
|
|
# try:
|
|
# with conn.connect() as connection:
|
|
# 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
|
|
|
|
return False
|
|
|
|
|
|
def get_version(ds: CoreDatasource | AssistantOutDsSchema):
|
|
version = ''
|
|
if isinstance(ds, CoreDatasource):
|
|
conf = DatasourceConf(
|
|
**json.loads(aes_decrypt(ds.configuration))) if not equals_ignore_case(ds.type,
|
|
"excel") else get_engine_config()
|
|
else:
|
|
conf = DatasourceConf(**json.loads(aes_decrypt(get_out_ds_conf(ds, 10))))
|
|
# if isinstance(ds, AssistantOutDsSchema):
|
|
# conf = DatasourceConf()
|
|
# conf.host = ds.host
|
|
# conf.port = ds.port
|
|
# conf.username = ds.user
|
|
# conf.password = ds.password
|
|
# conf.database = ds.dataBase
|
|
# conf.dbSchema = ds.db_schema
|
|
# conf.timeout = 10
|
|
db = DB.get_db(ds.type)
|
|
sql = get_version_sql(ds, conf)
|
|
if not sql:
|
|
return ''
|
|
try:
|
|
if db.connect_type == ConnectType.sqlalchemy:
|
|
with get_session(ds) as session:
|
|
with session.execute(text(sql)) as result:
|
|
res = result.fetchall()
|
|
version = res[0][0]
|
|
else:
|
|
extra_config_dict = get_extra_config(conf)
|
|
if equals_ignore_case(ds.type, 'dm'):
|
|
with dmPython.connect(user=conf.username, password=conf.password, server=conf.host,
|
|
port=conf.port) as conn, conn.cursor() as cursor:
|
|
cursor.execute(sql, timeout=10, **extra_config_dict)
|
|
res = cursor.fetchall()
|
|
version = res[0][0]
|
|
elif equals_ignore_case(ds.type, 'doris', 'starrocks'):
|
|
ssl_args = {'ssl': {'ssl_mode': 'REQUIRE'}} if conf.ssl else {}
|
|
with pymysql.connect(user=conf.username, passwd=conf.password, host=conf.host,
|
|
port=conf.port, db=conf.database, connect_timeout=10,
|
|
read_timeout=10, **extra_config_dict, **ssl_args) as conn, conn.cursor() as cursor:
|
|
cursor.execute(sql)
|
|
res = cursor.fetchall()
|
|
version = res[0][0]
|
|
elif equals_ignore_case(ds.type, 'redshift', 'es', 'hive'):
|
|
version = ''
|
|
except Exception as e:
|
|
print(e)
|
|
version = ''
|
|
return version.decode() if isinstance(version, bytes) else version
|
|
|
|
|
|
def get_schema(ds: CoreDatasource):
|
|
conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration))) if ds.type != "excel" else get_engine_config()
|
|
db = DB.get_db(ds.type)
|
|
if db.connect_type == ConnectType.sqlalchemy:
|
|
with get_session(ds) as session:
|
|
sql: str = ''
|
|
if equals_ignore_case(ds.type, "sqlServer"):
|
|
sql = """select name
|
|
from sys.schemas"""
|
|
elif equals_ignore_case(ds.type, "pg", "excel"):
|
|
sql = """SELECT nspname
|
|
FROM pg_namespace"""
|
|
elif equals_ignore_case(ds.type, "oracle"):
|
|
sql = """select *
|
|
from all_users"""
|
|
with session.execute(text(sql)) as result:
|
|
res = result.fetchall()
|
|
res_list = [item[0] for item in res]
|
|
return res_list
|
|
else:
|
|
extra_config_dict = get_extra_config(conf)
|
|
if equals_ignore_case(ds.type, 'dm'):
|
|
with dmPython.connect(user=conf.username, password=conf.password, server=conf.host,
|
|
port=conf.port, **extra_config_dict) as conn, conn.cursor() as cursor:
|
|
cursor.execute("""select OBJECT_NAME
|
|
from all_objects
|
|
where object_type = 'SCH'""", timeout=conf.timeout)
|
|
res = cursor.fetchall()
|
|
res_list = [item[0] for item in res]
|
|
return res_list
|
|
elif equals_ignore_case(ds.type, 'redshift'):
|
|
with redshift_connector.connect(host=conf.host, port=conf.port, database=conf.database, user=conf.username,
|
|
password=conf.password,
|
|
timeout=conf.timeout, **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
|
|
elif equals_ignore_case(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):
|
|
conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration))) if not equals_ignore_case(ds.type,
|
|
"excel") else get_engine_config()
|
|
db = DB.get_db(ds.type)
|
|
sql, sql_param = get_table_sql(ds, conf, get_version(ds))
|
|
if db.connect_type == ConnectType.sqlalchemy:
|
|
with get_session(ds) as session:
|
|
with session.execute(text(sql), {"param": sql_param}) as result:
|
|
res = result.fetchall()
|
|
res_list = [TableSchema(*item) for item in res]
|
|
return res_list
|
|
else:
|
|
extra_config_dict = get_extra_config(conf)
|
|
if equals_ignore_case(ds.type, 'dm'):
|
|
with dmPython.connect(user=conf.username, password=conf.password, server=conf.host,
|
|
port=conf.port, **extra_config_dict) as conn, conn.cursor() as cursor:
|
|
cursor.execute(sql, {"param": sql_param}, timeout=conf.timeout)
|
|
res = cursor.fetchall()
|
|
res_list = [TableSchema(*item) for item in res]
|
|
return res_list
|
|
elif equals_ignore_case(ds.type, 'doris', 'starrocks'):
|
|
ssl_args = {'ssl': {'ssl_mode': 'REQUIRE'}} if conf.ssl else {}
|
|
with pymysql.connect(user=conf.username, passwd=conf.password, host=conf.host,
|
|
port=conf.port, db=conf.database, connect_timeout=conf.timeout,
|
|
read_timeout=conf.timeout, **extra_config_dict,
|
|
**ssl_args) as conn, conn.cursor() as cursor:
|
|
cursor.execute(sql, (sql_param,))
|
|
res = cursor.fetchall()
|
|
res_list = [TableSchema(*item) for item in res]
|
|
return res_list
|
|
elif equals_ignore_case(ds.type, 'redshift'):
|
|
with redshift_connector.connect(host=conf.host, port=conf.port, database=conf.database, user=conf.username,
|
|
password=conf.password,
|
|
timeout=conf.timeout, **extra_config_dict) as conn, conn.cursor() as cursor:
|
|
cursor.execute(sql, (sql_param,))
|
|
res = cursor.fetchall()
|
|
res_list = [TableSchema(*item) for item in res]
|
|
return res_list
|
|
elif equals_ignore_case(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 equals_ignore_case(ds.type, 'es'):
|
|
res = get_es_index(conf)
|
|
res_list = [TableSchema(*item) for item in res]
|
|
return res_list
|
|
elif equals_ignore_case(ds.type, 'hive'):
|
|
with hive.connect(host=conf.host, port=conf.port, username=conf.username,
|
|
database=conf.database, **extra_config_dict) as conn, conn.cursor() as cursor:
|
|
cursor.execute(sql)
|
|
res = cursor.fetchall()
|
|
res_list = [TableSchema(*item) for item in res]
|
|
return res_list
|
|
|
|
|
|
def get_fields(ds: CoreDatasource, table_name: str = None):
|
|
conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration))) if not equals_ignore_case(ds.type,
|
|
"excel") else get_engine_config()
|
|
db = DB.get_db(ds.type)
|
|
sql, p1, p2 = get_field_sql(ds, conf, table_name)
|
|
if db.connect_type == ConnectType.sqlalchemy:
|
|
with get_session(ds) as session:
|
|
with session.execute(text(sql), {"param1": p1, "param2": p2}) as result:
|
|
res = result.fetchall()
|
|
res_list = [ColumnSchema(*item) for item in res]
|
|
return res_list
|
|
else:
|
|
extra_config_dict = get_extra_config(conf)
|
|
if equals_ignore_case(ds.type, 'dm'):
|
|
with dmPython.connect(user=conf.username, password=conf.password, server=conf.host,
|
|
port=conf.port, **extra_config_dict) as conn, conn.cursor() as cursor:
|
|
cursor.execute(sql, {"param1": p1, "param2": p2}, timeout=conf.timeout)
|
|
res = cursor.fetchall()
|
|
res_list = [ColumnSchema(*item) for item in res]
|
|
return res_list
|
|
elif equals_ignore_case(ds.type, 'doris', 'starrocks'):
|
|
ssl_args = {'ssl': {'ssl_mode': 'REQUIRE'}} if conf.ssl else {}
|
|
with pymysql.connect(user=conf.username, passwd=conf.password, host=conf.host,
|
|
port=conf.port, db=conf.database, connect_timeout=conf.timeout,
|
|
read_timeout=conf.timeout, **extra_config_dict,
|
|
**ssl_args) as conn, conn.cursor() as cursor:
|
|
cursor.execute(sql, (p1, p2))
|
|
res = cursor.fetchall()
|
|
res_list = [ColumnSchema(*item) for item in res]
|
|
return res_list
|
|
elif equals_ignore_case(ds.type, 'redshift'):
|
|
with redshift_connector.connect(host=conf.host, port=conf.port, database=conf.database, user=conf.username,
|
|
password=conf.password,
|
|
timeout=conf.timeout, **extra_config_dict) as conn, conn.cursor() as cursor:
|
|
cursor.execute(sql, (p1, p2))
|
|
res = cursor.fetchall()
|
|
res_list = [ColumnSchema(*item) for item in res]
|
|
return res_list
|
|
elif equals_ignore_case(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 equals_ignore_case(ds.type, 'es'):
|
|
res = get_es_fields(conf, table_name)
|
|
res_list = [ColumnSchema(*item) for item in res]
|
|
return res_list
|
|
elif equals_ignore_case(ds.type, 'hive'):
|
|
with hive.connect(host=conf.host, port=conf.port, username=conf.username,
|
|
database=conf.database, **extra_config_dict) as conn, conn.cursor() as cursor:
|
|
cursor.execute(sql)
|
|
res = cursor.fetchall()
|
|
res_list = [ColumnSchema(*item) for item in res]
|
|
return res_list
|
|
|
|
|
|
def convert_value(value, datetime_format='space'):
|
|
"""
|
|
将Python值转换为JSON可序列化的类型
|
|
|
|
:param value: 要转换的值
|
|
:param datetime_format: 日期时间格式
|
|
'iso' - 2024-01-15T14:30:45 (ISO标准，带T)
|
|
'space' - 2024-01-15 14:30:45 (空格分隔，更常见)
|
|
'auto' - 自动选择
|
|
"""
|
|
if value is None:
|
|
return None
|
|
# 处理 bytes 类型（包括 BIT 字段）
|
|
if isinstance(value, bytes):
|
|
# 1. 尝试判断是否是 BIT 类型
|
|
if len(value) <= 8: # BIT 类型通常不会很长
|
|
try:
|
|
# 转换为整数
|
|
int_val = int.from_bytes(value, 'big')
|
|
|
|
# 如果是 0 或 1，返回布尔值更直观
|
|
if int_val in (0, 1):
|
|
return bool(int_val)
|
|
else:
|
|
return int_val
|
|
except:
|
|
# 如果转换失败，尝试解码为字符串
|
|
pass
|
|
|
|
# 2. 尝试解码为 UTF-8 字符串
|
|
try:
|
|
return value.decode('utf-8')
|
|
except UnicodeDecodeError:
|
|
# 3. 如果包含非打印字符，返回十六进制
|
|
if any(b < 32 and b not in (9, 10, 13) for b in value): # 非打印字符
|
|
return f"0x{value.hex()}"
|
|
else:
|
|
# 4. 尝试 Latin-1 解码（不会失败）
|
|
return value.decode('latin-1')
|
|
|
|
elif isinstance(value, bytearray):
|
|
# 处理 bytearray
|
|
return convert_value(bytes(value))
|
|
|
|
if isinstance(value, timedelta):
|
|
# 将 timedelta 转换为秒数（整数）或字符串
|
|
return str(value) # 或 value.total_seconds()
|
|
elif isinstance(value, Decimal):
|
|
return float(value)
|
|
# 4. 处理 datetime
|
|
elif isinstance(value, datetime):
|
|
if datetime_format == 'iso':
|
|
return value.isoformat()
|
|
elif datetime_format == 'space':
|
|
return value.strftime('%Y-%m-%d %H:%M:%S')
|
|
else: # 'auto' 或其他
|
|
# 自动判断：没有时间部分只显示日期
|
|
if value.hour == 0 and value.minute == 0 and value.second == 0 and value.microsecond == 0:
|
|
return value.strftime('%Y-%m-%d')
|
|
else:
|
|
return value.strftime('%Y-%m-%d %H:%M:%S')
|
|
|
|
# 5. 处理 date
|
|
elif isinstance(value, date):
|
|
return value.isoformat() # 总是 YYYY-MM-DD
|
|
|
|
# 6. 处理 time
|
|
elif isinstance(value, time):
|
|
return str(value)
|
|
else:
|
|
return value
|
|
|
|
|
|
def exec_sql(ds: CoreDatasource | AssistantOutDsSchema, sql: str, origin_column=False):
|
|
while sql.endswith(';'):
|
|
sql = sql[:-1]
|
|
# check execute sql only contain read operations
|
|
if not check_sql_read(sql, ds):
|
|
raise ValueError(f"SQL can only contain read operations")
|
|
|
|
db = DB.get_db(ds.type)
|
|
if db.connect_type == ConnectType.sqlalchemy:
|
|
with get_session(ds) as session:
|
|
with session.execute(text(sql)) as result:
|
|
try:
|
|
columns = result.keys()._keys if origin_column else [item.lower() for item in result.keys()._keys]
|
|
res = result.fetchall()
|
|
result_list = [
|
|
{str(columns[i]): convert_value(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))
|
|
else:
|
|
conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration)))
|
|
extra_config_dict = get_extra_config(conf)
|
|
if equals_ignore_case(ds.type, 'dm'):
|
|
with dmPython.connect(user=conf.username, password=conf.password, server=conf.host,
|
|
port=conf.port, **extra_config_dict) as conn, conn.cursor() as cursor:
|
|
try:
|
|
cursor.execute(sql, timeout=conf.timeout)
|
|
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]): convert_value(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 equals_ignore_case(ds.type, 'doris', 'starrocks'):
|
|
ssl_args = {'ssl': {'ssl_mode': 'REQUIRE'}} if conf.ssl else {}
|
|
with pymysql.connect(user=conf.username, passwd=conf.password, host=conf.host,
|
|
port=conf.port, db=conf.database, connect_timeout=conf.timeout,
|
|
read_timeout=conf.timeout, **extra_config_dict,
|
|
**ssl_args) 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]): convert_value(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 equals_ignore_case(ds.type, 'redshift'):
|
|
with redshift_connector.connect(host=conf.host, port=conf.port, database=conf.database, user=conf.username,
|
|
password=conf.password,
|
|
timeout=conf.timeout, **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]): convert_value(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 equals_ignore_case(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]): convert_value(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 equals_ignore_case(ds.type, 'es'):
|
|
try:
|
|
res, columns = get_es_data_by_http(conf, sql)
|
|
columns = [field.get('name') for field in columns] if origin_column else [field.get('name').lower() for
|
|
field in
|
|
columns]
|
|
result_list = [
|
|
{str(columns[i]): convert_value(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 Exception(str(ex))
|
|
elif equals_ignore_case(ds.type, 'hive'):
|
|
with hive.connect(host=conf.host, port=conf.port, username=conf.username,
|
|
database=conf.database, **extra_config_dict) as conn, conn.cursor() as cursor:
|
|
try:
|
|
# Hive uses backticks for identifiers; normalize quoted identifiers as a compatibility fallback.
|
|
hive_sql = re.sub(r'"([A-Za-z_][A-Za-z0-9_]*)"', r'`\1`', sql)
|
|
cursor.execute(hive_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]): convert_value(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(hive_sql, 'utf-8')))}
|
|
except Exception as ex:
|
|
raise ParseSQLResultError(str(ex))
|
|
|
|
|
|
def check_sql_read(sql: str, ds: CoreDatasource | AssistantOutDsSchema):
|
|
try:
|
|
normalized_sql = sql.strip().lstrip("(").strip()
|
|
first_keyword = normalized_sql.split(None, 1)[0].upper() if normalized_sql else ""
|
|
allowed_read_commands = {"SELECT", "WITH", "SHOW", "DESCRIBE", "DESC", "EXPLAIN"}
|
|
denied_write_commands = {
|
|
"INSERT", "UPDATE", "DELETE", "CREATE", "DROP", "ALTER",
|
|
"TRUNCATE", "MERGE", "COPY", "REPLACE", "GRANT", "REVOKE",
|
|
"USE", "SET", "CALL"
|
|
}
|
|
|
|
if not first_keyword:
|
|
raise ValueError("Parse SQL Error")
|
|
if first_keyword in denied_write_commands:
|
|
return False
|
|
|
|
dialect = None
|
|
if equals_ignore_case(ds.type, 'mysql', 'doris', 'starrocks'):
|
|
dialect = 'mysql'
|
|
elif equals_ignore_case(ds.type, 'sqlServer'):
|
|
dialect = 'tsql'
|
|
elif equals_ignore_case(ds.type, 'hive'):
|
|
dialect = 'hive'
|
|
|
|
statements = sqlglot.parse(sql, dialect=dialect)
|
|
|
|
if not statements:
|
|
raise ValueError("Parse SQL Error")
|
|
|
|
write_types = (
|
|
exp.Insert, exp.Update, exp.Delete,
|
|
exp.Create, exp.Drop, exp.Alter,
|
|
exp.Merge, exp.Copy
|
|
)
|
|
|
|
for stmt in statements:
|
|
if stmt is None:
|
|
continue
|
|
if isinstance(stmt, write_types):
|
|
return False
|
|
|
|
return first_keyword in allowed_read_commands
|
|
|
|
except Exception as e:
|
|
raise ValueError(f"Parse SQL Error: {e}")
|
|
|
|
|
|
def checkParams(extraParams: str, illegalParams: List[str]):
|
|
kvs = extraParams.split('&')
|
|
for kv in kvs:
|
|
if kv and '=' in kv:
|
|
k, v = kv.split('=')
|
|
if k in illegalParams:
|
|
raise HTTPException(status_code=500, detail=f'Illegal Parameter: {k}')
|