mirror of
https://gh-proxy.org/https://github.com/klinsmaya/SQLBot.git
synced 2026-11-08 01:06:59 +08:00
fix: Fix the data source-related issues in the PR
This commit is contained in:
@@ -329,7 +329,7 @@ def preview(session: SessionDep, current_user: CurrentUser, id: int, data: Table
|
||||
|
||||
conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration))) if ds.type != "excel" else get_engine_config()
|
||||
sql: str = ""
|
||||
if ds.type == "mysql" or ds.type == "doris" or ds.type == "starrocks":
|
||||
if ds.type == "mysql" or ds.type == "doris" or ds.type == "starrocks" or ds.type == "hive":
|
||||
sql = f"""SELECT `{"`, `".join(fields)}` FROM `{data.table.table_name}`
|
||||
{where}
|
||||
LIMIT 100"""
|
||||
|
||||
+11
-32
@@ -38,7 +38,6 @@ 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(
|
||||
@@ -159,9 +158,10 @@ def get_engine(ds: CoreDatasource, timeout: int = 0) -> Engine:
|
||||
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
|
||||
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)
|
||||
engine = create_engine(get_uri(ds), connect_args={"connect_timeout": conf.timeout, "ssl": ssl_mode},
|
||||
poolclass=NullPool)
|
||||
elif equals_ignore_case(ds.type, 'sqlite'):
|
||||
engine = create_engine(get_uri(ds), connect_args={"check_same_thread": False}, poolclass=NullPool)
|
||||
else: # ck
|
||||
@@ -271,7 +271,7 @@ 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 equals_ignore_case(ds.type, 'es'):
|
||||
es_conn = get_es_connect(conf)
|
||||
if es_conn.ping():
|
||||
@@ -314,7 +314,7 @@ def get_version(ds: CoreDatasource | AssistantOutDsSchema):
|
||||
# conf.timeout = 10
|
||||
db = DB.get_db(ds.type)
|
||||
sql = get_version_sql(ds, conf)
|
||||
if equals_ignore_case(ds.type, 'sqlite'):
|
||||
if not sql:
|
||||
return ''
|
||||
try:
|
||||
if db.connect_type == ConnectType.sqlalchemy:
|
||||
@@ -397,30 +397,6 @@ def get_schema(ds: CoreDatasource):
|
||||
res = cursor.fetchall()
|
||||
res_list = [item[0] for item in res]
|
||||
return res_list
|
||||
elif equals_ignore_case(ds.type, 'hive'):
|
||||
conn = hive.connect(host=conf.host, port=conf.port, username=conf.username,
|
||||
database=conf.database, **extra_config_dict)
|
||||
cursor = conn.cursor()
|
||||
cursor.execute('SHOW DATABASES')
|
||||
res = cursor.fetchall()
|
||||
res_list = [item[0] for item in res]
|
||||
cursor.close()
|
||||
conn.close()
|
||||
return res_list
|
||||
elif equals_ignore_case(ds.type, 'doris', 'starrocks'):
|
||||
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) as conn, conn.cursor() as cursor:
|
||||
cursor.execute('SHOW DATABASES')
|
||||
res = cursor.fetchall()
|
||||
res_list = [item[0] for item in res]
|
||||
return res_list
|
||||
elif equals_ignore_case(ds.type, 'ck'):
|
||||
with get_session(ds) as session:
|
||||
with session.execute(text('SHOW DATABASES')) as result:
|
||||
res = result.fetchall()
|
||||
res_list = [item[0] for item in res]
|
||||
return res_list
|
||||
|
||||
|
||||
def get_tables(ds: CoreDatasource):
|
||||
@@ -456,7 +432,8 @@ def get_tables(ds: CoreDatasource):
|
||||
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:
|
||||
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]
|
||||
@@ -527,7 +504,8 @@ def get_fields(ds: CoreDatasource, table_name: str = None):
|
||||
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:
|
||||
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]
|
||||
@@ -684,7 +662,8 @@ def exec_sql(ds: CoreDatasource | AssistantOutDsSchema, sql: str, origin_column=
|
||||
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:
|
||||
read_timeout=conf.timeout, **extra_config_dict,
|
||||
**ssl_args) as conn, conn.cursor() as cursor:
|
||||
try:
|
||||
cursor.execute(sql)
|
||||
res = cursor.fetchall()
|
||||
|
||||
@@ -29,9 +29,7 @@ def get_version_sql(ds: CoreDatasource, conf: DatasourceConf):
|
||||
return """
|
||||
SELECT * FROM v$version
|
||||
"""
|
||||
elif equals_ignore_case(ds.type, "redshift"):
|
||||
return ''
|
||||
elif equals_ignore_case(ds.type, "sqlite"):
|
||||
elif equals_ignore_case(ds.type, "redshift", "sqlite", "hive"):
|
||||
return ''
|
||||
|
||||
|
||||
|
||||
@@ -54,7 +54,7 @@ dependencies = [
|
||||
"ldap3>=2.9.1",
|
||||
"sqlglot>=28.6.0",
|
||||
"numpy==2.3.5",
|
||||
"pyhive[hive]>=0.7.0",
|
||||
"pyhive[hive_pure_sasl]>=0.7.0",
|
||||
"thrift-sasl"
|
||||
]
|
||||
|
||||
|
||||
Reference in New Issue
Block a user