mirror of
https://gh-proxy.org/https://github.com/klinsmaya/SQLBot.git
synced 2026-11-08 09:17:06 +08:00
feat: support starrocks datasource
This commit is contained in:
@@ -290,7 +290,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":
|
||||
if ds.type == "mysql" or ds.type == "doris" or ds.type == "starrocks":
|
||||
sql = f"""SELECT `{"`, `".join(fields)}` FROM `{data.table.table_name}`
|
||||
{where}
|
||||
LIMIT 100"""
|
||||
|
||||
@@ -24,6 +24,7 @@ class DB(Enum):
|
||||
redshift = ('redshift', 'AWS Redshift', '"', '"', ConnectType.py_driver)
|
||||
es = ('es', 'Elasticsearch', '"', '"', ConnectType.py_driver)
|
||||
kingbase = ('kingbase', 'Kingbase', '"', '"', ConnectType.py_driver)
|
||||
starrocks = ('starrocks', 'StarRocks', '"', '"', ConnectType.py_driver)
|
||||
|
||||
def __init__(self, type, db_name, prefix, suffix, connect_type: ConnectType):
|
||||
self.type = type
|
||||
|
||||
@@ -164,7 +164,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 ds.type == 'doris':
|
||||
elif ds.type == 'doris' or ds.type == "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:
|
||||
@@ -259,7 +259,7 @@ def get_version(ds: CoreDatasource | AssistantOutDsSchema):
|
||||
cursor.execute(sql, timeout=10, **extra_config_dict)
|
||||
res = cursor.fetchall()
|
||||
version = res[0][0]
|
||||
elif ds.type == 'doris':
|
||||
elif ds.type == 'doris' or ds.type == "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:
|
||||
@@ -337,7 +337,7 @@ def get_tables(ds: CoreDatasource):
|
||||
res = cursor.fetchall()
|
||||
res_list = [TableSchema(*item) for item in res]
|
||||
return res_list
|
||||
elif ds.type == 'doris':
|
||||
elif ds.type == 'doris' or ds.type == "starrocks":
|
||||
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) as conn, conn.cursor() as cursor:
|
||||
@@ -387,7 +387,7 @@ 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 == 'doris':
|
||||
elif ds.type == 'doris' or ds.type == "starrocks":
|
||||
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) as conn, conn.cursor() as cursor:
|
||||
@@ -459,7 +459,7 @@ 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 == 'doris':
|
||||
elif ds.type == 'doris' or ds.type == "starrocks":
|
||||
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) as conn, conn.cursor() as cursor:
|
||||
|
||||
@@ -4,7 +4,7 @@ from apps.datasource.models.datasource import CoreDatasource, DatasourceConf
|
||||
|
||||
|
||||
def get_version_sql(ds: CoreDatasource, conf: DatasourceConf):
|
||||
if ds.type == "mysql" or ds.type == "doris":
|
||||
if ds.type == "mysql" or ds.type == "doris" or ds.type == "starrocks":
|
||||
return """
|
||||
SELECT VERSION()
|
||||
"""
|
||||
@@ -134,7 +134,7 @@ def get_table_sql(ds: CoreDatasource, conf: DatasourceConf, db_version: str = ''
|
||||
relkind in ('r','p', 'f')
|
||||
AND relnamespace = (SELECT oid FROM pg_namespace WHERE nspname = %s)
|
||||
""", conf.dbSchema
|
||||
elif ds.type == "doris":
|
||||
elif ds.type == "doris" or ds.type == "starrocks":
|
||||
return """
|
||||
SELECT
|
||||
TABLE_NAME,
|
||||
@@ -281,7 +281,7 @@ def get_field_sql(ds: CoreDatasource, conf: DatasourceConf, table_name: str = No
|
||||
"""
|
||||
sql2 = " AND c.TABLE_NAME = :param2" if table_name is not None and table_name != "" else ""
|
||||
return sql1 + sql2, conf.dbSchema, table_name
|
||||
elif ds.type == "doris":
|
||||
elif ds.type == "doris" or ds.type == "starrocks":
|
||||
sql1 = """
|
||||
SELECT
|
||||
COLUMN_NAME,
|
||||
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 974 B |
@@ -9,6 +9,7 @@ import doris from '@/assets/datasource/icon_doris.png'
|
||||
import redshift from '@/assets/datasource/icon_redshift.png'
|
||||
import es from '@/assets/datasource/icon_es.png'
|
||||
import kingbase from '@/assets/datasource/icon_kingbase.png'
|
||||
import starrocks from '@/assets/datasource/icon_starrocks.png'
|
||||
import { i18n } from '@/i18n'
|
||||
|
||||
const t = i18n.global.t
|
||||
@@ -24,6 +25,7 @@ export const dsType = [
|
||||
{ label: 'AWS Redshift', value: 'redshift' },
|
||||
{ label: 'Elasticsearch', value: 'es' },
|
||||
{ label: 'Kingbase', value: 'kingbase' },
|
||||
{ label: 'StarRocks', value: 'starrocks' },
|
||||
]
|
||||
|
||||
export const dsTypeWithImg = [
|
||||
@@ -38,6 +40,7 @@ export const dsTypeWithImg = [
|
||||
{ name: 'AWS Redshift', type: 'redshift', img: redshift },
|
||||
{ name: 'Elasticsearch', type: 'es', img: es },
|
||||
{ name: 'Kingbase', type: 'kingbase', img: kingbase },
|
||||
{ name: 'StarRocks', type: 'starrocks', img: starrocks },
|
||||
]
|
||||
|
||||
export const haveSchema = ['sqlServer', 'pg', 'oracle', 'dm', 'redshift', 'kingbase']
|
||||
|
||||
Reference in New Issue
Block a user