feat: support AWS Redshift datasource

Author: XiaJunjie2020

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"""

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

backend/apps/db/db.py

@@ -9,6 +9,7 @@
 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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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']