add support mssql

Author: junjie.miao

README.md

@@ -6,6 +6,10 @@
 
 ### Demonstration
 
+Currently supported database types: mysql, oracle, postgresql, or mssql.
+
+目前支持的数据库类型：mysql、oracle、postgresql、mssql。
+
 ![db_query](db_query/_assets/db_query.png)
 
 ![db_query_sql_query](db_query/_assets/db_query_sql_query.png)

db_query/requirements.txt

@@ -5,4 +5,5 @@ pandas==2.2.3
 SQLAlchemy==2.0.35
 PyMySQL==1.1.1
 oracledb==2.2.1
-psycopg2-binary==2.9.10
+psycopg2-binary==2.9.10
+pymssql==2.3.4

db_query/tools/db_util.py

@@ -34,6 +34,8 @@ def get_driver_name(self):
             driver_name = 'oracle+oracledb'
         elif db_type == 'postgresql':
             driver_name = 'postgresql+psycopg2'
+        elif db_type == 'mssql':
+            driver_name = 'mssql+pymssql'
         return driver_name
 
     def get_url(self):
@@ -76,10 +78,11 @@ def run_query(self, query_sql: str) -> list[dict]:
         return records
 
     def test_sql(self):
-        if self.db_type in {'mysql', 'postgresql'}:
-            return "SELECT 1"
-        elif self.db_type == 'oracle':
+        db_type = self.db_type.lower()
+        if db_type == 'oracle':
             return "SELECT 1 FROM DUAL"
+        else:
+            return "SELECT 1"
 
     @staticmethod
     def is_not_empty(s: str):

db_query/tools/sql_query.yaml

@@ -29,13 +29,17 @@ parameters:
         label:
           en_US: PostgreSQL
           zh_Hans: PostgreSQL
+      - value: mssql
+        label:
+          en_US: Microsoft SQL Server
+          zh_Hans: Microsoft SQL Server
     default: mysql
     label:
       en_US: Database type
       zh_Hans: 数据库类型
     human_description:
-      en_US: Used for selecting the database type, mysql, oracle or postgresql.
-      zh_Hans: 用于选择数据库类型，mysql、oracle或postgresql。
+      en_US: Used for selecting the database type, mysql, oracle, postgresql or mssql.
+      zh_Hans: 用于选择数据库类型，mysql、oracle、postgresql或mssql。
     form: llm
   - name: db_host
     type: string

db_query_pre_auth/provider/db_query.yaml

@@ -38,8 +38,8 @@ credentials_for_provider:
       en_US: Database type
       zh_Hans: 数据库类型
     help:
-      en_US: Used for selecting the database type, mysql, oracle or postgresql.
-      zh_Hans: 用于选择数据库类型，mysql、oracle或postgresql。
+      en_US: Used for selecting the database type, mysql, oracle, postgresql or mssql.
+      zh_Hans: 用于选择数据库类型，mysql、oracle、postgresql或mssql。
   db_host:
     type: text-input
     required: true

db_query_pre_auth/requirements.txt

@@ -5,4 +5,5 @@ pandas==2.2.3
 SQLAlchemy==2.0.35
 PyMySQL==1.1.1
 oracledb==2.2.1
-psycopg2-binary==2.9.10
+psycopg2-binary==2.9.10
+pymssql==2.3.4

db_query_pre_auth/tools/db_util.py

@@ -34,6 +34,8 @@ def get_driver_name(self):
             driver_name = 'oracle+oracledb'
         elif db_type == 'postgresql':
             driver_name = 'postgresql+psycopg2'
+        elif db_type == 'mssql':
+            driver_name = 'mssql+pymssql'
         return driver_name
 
     def get_url(self):
@@ -76,10 +78,11 @@ def run_query(self, query_sql: str) -> list[dict]:
         return records
 
     def test_sql(self):
-        if self.db_type in {'mysql', 'postgresql'}:
-            return "SELECT 1"
-        elif self.db_type == 'oracle':
+        db_type = self.db_type.lower()
+        if db_type == 'oracle':
             return "SELECT 1 FROM DUAL"
+        else:
+            return "SELECT 1"
 
     @staticmethod
     def is_not_empty(s: str):