mirror of
https://gh-proxy.org/https://github.com/klinsmaya/SQLBot.git
synced 2026-11-08 01:06:59 +08:00
104 lines
3.5 KiB
Python
104 lines
3.5 KiB
Python
from sqlalchemy import create_engine, text, Result
|
|
from sqlalchemy.orm import sessionmaker
|
|
from apps.datasource.models.datasource import DatasourceConf, CoreDatasource, TableSchema, ColumnSchema
|
|
import urllib.parse
|
|
from typing import Any
|
|
import json
|
|
from apps.datasource.utils.utils import aes_decrypt
|
|
from common.core.deps import SessionDep
|
|
|
|
|
|
def get_uri(ds: CoreDatasource):
|
|
conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration)))
|
|
db_url: str
|
|
if ds.type == "mysql":
|
|
db_url = f"mysql+pymysql://{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_session(ds: CoreDatasource):
|
|
conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration)))
|
|
db_url: str
|
|
if ds.type == "mysql":
|
|
db_url = f"mysql+pymysql://{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.'
|
|
engine = create_engine(db_url)
|
|
session_maker = sessionmaker(bind=engine)
|
|
session = session_maker()
|
|
return session
|
|
|
|
|
|
def get_tables(ds: CoreDatasource):
|
|
conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration)))
|
|
session = get_session(ds)
|
|
result: Result[Any]
|
|
try:
|
|
if ds.type == "mysql":
|
|
sql = f"""SELECT
|
|
TABLE_NAME,
|
|
TABLE_COMMENT
|
|
FROM
|
|
information_schema.TABLES
|
|
WHERE
|
|
TABLE_SCHEMA = '{conf.database}';"""
|
|
result = session.execute(text(sql))
|
|
|
|
res = result.fetchall()
|
|
res_list = [TableSchema(*item) for item in res]
|
|
return res_list
|
|
finally:
|
|
if result is not None:
|
|
result.close()
|
|
if session is not None:
|
|
session.close()
|
|
|
|
|
|
def get_fields(ds: CoreDatasource, table_name: str = None):
|
|
conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration)))
|
|
session = get_session(ds)
|
|
result: Result[Any]
|
|
try:
|
|
if ds.type == "mysql":
|
|
sql1 = f"""SELECT
|
|
COLUMN_NAME,
|
|
DATA_TYPE,
|
|
COLUMN_COMMENT
|
|
FROM
|
|
INFORMATION_SCHEMA.COLUMNS
|
|
WHERE
|
|
TABLE_SCHEMA = '{conf.database}'"""
|
|
sql2 = f""" AND TABLE_NAME = '{table_name}';""" if table_name is not None and table_name != "" else ";"
|
|
sql = sql1 + sql2
|
|
result = session.execute(text(sql))
|
|
|
|
res = result.fetchall()
|
|
res_list = [ColumnSchema(*item) for item in res]
|
|
return res_list
|
|
finally:
|
|
if result is not None:
|
|
result.close()
|
|
if session is not None:
|
|
session.close()
|
|
|
|
|
|
def exec_sql(ds: CoreDatasource, sql: str):
|
|
session = get_session(ds)
|
|
result = session.execute(text(sql))
|
|
try:
|
|
columns = result.keys()._keys
|
|
res = result.fetchall()
|
|
result_list = [
|
|
{columns[i]: value for i, value in enumerate(tuple_item)}
|
|
for tuple_item in res
|
|
]
|
|
return {"fields": columns, "data": result_list}
|
|
finally:
|
|
if result is not None:
|
|
result.close()
|
|
if session is not None:
|
|
session.close()
|
|
|