Merge pull request #54 from Qiao-gengle/main

fix bug

Author: jaguarliuu

database_schema/connector.py

@@ -79,4 +79,4 @@ def get_db_schema(
         return None
     finally:
         if engine:
-            engine.dispose()
+            engine.dispose()

database_schema/inspectors/sqlserver.py

@@ -4,24 +4,21 @@
 from urllib.parse import quote_plus
 
 class SQLServerInspector(BaseInspector):
-    """SQLServer元数据获取实现"""
-    
-    def __init__(self, host: str, port: int, database: str,
-                 username: str, password: str, schema_name: str = None, **kwargs):
-        # 在SQL Server中，schema和database是不同的概念
-        # 如果未指定schema，默认使用"dbo"
-        schema_name = schema_name or "dbo"
-        super().__init__(host, port, database, username, password, schema_name)
+    """SQL Server元数据获取实现"""
+
+    def __init__(self, host, port, database, username, password, schema_name = None, **kwargs):
+        super().__init__(host, port, database, username, password, schema_name, **kwargs)
+        self.schema_name = schema_name if schema_name != None else 'dbo'
 
     def build_conn_str(self, host: str, port: int, database: str,
                       username: str, password: str) -> str:
-        import os
-        driver = 'ODBC+Driver+17+for+SQL+Server' if os.name == 'posix' else 'SQL Server'
+
+        # import os
+        # driver = 'ODBC+Driver+17+for+SQL+Server' if os.name == 'posix' else 'SQL Server'
         # driver = 'ODBC+Driver+17+for+SQL+Server'
         return (
-            f"mssql+pyodbc://{quote_plus(username)}:{quote_plus(password)}"
-            f"@{host},{port}/{database}?"
-            f"driver={driver}"
+            f"mssql+pymssql://{quote_plus(username)}:{quote_plus(password)}"
+            f"@{host}:{port}/{database}"
         )
 
     def get_table_names(self, inspector: reflection.Inspector) -> list[str]:

provider/rookie_text2data.py

@@ -3,4 +3,4 @@
 
 class RookieText2dataProvider(ToolProvider):
     def _validate_credentials(self, credentials: dict[str, Any]) -> None:
-        pass
+        pass

tools/rookie_excute_sql.py

@@ -46,7 +46,7 @@ def _validate_and_prepare_params(self, params: dict) -> dict:
 
         if self._contains_risk_commands(params['sql']):
             raise ValueError("SQL语句包含危险操作")
-
+        params['schema'] = params.get('schema')if params.get('schema') != None else 'dbo' if params['db_type'] == 'sqlserver' else 'public'
         # 数据库执行参数
         execute_params = {
             'db_type': params['db_type'],
@@ -197,4 +197,4 @@ def _safe_serialize(self, data: Any) -> Any:
         """安全的数据序列化"""
         return json.loads(
             json.dumps(data, default=self._custom_serializer, ensure_ascii=False)
-        )
+        )

tools/rookie_text2data.py

@@ -23,7 +23,7 @@ def _invoke(self, tool_parameters: dict[str, Any]) -> Generator[ToolInvokeMessag
             username=tool_parameters['username'],
             password=tool_parameters['password'],
             table_names=tool_parameters['table_names'],
-            schema_name=tool_parameters.get('schema_name', 'public')
+            schema_name=tool_parameters.get('schema_name')
         )
         with_comment = tool_parameters.get('with_comment', False)
         dsl_text = format_schema_dsl(meta_data, with_type=True, with_comment=with_comment)

utils/alchemy_db_client.py

@@ -4,108 +4,108 @@
 from urllib.parse import quote_plus # 用于对URL进行编码
 from typing import Any, Optional, Union
 
-def get_db_schema(
-        db_type: str,
-        host: str,
-        port: int,
-        database: str,
-        username: str,
-        password: str,
-        table_names: str | None = None
-) -> dict[str, Any] | None:
-    """
-    获取数据库表结构信息
-    :param db_type: 数据库类型 (mysql/oracle/sqlserver/postgresql)
-    :param host: 主机地址
-    :param port: 端口号
-    :param database: 数据库名
-    :param username: 用户名
-    :param password: 密码
-    :param table_names: 要查询的表名，以逗号分隔的字符串，如果为None则查询所有表
-    :return: 包含所有表结构信息的字典
-    """
-    result: dict[str, Any] = {}
-    db_type = 'mssql' if db_type == 'sqlserver' else db_type
-    # 构建连接URL
-    driver = {
-        'mysql': 'pymysql',
-        'oracle': 'cx_oracle',
-        'mssql': 'pyodbc',
-        'postgresql': 'psycopg2'
-    }.get(db_type.lower(), '')
-
-    encoded_username = quote_plus(username)
-    encoded_password = quote_plus(password)
-    # separator = ':' if db_type is 'mssql' else ','
-
-    engine = create_engine(f'{db_type.lower()}+{driver}://{encoded_username}:{encoded_password}@{host}{separator}{port}/{database}')
-    inspector = inspect(engine)
-
-    # 获取字段注释的SQL语句
-    column_comment_sql = {
-        'mysql': f"SELECT COLUMN_COMMENT FROM information_schema.COLUMNS WHERE TABLE_SCHEMA = '{database}' AND TABLE_NAME = :table_name AND COLUMN_NAME = :column_name",
-        'oracle': "SELECT COMMENTS FROM ALL_COL_COMMENTS WHERE TABLE_NAME = :table_name AND COLUMN_NAME = :column_name",
-        'mssql': "SELECT CAST(ep.value AS NVARCHAR(MAX)) FROM sys.columns c LEFT JOIN sys.extended_properties ep ON ep.major_id = c.object_id AND ep.minor_id = c.column_id WHERE OBJECT_NAME(c.object_id) = :table_name AND c.name = :column_name",
-        'postgresql': """
-            SELECT pg_catalog.col_description(c.oid, cols.ordinal_position::int)
-            FROM pg_catalog.pg_class c
-            JOIN information_schema.columns cols
-            ON c.relname = cols.table_name
-            WHERE c.relname = :table_name AND cols.column_name = :column_name
-        """
-    }.get(db_type.lower(), "")
-
-    try:
-        # 获取所有表名
-        all_tables = inspector.get_table_names()
-
-        # 如果指定了table_names，则过滤表名
-        target_tables = all_tables
-
-        if table_names:
-            target_tables = [table.strip() for table in table_names.split(',')]
-            # 过滤出实际存在的表
-            target_tables = [table for table in target_tables if table in all_tables]
-        print(f"Retrieving table metadata for {len(target_tables)} tables...")
-        for table_name in target_tables:
-            # 获取表注释
-            table_comment = ""
-            try:
-                table_comment = inspector.get_table_comment(table_name).get("text") or ""
-            except SQLAlchemyError as e:
-                raise ValueError(f"Failed to retrieve table comments: {str(e)}")
-
-            table_info = {
-                'comment': table_comment,
-                'columns': []
-            }
-
-            for column in inspector.get_columns(table_name):
-                # 获取字段注释
-                column_comment = ""
-                try:
-                    with engine.connect() as conn:
-                        stmt = text(column_comment_sql)
-                        column_comment = conn.execute(stmt, {
-                            'table_name': table_name,
-                            'column_name': column['name']
-                        }).scalar() or ""
-                except SQLAlchemyError as e:
-                    print(f"Warning: failed to get comment for {table_name}.{column['name']} - {e}")
-                    column_comment = ""
-
-                table_info['columns'].append({
-                    'name': column['name'],
-                    'comment': column_comment,
-                    'type': str(column['type'])
-                })
-
-            result[table_name] = table_info
-        return result
-    except SQLAlchemyError as e:
-        raise ValueError(f"Failed to retrieve database table metadata: {str(e)}")
-    finally:
-        engine.dispose()
+#def get_db_schema(
+#        db_type: str,
+#        host: str,
+#        port: int,
+#        database: str,
+#        username: str,
+#        password: str,
+#        table_names: str | None = None
+#) -> dict[str, Any] | None:
+#    """
+#    获取数据库表结构信息
+#    :param db_type: 数据库类型 (mysql/oracle/sqlserver/postgresql)
+#    :param host: 主机地址
+#    :param port: 端口号
+#    :param database: 数据库名
+#    :param username: 用户名
+#    :param password: 密码
+#    :param table_names: 要查询的表名，以逗号分隔的字符串，如果为None则查询所有表
+#    :return: 包含所有表结构信息的字典
+#    """
+#    result: dict[str, Any] = {}
+#    db_type = 'mssql' if db_type == 'sqlserver' else db_type
+#    # 构建连接URL
+#    driver = {
+#        'mysql': 'pymysql',
+#        'oracle': 'cx_oracle',
+#        'mssql': 'pyodbc',
+#        'postgresql': 'psycopg2'
+#    }.get(db_type.lower(), '')
+#
+#    encoded_username = quote_plus(username)
+#    encoded_password = quote_plus(password)
+#    # separator = ':' if db_type is 'mssql' else ','
+#
+#    engine = create_engine(f'{db_type.lower()}+{driver}://{encoded_username}:{encoded_password}@{host}{separator}{port}/{database}')
+#    inspector = inspect(engine)
+#
+#    # 获取字段注释的SQL语句
+#    column_comment_sql = {
+#        'mysql': f"SELECT COLUMN_COMMENT FROM information_schema.COLUMNS WHERE TABLE_SCHEMA = '{database}' AND TABLE_NAME = :table_name AND COLUMN_NAME = :column_name",
+#        'oracle': "SELECT COMMENTS FROM ALL_COL_COMMENTS WHERE TABLE_NAME = :table_name AND COLUMN_NAME = :column_name",
+#        'mssql': "SELECT CAST(ep.value AS NVARCHAR(MAX)) FROM sys.columns c LEFT JOIN sys.extended_properties ep ON ep.major_id = c.object_id AND ep.minor_id = c.column_id WHERE OBJECT_NAME(c.object_id) = :table_name AND c.name = :column_name",
+#        'postgresql': """
+#            SELECT pg_catalog.col_description(c.oid, cols.ordinal_position::int)
+#            FROM pg_catalog.pg_class c
+#            JOIN information_schema.columns cols
+#            ON c.relname = cols.table_name
+#            WHERE c.relname = :table_name AND cols.column_name = :column_name
+#        """
+#    }.get(db_type.lower(), "")
+#
+#    try:
+#        # 获取所有表名
+#        all_tables = inspector.get_table_names()
+#
+#        # 如果指定了table_names，则过滤表名
+#        target_tables = all_tables
+#
+#        if table_names:
+#            target_tables = [table.strip() for table in table_names.split(',')]
+#            # 过滤出实际存在的表
+#            target_tables = [table for table in target_tables if table in all_tables]
+#        print(f"Retrieving table metadata for {len(target_tables)} tables...")
+#        for table_name in target_tables:
+#            # 获取表注释
+#            table_comment = ""
+#            try:
+#                table_comment = inspector.get_table_comment(table_name).get("text") or ""
+#            except SQLAlchemyError as e:
+#                raise ValueError(f"Failed to retrieve table comments: {str(e)}")
+#
+#            table_info = {
+#                'comment': table_comment,
+#                'columns': []
+#            }
+#
+#            for column in inspector.get_columns(table_name):
+#                # 获取字段注释
+#                column_comment = ""
+#                try:
+#                    with engine.connect() as conn:
+#                        stmt = text(column_comment_sql)
+#                        column_comment = conn.execute(stmt, {
+#                            'table_name': table_name,
+#                            'column_name': column['name']
+#                        }).scalar() or ""
+#                except SQLAlchemyError as e:
+#                    print(f"Warning: failed to get comment for {table_name}.{column['name']} - {e}")
+#                    column_comment = ""
+#
+#                table_info['columns'].append({
+#                    'name': column['name'],
+#                    'comment': column_comment,
+#                    'type': str(column['type'])
+#                })
+#
+#            result[table_name] = table_info
+#        return result
+#    except SQLAlchemyError as e:
+#        raise ValueError(f"Failed to retrieve database table metadata: {str(e)}")
+#    finally:
+#        engine.dispose()
 
 def format_schema_dsl(schema: dict[str, Any], with_type: bool = True, with_comment: bool = False) -> str:
     """
@@ -168,18 +168,18 @@ def execute_sql(
     encoded_username = quote_plus(username)
     encoded_password = quote_plus(password)
     connect_args = {}
-    driver_extra_info = None
     # PostgreSQL 特殊处理
     if db_type.lower() == 'postgresql' and schema:
         connect_args['options'] = f"-c search_path={schema}"
 
-    if db_type.lower() == 'sqlserver':
-        import os
-        driver_extra_info = 'ODBC+Driver+17+for+SQL+Server' if os.name == 'posix' else 'SQL Server'
+    #if db_type.lower() == 'sqlserver':
+    #    import os
+    #    driver_extra_info = 'ODBC+Driver+17+for+SQL+Server' if os.name == 'posix' else 'SQL Server'
+    #    print(driver_extra_info)
     # 构建连接字符串
     connection_uri = _build_connection_uri(
         db_type, driver, encoded_username, encoded_password,
-        host, port, database, driver_extra_info
+        host, port, database
     )
 
     try:
@@ -204,7 +204,7 @@ def _get_driver(db_type: str) -> str:
     drivers = {
         'mysql': 'pymysql',
         'oracle': 'cx_oracle',
-        'sqlserver': 'pyodbc',
+        'sqlserver': 'pymssql',
         'postgresql': 'psycopg2'
     }
     return drivers.get(db_type.lower(), '')
@@ -217,17 +217,21 @@ def _build_connection_uri(
     host: str,
     port: int,
     database: str,
-    driver_info: str
 ) -> str:
     """构建数据库连接字符串"""
+<<<<<<< HEAD
     separator = ':' if db_type != 'sqlserver' else ','
     db_type = db_type if db_type != 'sqlserver' else 'mssql'
     extra_info = '' if driver_info == None else f'?driver={driver_info}'
 
     return f"{db_type}+{driver}://{username}:{password}@{host}{separator}{port}/{database}{extra_info}"
+=======
+    db_type = db_type if db_type != 'sqlserver' else 'mssql'
+    return f"{db_type}+{driver}://{username}:{password}@{host}:{port}/{database}"
+>>>>>>> debug
 
 def _process_result(result_proxy) -> Union[list[dict], dict, None]:
     """处理执行结果"""
     if result_proxy.returns_rows:
         return [dict(row._mapping) for row in result_proxy]
-    return {"rowcount": result_proxy.rowcount}
+    return {"rowcount": result_proxy.rowcount}