feat: support starrocks datasource

Author: XiaJunjie2020

backend/apps/datasource/crud/datasource.py

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

backend/apps/db/constant.py

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

backend/apps/db/db.py

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

backend/apps/db/db_sql.py

@@ -4,7 +4,7 @@
 
 
 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,

frontend/src/assets/datasource/icon_starrocks.png

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