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 AWS Redshift datasource

Browse Source
This commit is contained in:
junjun
2025-08-28 13:56:20 +08:00
parent 79f7630537
commit 6c975b6dc8
9 changed files with 84 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/apps/datasource/crud/datasource.py
+1 -1
View File
@@ -275,7 +275,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":
elif ds.type == "pg" or ds.type == "excel" or ds.type == "redshift":
sql = f"""SELECT "{'", "'.join(fields)}" FROM "{conf.dbSchema}"."{data.table.table_name}"
{where}
LIMIT 100"""
backend/apps/db/constant.py
+1
View File
@@ -21,6 +21,7 @@ class DB(Enum):
ck = ('ck', '"', '"', ConnectType.sqlalchemy)
dm = ('dm', '"', '"', ConnectType.py_driver)
doris = ('doris', '`', '`', ConnectType.py_driver)
redshift = ('redshift', '"', '"', ConnectType.py_driver)
def __init__(self, type, prefix, suffix, connect_type: ConnectType):
self.type = type
backend/apps/db/db.py
+59
View File
@@ -9,6 +9,7 @@ from apps.db.db_sql import get_table_sql, get_field_sql, get_version_sql
if platform.system() != "Darwin":
import dmPython
import pymysql
import redshift_connector
from sqlalchemy import create_engine, text, Engine
from sqlalchemy.orm import sessionmaker
@@ -139,6 +140,19 @@ def check_connection(trans: Trans, ds: CoreDatasource, is_raise: bool = False):
if is_raise:
raise HTTPException(status_code=500, detail=trans('i18n_ds_invalid') + f': {e.args}')
return False
elif ds.type == 'redshift':
with redshift_connector.connect(host=conf.host, port=conf.port, database=conf.database, user=conf.username,
password=conf.password,
timeout=10) 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
def get_version(ds: CoreDatasource):
@@ -165,6 +179,8 @@ def get_version(ds: CoreDatasource):
cursor.execute(sql)
res = cursor.fetchall()
return res[0][0]
elif ds.type == 'redshift':
return ''
except Exception as e:
print(e)
return ''
@@ -194,6 +210,14 @@ def get_schema(ds: CoreDatasource):
res = cursor.fetchall()
res_list = [item[0] for item in res]
return res_list
elif 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) as conn, conn.cursor() as cursor:
cursor.execute(f"""SELECT nspname FROM pg_namespace""")
res = cursor.fetchall()
res_list = [item[0] for item in res]
return res_list
def get_tables(ds: CoreDatasource):
@@ -222,6 +246,14 @@ def get_tables(ds: CoreDatasource):
res = cursor.fetchall()
res_list = [TableSchema(*item) for item in res]
return res_list
elif 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) 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):
@@ -250,6 +282,14 @@ 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 == 'redshift':
with redshift_connector.connect(host=conf.host, port=conf.port, database=conf.database, user=conf.username,
password=conf.password,
timeout=conf.timeout) as conn, conn.cursor() as cursor:
cursor.execute(sql)
res = cursor.fetchall()
res_list = [ColumnSchema(*item) for item in res]
return res_list
def exec_sql(ds: CoreDatasource | AssistantOutDsSchema, sql: str, origin_column=False):
@@ -311,3 +351,22 @@ def exec_sql(ds: CoreDatasource | AssistantOutDsSchema, sql: str, origin_column=
"sql": bytes.decode(base64.b64encode(bytes(sql, 'utf-8')))}
except Exception as ex:
raise ex
elif 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) 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 ex
backend/apps/db/db_sql.py
+14 -1
View File
@@ -28,6 +28,8 @@ def get_version_sql(ds: CoreDatasource, conf: DatasourceConf):
return f"""
SELECT * FROM v$version
"""
elif ds.type == 'redshift':
return ''
def get_table_sql(ds: CoreDatasource, conf: DatasourceConf):
@@ -107,6 +109,17 @@ def get_table_sql(ds: CoreDatasource, conf: DatasourceConf):
where owner='{conf.dbSchema}'
AND (table_type = 'TABLE' or table_type = 'VIEW')
"""
elif ds.type == 'redshift':
return f"""
SELECT
relname AS TableName,
obj_description(relfilenode::regclass, 'pg_class') AS TableDescription
FROM
pg_class
WHERE
relkind in ('r','p', 'f')
AND relnamespace = (SELECT oid FROM pg_namespace WHERE nspname = '{conf.dbSchema}')
"""
def get_field_sql(ds: CoreDatasource, conf: DatasourceConf, table_name: str = None):
@@ -141,7 +154,7 @@ def get_field_sql(ds: CoreDatasource, conf: DatasourceConf, table_name: str = No
"""
sql2 = f" AND C.TABLE_NAME = '{table_name}'" if table_name is not None and table_name != "" else ""
return sql1 + sql2
elif ds.type == "pg" or ds.type == "excel":
elif ds.type == "pg" or ds.type == "excel" or ds.type == "redshift":
sql1 = f"""
SELECT a.attname AS COLUMN_NAME,
pg_catalog.format_type(a.atttypid, a.atttypmod) AS DATA_TYPE,
backend/apps/db/type.py
+2 -1
View File
@@ -12,5 +12,6 @@ def db_type_relation() -> Dict:
"oracle": "Oracle",
"ck": "ClickHouse",
"dm": "达梦",
"doris": "Apache Doris"
"doris": "Apache Doris",
"redshift": "AWS Redshift"
}
backend/pyproject.toml
+1
View File
@@ -49,6 +49,7 @@ dependencies = [
"clickhouse-sqlalchemy>=0.3.2",
"dicttoxml>=1.7.16",
"dmpython>=2.5.22; platform_system != 'Darwin'",
"redshift-connector>=2.1.8",
]
[project.optional-dependencies]
backend/template.yaml
+2 -2
View File
@@ -39,7 +39,7 @@ template:
- SQL查询的字段若是函数字段，如 COUNT(),CAST() 等，必须加上别名
- 计算占比，百分比类型字段，保留两位小数，以%结尾。
- 生成SQL时，必须避免关键字冲突。
- 如数据库引擎是 PostgreSQL、Oracle、ClickHouse、达梦（DM），则在schema、表名、字段名、别名外层加双引号；
- 如数据库引擎是 PostgreSQL、Oracle、ClickHouse、达梦（DM）、AWS Redshift，则在schema、表名、字段名、别名外层加双引号；
- 如数据库引擎是 MySQL、Doris，则在表名、字段名、别名外层加反引号；
- 如数据库引擎是 Microsoft SQL Server，则在schema、表名、字段名、别名外层加方括号。
- 以PostgreSQL为例，查询Schema为TEST表TABLE下所有数据，则生成的SQL为：
@@ -224,7 +224,7 @@ template:
- 如果存在冗余的过滤条件则进行去重后再生成新SQL。
- 给过滤条件中的字段前加上表别名（如果没有表别名则加表名），如：table.field。
- 生成SQL时，必须避免关键字冲突:
- 如数据库引擎是 PostgreSQL、Oracle、ClickHouse、达梦（DM），则在schema、表名、字段名、别名外层加双引号；
- 如数据库引擎是 PostgreSQL、Oracle、ClickHouse、达梦（DM）、AWS Redshift，则在schema、表名、字段名、别名外层加双引号；
- 如数据库引擎是 MySQL、Doris，则在表名、字段名、别名外层加反引号；
- 如数据库引擎是 Microsoft SQL Server，则在schema、表名、字段名、别名外层加方括号。
- 生成的SQL使用JSON格式返回：
frontend/src/assets/datasource/icon_redshift.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 486 B

frontend/src/views/ds/js/ds-type.ts
+4 -1
View File
@@ -6,6 +6,7 @@ import sqlServer from '@/assets/datasource/icon_SQL_Server.png'
import ck from '@/assets/datasource/icon_ck.png'
import dm from '@/assets/datasource/icon_dm.png'
import doris from '@/assets/datasource/icon_doris.png'
import redshift from '@/assets/datasource/icon_redshift.png'
import { i18n } from '@/i18n'
const t = i18n.global.t
@@ -18,6 +19,7 @@ export const dsType = [
{ label: 'ClickHouse', value: 'ck' },
{ label: '达梦', value: 'dm' },
{ label: 'Apache Doris', value: 'doris' },
{ label: 'AWS Redshift', value: 'redshift' },
]
export const dsTypeWithImg = [
@@ -29,6 +31,7 @@ export const dsTypeWithImg = [
{ name: 'ClickHouse', type: 'ck', img: ck },
{ name: '达梦', type: 'dm', img: dm },
{ name: 'Apache Doris', type: 'doris', img: doris },
{ name: 'AWS Redshift', type: 'redshift', img: redshift },
]
export const haveSchema = ['sqlServer', 'pg', 'oracle', 'dm']
export const haveSchema = ['sqlServer', 'pg', 'oracle', 'dm', 'redshift']