diff --git a/backend/apps/datasource/crud/datasource.py b/backend/apps/datasource/crud/datasource.py index a3f9b65f..efb6bb97 100644 --- a/backend/apps/datasource/crud/datasource.py +++ b/backend/apps/datasource/crud/datasource.py @@ -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""" diff --git a/backend/apps/db/constant.py b/backend/apps/db/constant.py index fc118d9d..c4de1cd0 100644 --- a/backend/apps/db/constant.py +++ b/backend/apps/db/constant.py @@ -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 diff --git a/backend/apps/db/db.py b/backend/apps/db/db.py index 73aea904..47147d77 100644 --- a/backend/apps/db/db.py +++ b/backend/apps/db/db.py @@ -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 diff --git a/backend/apps/db/db_sql.py b/backend/apps/db/db_sql.py index 48282bf6..34982f6f 100644 --- a/backend/apps/db/db_sql.py +++ b/backend/apps/db/db_sql.py @@ -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, diff --git a/backend/apps/db/type.py b/backend/apps/db/type.py index 3a033898..e2c6c4fb 100644 --- a/backend/apps/db/type.py +++ b/backend/apps/db/type.py @@ -12,5 +12,6 @@ def db_type_relation() -> Dict: "oracle": "Oracle", "ck": "ClickHouse", "dm": "达梦", - "doris": "Apache Doris" + "doris": "Apache Doris", + "redshift": "AWS Redshift" } diff --git a/backend/pyproject.toml b/backend/pyproject.toml index 7afa8e96..d031d8be 100644 --- a/backend/pyproject.toml +++ b/backend/pyproject.toml @@ -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] diff --git a/backend/template.yaml b/backend/template.yaml index 98e1bf21..1fc0fed6 100644 --- a/backend/template.yaml +++ b/backend/template.yaml @@ -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格式返回： diff --git a/frontend/src/assets/datasource/icon_redshift.png b/frontend/src/assets/datasource/icon_redshift.png new file mode 100644 index 00000000..a6c6bdcf Binary files /dev/null and b/frontend/src/assets/datasource/icon_redshift.png differ diff --git a/frontend/src/views/ds/js/ds-type.ts b/frontend/src/views/ds/js/ds-type.ts index e1ce6c91..ab43dc9c 100644 --- a/frontend/src/views/ds/js/ds-type.ts +++ b/frontend/src/views/ds/js/ds-type.ts @@ -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']