From d7876255a9d1e275defe02368a8fece5c9ea641a Mon Sep 17 00:00:00 2001 From: Guy Korland Date: Thu, 13 Aug 2026 21:09:22 +0300 Subject: [PATCH 01/10] feat(loaders): add Microsoft SQL Server support Adds a `sqlserver://` loader so QueryWeaver can introspect Microsoft SQL Server and Azure SQL instances and answer natural-language questions against them. Rebuilt on top of current staging and reworked to address the review findings on #538. - api/loaders/sqlserver_loader.py: new pymssql-based loader. Connections use `as_dict=True`, so rows are read by column name; positional access raises KeyError with that setting. - Schema scoping: `parse_schema_from_url` reads `?schema=` (default `dbo`). All catalog queries join `sys.schemas` and bind the schema as a parameter, and sample queries are schema-qualified, so same-named tables in other schemas can no longer collide. - Identifier quoting: `quote_ident` doubles a literal `]` so it cannot terminate a bracket delimiter early. - Connections are released in `finally` via `_close_quietly` / `_rollback_quietly` instead of `if 'conn' in locals()`. - api/core/pipeline.py: dispatch `sqlserver://` with an `sdk_only` guard and a lazy import, and map `sqlserver`/`mssql` to the `tsql` sqlglot dialect. Without the mapping the fail-closed destructive-operation guard classified ordinary reads such as `SELECT TOP 10 ...` as destructive. - api/core/schema_loader.py: accept the `sqlserver://` scheme. - api/sql_utils/sql_sanitizer.py: "already quoted" is now dialect-scoped, so `[weird]` is still quoted on PostgreSQL/MySQL where brackets are data; `get_quote_char` returns `[` for sqlserver/mssql. - pyproject.toml: pymssql lives in the `server` extra, not core deps, so the published SDK wheel is unaffected. - DatabaseModal.tsx: replace the nested protocol/port/placeholder ternaries with a `DB_PROFILES` map and expose the SQL Server option and its schema field. - tests: new `tests/test_sqlserver_loader.py` uses fakes that mimic pymssql dict rows, so the cursor contract is actually exercised; added T-SQL dialect and bracket-quoting regression tests. - docs/sqlserver_loader.md and README updated. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .github/wordlist.txt | 8 + .gitignore | 1 + README.md | 6 +- api/core/pipeline.py | 18 +- api/core/schema_loader.py | 4 +- api/loaders/sqlserver_loader.py | 715 ++++++++++++++++++++ api/sql_utils/sql_sanitizer.py | 52 +- app/src/components/modals/DatabaseModal.tsx | 70 +- docs/sqlserver_loader.md | 144 ++++ pyproject.toml | 1 + tests/test_destructive_detection.py | 34 + tests/test_sql_sanitizer.py | 55 +- tests/test_sqlserver_loader.py | 493 ++++++++++++++ uv.lock | 33 + 14 files changed, 1597 insertions(+), 37 deletions(-) create mode 100644 api/loaders/sqlserver_loader.py create mode 100644 docs/sqlserver_loader.md create mode 100644 tests/test_sqlserver_loader.py diff --git a/.github/wordlist.txt b/.github/wordlist.txt index 9dce194d..96b2ff25 100644 --- a/.github/wordlist.txt +++ b/.github/wordlist.txt @@ -124,3 +124,11 @@ SDK Dependabot PyPI pypi +pymssql +sqlserver +SQLServerLoader +dbo +tsql +hostname +TLS +sqlglot diff --git a/.gitignore b/.gitignore index 3bf6a813..7964d0ed 100644 --- a/.gitignore +++ b/.gitignore @@ -35,3 +35,4 @@ e2e/.auth/ # Build artifacts clients/python/queryweaver_client.egg-info/ clients/ts/dist/ +wordlist.dic diff --git a/README.md b/README.md index fa295322..c45bb846 100644 --- a/README.md +++ b/README.md @@ -266,7 +266,7 @@ async def main(): # Initialize with FalkorDB connection qw = QueryWeaver(falkordb_url="redis://localhost:6379") - # Connect a PostgreSQL or MySQL database + # Connect a PostgreSQL, MySQL, SQL Server or Snowflake database conn = await qw.connect_database("postgresql://user:pass@host:5432/mydb") print(f"Connected: {conn.database_id}") # "mydb" @@ -310,7 +310,7 @@ async with QueryWeaver(falkordb_url="redis://host-a:6379", user_id="tenant_a") a | Method | Description | |--------|-------------| -| `connect_database(db_url)` | Connect PostgreSQL/MySQL and load schema | +| `connect_database(db_url)` | Connect PostgreSQL/MySQL/SQL Server/Snowflake and load schema | | `query(database, question)` | Convert natural language to SQL and execute | | `get_schema(database)` | Retrieve database schema (tables and relationships) | | `list_databases()` | List all connected databases | @@ -356,7 +356,7 @@ if result.requires_confirmation: - Python 3.12+ - FalkorDB instance (local or remote) - OpenAI or Azure OpenAI API key (for LLM) -- Target SQL database (PostgreSQL or MySQL) +- Target SQL database (PostgreSQL, MySQL, SQL Server or Snowflake) ## Development diff --git a/api/core/pipeline.py b/api/core/pipeline.py index 2aca3998..5ba8d426 100644 --- a/api/core/pipeline.py +++ b/api/core/pipeline.py @@ -115,8 +115,9 @@ def get_database_type_and_loader( PostgreSQL for backward compatibility on the server path. When ``sdk_only`` is True, raises ``InvalidArgumentError`` for vendors - that need the ``[server]`` extra (snowflake) or for unknown URL schemes, - so SDK callers get a clean error instead of a deferred ``ImportError``. + that need the ``[server]`` extra (snowflake, sqlserver) or for unknown URL + schemes, so SDK callers get a clean error instead of a deferred + ``ImportError``. """ if not db_url or db_url == "No URL available for this database.": return None, None @@ -138,6 +139,17 @@ def get_database_type_and_loader( # pylint: disable=import-outside-toplevel from api.loaders.snowflake_loader import SnowflakeLoader return 'snowflake', SnowflakeLoader + if db_url_lower.startswith('sqlserver://'): + if sdk_only: + raise InvalidArgumentError( + "SQL Server requires the [server] extra: " + "pip install queryweaver[server]" + ) + # Lazy-import: pymssql is in the [server] extra, not in the core SDK + # install. + # pylint: disable=import-outside-toplevel + from api.loaders.sqlserver_loader import SQLServerLoader + return 'sqlserver', SQLServerLoader if sdk_only: raise InvalidArgumentError( @@ -205,6 +217,8 @@ def truncate_for_log(query: str, max_length: int = 200) -> str: "postgres": "postgres", "mysql": "mysql", "snowflake": "snowflake", + "sqlserver": "tsql", + "mssql": "tsql", } # sqlglot expression class names that represent a write, DDL, privilege change, diff --git a/api/core/schema_loader.py b/api/core/schema_loader.py index edb44d6c..21d13f78 100644 --- a/api/core/schema_loader.py +++ b/api/core/schema_loader.py @@ -32,7 +32,9 @@ def _step_start(steps_counter: int) -> dict[str, str]: "message": f"Step {steps_counter}: Starting database connection", } -_KNOWN_DB_SCHEMES = ("postgresql://", "postgres://", "mysql://", "snowflake://") +_KNOWN_DB_SCHEMES = ( + "postgresql://", "postgres://", "mysql://", "snowflake://", "sqlserver://", +) def _step_detect_db_type(steps_counter: int, url: str) -> tuple[type[BaseLoader], dict[str, str]]: diff --git a/api/loaders/sqlserver_loader.py b/api/loaders/sqlserver_loader.py new file mode 100644 index 00000000..239f4943 --- /dev/null +++ b/api/loaders/sqlserver_loader.py @@ -0,0 +1,715 @@ +"""SQL Server loader for loading database schemas into FalkorDB graphs.""" + +import datetime +import decimal +import logging +import re +from typing import AsyncGenerator, Dict, Any, List, Tuple +from urllib.parse import urlparse, parse_qs, unquote + +import tqdm +import pymssql + +from api.loaders.base_loader import BaseLoader +from api.loaders.graph_loader import load_to_graph + +logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") + +DEFAULT_SCHEMA = "dbo" +DEFAULT_PORT = 1433 + + +class SQLServerQueryError(Exception): + """Exception raised for SQL Server query execution errors.""" + + +class SQLServerConnectionError(Exception): + """Exception raised for SQL Server connection errors.""" + + +def quote_ident(identifier: str) -> str: + """Bracket-quote a T-SQL identifier, escaping any embedded ``]``. + + SQL Server escapes a closing bracket inside a delimited identifier by + doubling it, so ``my]table`` must become ``[my]]table]``. Without this a + crafted identifier would terminate the quote early. + + Args: + identifier: Raw identifier as read from the system catalog. + + Returns: + The bracket-quoted identifier. + """ + return f"[{identifier.replace(']', ']]')}]" + + +class SQLServerLoader(BaseLoader): + """ + Loader for SQL Server databases that connects and extracts schema information. + """ + + # DDL operations that modify database schema # pylint: disable=duplicate-code + SCHEMA_MODIFYING_OPERATIONS = { + 'CREATE', 'ALTER', 'DROP', 'RENAME', 'TRUNCATE' + } + + # More specific patterns for schema-affecting operations + SCHEMA_PATTERNS = [ # pylint: disable=duplicate-code + r'^\s*CREATE\s+TABLE', + r'^\s*CREATE\s+INDEX', + r'^\s*CREATE\s+UNIQUE\s+INDEX', + r'^\s*ALTER\s+TABLE', + r'^\s*DROP\s+TABLE', + r'^\s*DROP\s+INDEX', + r'^\s*RENAME\s+TABLE', + r'^\s*TRUNCATE\s+TABLE', + r'^\s*CREATE\s+VIEW', + r'^\s*DROP\s+VIEW', + r'^\s*CREATE\s+SCHEMA', + r'^\s*DROP\s+SCHEMA', + ] + + @staticmethod + def _execute_sample_query( + cursor, table_name: str, col_name: str, sample_size: int = 3 + ) -> List[Any]: + """ + Execute query to get random sample values for a column. + SQL Server implementation using TOP with NEWID() for random sampling. + + ``table_name`` may be schema-qualified (``schema.table``); each part is + bracket-quoted separately so the schema prefix survives. + """ + schema, _, bare_table = table_name.rpartition('.') + qualified = quote_ident(bare_table) + if schema: + qualified = f"{quote_ident(schema)}.{qualified}" + + col = quote_ident(col_name) + # ``sample_size`` is coerced to int; identifiers are bracket-quoted with + # ``]`` escaped, since T-SQL cannot bind identifiers as parameters. + query = ( + f"SELECT DISTINCT TOP {int(sample_size)} {col}" + f" FROM {qualified}" + f" WHERE {col} IS NOT NULL" + f" ORDER BY NEWID()" + ) + cursor.execute(query) + + # The cursor is opened with ``as_dict=True`` so rows are keyed by column + # name only — pymssql's ``row2dict`` strips positional keys. + sample_results = cursor.fetchall() + return [row[col_name] for row in sample_results if row[col_name] is not None] + + @staticmethod + def _serialize_value(value): + """ + Convert non-JSON serializable values to JSON serializable format. + + Args: + value: The value to serialize + + Returns: + JSON serializable version of the value + """ + if isinstance(value, (datetime.date, datetime.datetime)): + return value.isoformat() + if isinstance(value, datetime.time): + return value.isoformat() + if isinstance(value, decimal.Decimal): + return float(value) + if isinstance(value, bytes): + return value.hex() + if value is None: + return None + return value + + @staticmethod + def parse_schema_from_url(connection_url: str) -> str: + """ + Parse the target schema from the connection URL's ``schema`` parameter. + + Expected format: + ``sqlserver://user:pass@host:port/database?schema=schema_name`` + + Args: + connection_url: SQL Server connection URL + + Returns: + The requested schema, or ``dbo`` when not specified. + """ + try: + parsed = urlparse(connection_url) + schema = parse_qs(parsed.query).get('schema', [''])[0] + return unquote(schema).strip() or DEFAULT_SCHEMA + except (ValueError, AttributeError): + return DEFAULT_SCHEMA + + @staticmethod + def _parse_sqlserver_url(connection_url: str) -> Dict[str, Any]: + """ + Parse SQL Server connection URL into connection parameters. + + Args: + connection_url: SQL Server connection URL in format: + sqlserver://user:password@host:port/database + + Returns: + Dict with connection parameters accepted by ``pymssql.connect``. + + Raises: + ValueError: If the URL is malformed. + """ + if not connection_url.lower().startswith('sqlserver://'): + raise ValueError( + "Invalid SQL Server URL format. Expected " + "sqlserver://user:password@host:port/database" + ) + + parsed = urlparse(connection_url) + + if not parsed.hostname: + raise ValueError("SQL Server URL must include a host") + + database = unquote(parsed.path or '').lstrip('/') + if not database: + raise ValueError("SQL Server URL must include database name") + + if not parsed.username: + raise ValueError("SQL Server URL must include username and host") + + params: Dict[str, Any] = { + 'server': parsed.hostname, + 'port': parsed.port or DEFAULT_PORT, + 'user': unquote(parsed.username), + 'password': unquote(parsed.password) if parsed.password else "", + 'database': database, + } + + # Opt-in transport encryption: ``?encrypt=true`` maps to FreeTDS' TLS + # negotiation. Left unset otherwise to preserve driver defaults. + encrypt = parse_qs(parsed.query).get('encrypt', [''])[0].strip().lower() + if encrypt in ('true', '1', 'yes', 'require'): + params['encryption'] = 'require' + elif encrypt in ('false', '0', 'no', 'off'): + params['encryption'] = 'off' + + return params + + @staticmethod + async def load( # pylint: disable=arguments-differ + prefix: str, + connection_url: str, + db=None, + ) -> AsyncGenerator[tuple[bool, str], None]: + """ + Load the graph data from a SQL Server database into the graph database. + + Args: + prefix: Graph name prefix (typically the user id). + connection_url: SQL Server connection URL in format: + sqlserver://user:password@host:port/database + db: Optional FalkorDB handle; falls back to the server singleton. + + Yields: + Tuple[bool, str]: Success status and message + """ + conn = None + cursor = None + try: + # Parse connection URL + conn_params = SQLServerLoader._parse_sqlserver_url(connection_url) + schema = SQLServerLoader.parse_schema_from_url(connection_url) + + # Connect to SQL Server database + conn = pymssql.connect(**conn_params) # pylint: disable=no-member + cursor = conn.cursor(as_dict=True) + + # Get database name + db_name = conn_params['database'] + + # Get all table information + yield True, "Extracting table information..." + entities = SQLServerLoader.extract_tables_info(cursor, schema) + + # Get all relationship information + yield True, "Extracting relationship information..." + relationships = SQLServerLoader.extract_relationships(cursor, schema) + + # Close database connection + cursor.close() + cursor = None + conn.close() + conn = None + + # Load data into graph + yield True, "Loading data into graph..." + await load_to_graph(f"{prefix}_{db_name}", entities, relationships, + db_name=db_name, db_url=connection_url, db=db) + + yield True, (f"SQL Server schema loaded successfully. " + f"Found {len(entities)} tables.") + + except pymssql.Error as e: + logging.error("SQL Server connection error: %s", e) + yield False, "Failed to connect to SQL Server database" + except Exception as e: # pylint: disable=broad-exception-caught + logging.error("Error loading SQL Server schema: %s", e) + yield False, "Failed to load SQL Server database schema" + finally: + SQLServerLoader._close_quietly(cursor, conn) + + @staticmethod + def _close_quietly(cursor, conn) -> None: + """Close *cursor* and *conn* if still open, ignoring teardown errors.""" + for handle in (cursor, conn): + if handle is None: + continue + try: + handle.close() + except Exception: # pylint: disable=broad-exception-caught + logging.debug("Ignoring error while closing SQL Server handle", exc_info=True) + + @staticmethod + def extract_tables_info(cursor, schema: str = DEFAULT_SCHEMA) -> Dict[str, Any]: + """ + Extract table and column information from a SQL Server schema. + + Args: + cursor: Database cursor + schema: Schema to extract tables from (default: ``dbo``) + + Returns: + Dict containing table information + """ + entities = {} + + # Get all tables in the requested schema + cursor.execute(""" + SELECT + t.name AS table_name, + ISNULL(CAST(ep.value AS NVARCHAR(MAX)), '') AS table_comment + FROM sys.tables t + JOIN sys.schemas s ON t.schema_id = s.schema_id + LEFT JOIN sys.extended_properties ep + ON ep.major_id = t.object_id + AND ep.minor_id = 0 + AND ep.class = 1 + AND ep.name = 'MS_Description' + WHERE t.is_ms_shipped = 0 + AND s.name = %s + ORDER BY t.name; + """, (schema,)) + + tables = cursor.fetchall() + + for table_info in tqdm.tqdm(tables, desc="Extracting table information"): + table_name = table_info['table_name'] + table_comment = table_info['table_comment'] + + # Get column information for this table + columns_info = SQLServerLoader.extract_columns_info(cursor, schema, table_name) + + # Get foreign keys for this table + foreign_keys = SQLServerLoader.extract_foreign_keys(cursor, schema, table_name) + + # Generate table description + table_description = table_comment if table_comment else f"Table: {table_name}" + + # Get column descriptions for batch embedding + col_descriptions = [col_info['description'] for col_info in columns_info.values()] + + entities[table_name] = { + 'description': table_description, + 'columns': columns_info, + 'foreign_keys': foreign_keys, + 'col_descriptions': col_descriptions + } + + return entities + + @staticmethod + def extract_columns_info(cursor, schema: str, table_name: str) -> Dict[str, Any]: + """ + Extract column information for a specific table. + + Args: + cursor: Database cursor + schema: Schema owning the table + table_name: Name of the table + + Returns: + Dict containing column information + """ + cursor.execute(""" + SELECT + c.name AS column_name, + tp.name AS data_type, + c.is_nullable, + dc.definition AS column_default, + CASE + WHEN pk.column_id IS NOT NULL THEN 'PRI' + WHEN fk.parent_column_id IS NOT NULL THEN 'MUL' + WHEN uc.column_id IS NOT NULL THEN 'UNI' + ELSE '' + END AS column_key, + ISNULL(CAST(ep.value AS NVARCHAR(MAX)), '') AS column_comment + FROM sys.columns c + JOIN sys.types tp ON c.user_type_id = tp.user_type_id + JOIN sys.tables t ON c.object_id = t.object_id + JOIN sys.schemas s ON t.schema_id = s.schema_id + LEFT JOIN sys.default_constraints dc ON c.default_object_id = dc.object_id + LEFT JOIN ( + SELECT ic.object_id, ic.column_id + FROM sys.index_columns ic + JOIN sys.indexes i ON ic.object_id = i.object_id AND ic.index_id = i.index_id + WHERE i.is_primary_key = 1 + ) pk ON c.object_id = pk.object_id AND c.column_id = pk.column_id + LEFT JOIN sys.foreign_key_columns fk + ON fk.parent_object_id = c.object_id AND fk.parent_column_id = c.column_id + LEFT JOIN ( + SELECT ic.object_id, ic.column_id + FROM sys.index_columns ic + JOIN sys.indexes i ON ic.object_id = i.object_id AND ic.index_id = i.index_id + WHERE i.is_unique = 1 AND i.is_primary_key = 0 + ) uc ON c.object_id = uc.object_id AND c.column_id = uc.column_id + LEFT JOIN sys.extended_properties ep + ON ep.major_id = c.object_id + AND ep.minor_id = c.column_id + AND ep.class = 1 + AND ep.name = 'MS_Description' + WHERE s.name = %s AND t.name = %s + ORDER BY c.column_id; + """, (schema, table_name)) + + columns = cursor.fetchall() + columns_info = {} + + for col_info in columns: + col_name = col_info['column_name'] + data_type = col_info['data_type'] + is_nullable = 'YES' if col_info['is_nullable'] else 'NO' + column_default = col_info['column_default'] + column_key = col_info['column_key'] + column_comment = col_info['column_comment'] + + # Determine key type + if column_key == 'PRI': + key_type = 'PRIMARY KEY' + elif column_key == 'MUL': + key_type = 'FOREIGN KEY' + elif column_key == 'UNI': + key_type = 'UNIQUE KEY' + else: + key_type = 'NONE' + + # Generate column description + description_parts = [] + if column_comment: + description_parts.append(str(column_comment)) + else: + description_parts.append(f"Column {col_name} of type {data_type}") + + if key_type != 'NONE': + description_parts.append(f"({key_type})") + + if is_nullable == 'NO': + description_parts.append("(NOT NULL)") + + if column_default is not None: + description_parts.append(f"(Default: {column_default})") + + # Extract sample values for the column (stored separately, not in description) + sample_values = SQLServerLoader.extract_sample_values_for_column( + cursor, f"{schema}.{table_name}", col_name + ) + + columns_info[col_name] = { + 'type': data_type, + 'null': is_nullable, + 'key': key_type, + 'description': ' '.join(description_parts), + 'default': column_default, + 'sample_values': sample_values + } + + return columns_info + + @staticmethod + def extract_foreign_keys(cursor, schema: str, table_name: str) -> List[Dict[str, str]]: + """ + Extract foreign key information for a specific table. + + Args: + cursor: Database cursor + schema: Schema owning the table + table_name: Name of the table + + Returns: + List of foreign key dictionaries + """ + cursor.execute(""" + SELECT + fk.name AS constraint_name, + cp.name AS column_name, + rt.name AS referenced_table_name, + rs.name AS referenced_schema_name, + cr.name AS referenced_column_name + FROM sys.foreign_keys fk + JOIN sys.foreign_key_columns fkc + ON fk.object_id = fkc.constraint_object_id + JOIN sys.columns cp + ON fkc.parent_object_id = cp.object_id + AND fkc.parent_column_id = cp.column_id + JOIN sys.tables rt + ON fkc.referenced_object_id = rt.object_id + JOIN sys.schemas rs ON rt.schema_id = rs.schema_id + JOIN sys.columns cr + ON fkc.referenced_object_id = cr.object_id + AND fkc.referenced_column_id = cr.column_id + JOIN sys.tables pt + ON fkc.parent_object_id = pt.object_id + JOIN sys.schemas ps ON pt.schema_id = ps.schema_id + WHERE ps.name = %s AND pt.name = %s + ORDER BY fk.name; + """, (schema, table_name)) + + foreign_keys = [] + for fk_info in cursor.fetchall(): + foreign_keys.append({ + 'constraint_name': fk_info['constraint_name'], + 'column': fk_info['column_name'], + 'referenced_table': fk_info['referenced_table_name'], + 'referenced_column': fk_info['referenced_column_name'] + }) + + return foreign_keys + + @staticmethod + def extract_relationships( + cursor, schema: str = DEFAULT_SCHEMA + ) -> Dict[str, List[Dict[str, str]]]: + """ + Extract all relationship information from a schema. + + Only foreign keys whose parent *and* referenced tables both live in + *schema* are returned, so relationships always point at entities that + were actually loaded. + + Args: + cursor: Database cursor + schema: Schema to extract relationships from (default: ``dbo``) + + Returns: + Dict containing relationship information + """ + cursor.execute(""" + SELECT + pt.name AS table_name, + fk.name AS constraint_name, + cp.name AS column_name, + rt.name AS referenced_table_name, + cr.name AS referenced_column_name + FROM sys.foreign_keys fk + JOIN sys.foreign_key_columns fkc + ON fk.object_id = fkc.constraint_object_id + JOIN sys.columns cp + ON fkc.parent_object_id = cp.object_id + AND fkc.parent_column_id = cp.column_id + JOIN sys.tables pt + ON fkc.parent_object_id = pt.object_id + JOIN sys.schemas ps ON pt.schema_id = ps.schema_id + JOIN sys.tables rt + ON fkc.referenced_object_id = rt.object_id + JOIN sys.schemas rs ON rt.schema_id = rs.schema_id + JOIN sys.columns cr + ON fkc.referenced_object_id = cr.object_id + AND fkc.referenced_column_id = cr.column_id + WHERE ps.name = %s AND rs.name = %s + ORDER BY pt.name, fk.name; + """, (schema, schema)) + + relationships: Dict[str, List[Dict[str, str]]] = {} + for rel_info in cursor.fetchall(): + constraint_name = rel_info['constraint_name'] + + if constraint_name not in relationships: + relationships[constraint_name] = [] + + relationships[constraint_name].append({ + 'from': rel_info['table_name'], + 'to': rel_info['referenced_table_name'], + 'source_column': rel_info['column_name'], + 'target_column': rel_info['referenced_column_name'], + 'note': f'Foreign key constraint: {constraint_name}' + }) + + return relationships + + @staticmethod + def is_schema_modifying_query(sql_query: str) -> Tuple[bool, str]: + """ + Check if a SQL query modifies the database schema. + + Args: + sql_query: The SQL query to check + + Returns: + Tuple of (is_schema_modifying, operation_type) + """ + if not sql_query or not sql_query.strip(): + return False, "" + + # Clean and normalize the query + normalized_query = sql_query.strip().upper() + + # Check for basic DDL operations + first_word = normalized_query.split()[0] if normalized_query.split() else "" + if first_word in SQLServerLoader.SCHEMA_MODIFYING_OPERATIONS: + # Additional pattern matching for more precise detection + for pattern in SQLServerLoader.SCHEMA_PATTERNS: + if re.match(pattern, normalized_query, re.IGNORECASE): + return True, first_word + + # If it's a known DDL operation but doesn't match specific patterns, + # still consider it schema-modifying (better safe than sorry) + return True, first_word + + return False, "" + + @staticmethod + async def refresh_graph_schema(graph_id: str, db_url: str, db=None) -> Tuple[bool, str]: + """ + Refresh the graph schema by clearing existing data and reloading from the database. + + Args: + graph_id: The graph ID to refresh + db_url: Database connection URL + db: Optional FalkorDB handle; falls back to the server singleton. + + Returns: + Tuple of (success, message) + """ + try: + logging.info("Schema modification detected. Refreshing graph schema.") + + from api.core.db_resolver import resolve_db # pylint: disable=import-outside-toplevel + + # Clear existing graph data + # Drop current graph before reloading + graph = resolve_db(db).select_graph(graph_id) + await graph.delete() + + # Extract prefix from graph_id (remove database name part) + # graph_id format is typically "prefix_database_name" + parts = graph_id.split('_') + if len(parts) >= 2: + # Reconstruct prefix by joining all parts except the last one + prefix = '_'.join(parts[:-1]) + else: + prefix = graph_id + + # Reuse the existing load method to reload the schema + success, message = False, "" + async for progress in SQLServerLoader.load(prefix, db_url, db=db): + success, message = progress + + if success: + logging.info("Graph schema refreshed successfully.") + return True, message + + logging.error("Schema refresh failed") + return False, "Failed to reload schema" + + except Exception as e: # pylint: disable=broad-exception-caught + # Log the error and return failure + logging.error("Error refreshing graph schema: %s", str(e)) + error_msg = "Error refreshing graph schema" + logging.error(error_msg) + return False, error_msg + + @staticmethod + def execute_sql_query(sql_query: str, db_url: str) -> List[Dict[str, Any]]: + """ + Execute a SQL query on the SQL Server database and return the results. + + Args: + sql_query: The SQL query to execute + db_url: SQL Server connection URL in format: + sqlserver://user:password@host:port/database + + Returns: + List of dictionaries containing the query results + + Raises: + SQLServerQueryError: If the query fails. + """ + conn = None + cursor = None + try: + # Parse connection URL + conn_params = SQLServerLoader._parse_sqlserver_url(db_url) + + # Connect to SQL Server database + conn = pymssql.connect(**conn_params) # pylint: disable=no-member + cursor = conn.cursor(as_dict=True) + + # Execute the SQL query + cursor.execute(sql_query) + + # Check if the query returns results (SELECT queries) + if cursor.description is not None: + # This is a SELECT query or similar that returns rows + results = cursor.fetchall() + result_list = [] + for row in results: + # Serialize each value to ensure JSON compatibility + serialized_row = { + key: SQLServerLoader._serialize_value(value) + for key, value in row.items() + } + result_list.append(serialized_row) + else: + # This is an INSERT, UPDATE, DELETE, or other non-SELECT query + # Return information about the operation + affected_rows = cursor.rowcount + sql_type = sql_query.strip().split()[0].upper() + + if sql_type in ['INSERT', 'UPDATE', 'DELETE']: + result_list = [{ + "operation": sql_type, + "affected_rows": affected_rows, + "status": "success" + }] + else: + # For other types of queries (CREATE, DROP, etc.) + result_list = [{ + "operation": sql_type, + "status": "success" + }] + + # Commit the transaction for write operations + conn.commit() + + return result_list + + except pymssql.Error as e: + SQLServerLoader._rollback_quietly(conn) + logging.error("SQL Server query execution error: %s", e) + raise SQLServerQueryError(f"SQL Server query execution error: {str(e)}") from e + except Exception as e: + SQLServerLoader._rollback_quietly(conn) + logging.error("Error executing SQL query: %s", e) + raise SQLServerQueryError(f"Error executing SQL query: {str(e)}") from e + finally: + SQLServerLoader._close_quietly(cursor, conn) + + @staticmethod + def _rollback_quietly(conn) -> None: + """Roll *conn* back if it exists, ignoring rollback failures.""" + if conn is None: + return + try: + conn.rollback() + except Exception: # pylint: disable=broad-exception-caught + logging.debug("Ignoring error during SQL Server rollback", exc_info=True) diff --git a/api/sql_utils/sql_sanitizer.py b/api/sql_utils/sql_sanitizer.py index 6f7d127e..421ee472 100644 --- a/api/sql_utils/sql_sanitizer.py +++ b/api/sql_utils/sql_sanitizer.py @@ -24,20 +24,39 @@ class SQLIdentifierQuoter: 'EXCEPT', 'CASE', 'WHEN', 'THEN', 'ELSE', 'END', 'CAST', 'ASC', 'DESC' } + @staticmethod + def _is_already_quoted(identifier: str, quote_char: str = '"') -> bool: + """Check if an identifier is already quoted for the active dialect. + + The pair is scoped to *quote_char* so that a bracketed identifier is + only treated as pre-quoted for SQL Server; on PostgreSQL/MySQL a name + such as ``[weird]`` is data, not a delimiter, and must still be quoted. + + Args: + identifier: The identifier to inspect. + quote_char: Opening delimiter of the active dialect. + + Returns: + True if *identifier* is already delimited. + """ + if quote_char == '[': + return identifier.startswith('[') and identifier.endswith(']') + return identifier.startswith(quote_char) and identifier.endswith(quote_char) + @classmethod - def needs_quoting(cls, identifier: str) -> bool: + def needs_quoting(cls, identifier: str, quote_char: str = '"') -> bool: """ Check if an identifier needs quoting based on special characters. - + Args: identifier: The table or column name to check - + quote_char: Quote character of the active dialect + Returns: True if the identifier needs quoting, False otherwise """ # Already quoted - if (identifier.startswith('"') and identifier.endswith('"')) or \ - (identifier.startswith('`') and identifier.endswith('`')): + if cls._is_already_quoted(identifier, quote_char): return False # Check if it's a SQL keyword @@ -51,21 +70,28 @@ def needs_quoting(cls, identifier: str) -> bool: def quote_identifier(identifier: str, quote_char: str = '"') -> str: """ Quote an identifier if not already quoted. - + Args: identifier: The identifier to quote - quote_char: The quote character to use (default: " for PostgreSQL/standard SQL) - + quote_char: The quote character to use (default: " for PostgreSQL/standard SQL, + use ` for MySQL, [ for SQL Server) + Returns: Quoted identifier """ identifier = identifier.strip() # Don't double-quote - if (identifier.startswith('"') and identifier.endswith('"')) or \ - (identifier.startswith('`') and identifier.endswith('`')): + if SQLIdentifierQuoter._is_already_quoted(identifier, quote_char): return identifier + # SQL Server uses bracket pairs: [identifier]. A literal ``]`` inside the + # name is escaped by doubling it, otherwise it would close the delimiter + # early and change the meaning of the statement. + if quote_char == '[': + escaped = identifier.replace(']', ']]') + return f'[{escaped}]' + return f'{quote_char}{identifier}{quote_char}' @classmethod @@ -130,7 +156,7 @@ def auto_quote_identifiers( # For each table that needs quoting for table in query_tables: # Check if this table exists in known schema and needs quoting - if table in known_tables and cls.needs_quoting(table): + if table in known_tables and cls.needs_quoting(table, quote_char): # Quote the table name quoted = cls.quote_identifier(table, quote_char) @@ -167,5 +193,7 @@ def get_quote_char(db_type: str) -> str: """ if db_type.lower() in ['mysql', 'mariadb']: return '`' - # PostgreSQL, SQLite, SQL Server (standard SQL) use double quotes + if db_type.lower() in ['sqlserver', 'mssql']: + return '[' + # PostgreSQL, SQLite use double quotes (standard SQL) return '"' diff --git a/app/src/components/modals/DatabaseModal.tsx b/app/src/components/modals/DatabaseModal.tsx index e3a7f47a..a7158e29 100644 --- a/app/src/components/modals/DatabaseModal.tsx +++ b/app/src/components/modals/DatabaseModal.tsx @@ -20,6 +20,36 @@ interface ConnectionStep { status: 'pending' | 'success' | 'error'; } +/** + * Per-vendor connection defaults. Keeping these in one map means adding a new + * database only requires a single entry rather than editing several ternaries. + */ +const DB_PROFILES = { + postgresql: { + protocol: 'postgresql', + port: '5432', + urlPlaceholder: 'postgresql://user:password@host:5432/database', + }, + mysql: { + protocol: 'mysql', + port: '3306', + urlPlaceholder: 'mysql://user:password@host:3306/database', + }, + sqlserver: { + protocol: 'sqlserver', + port: '1433', + urlPlaceholder: 'sqlserver://user:password@host:1433/database', + }, +} as const satisfies Record; + +type DbProfileKey = keyof typeof DB_PROFILES; + +const snowflakeUrlPlaceholder = + 'snowflake://user:password@account/database/schema?warehouse=warehouse_name'; + +const getDbProfile = (dbType: string) => + DB_PROFILES[dbType as DbProfileKey] ?? DB_PROFILES.postgresql; + const DatabaseModal = ({ open, onOpenChange }: DatabaseModalProps) => { const [connectionMode, setConnectionMode] = useState<'url' | 'manual'>('url'); const [selectedDatabase, setSelectedDatabase] = useState(""); @@ -134,17 +164,21 @@ const DatabaseModal = ({ open, onOpenChange }: DatabaseModalProps) => { builtUrl.searchParams.set('warehouse', warehouse); dbUrl = builtUrl.toString(); } else { - const protocol = selectedDatabase === 'mysql' ? 'mysql' : 'postgresql'; - const builtUrl = new URL(`${protocol}://${host}:${port}/${database}`); + const profile = getDbProfile(selectedDatabase); + const builtUrl = new URL(`${profile.protocol}://${host}:${port}/${database}`); builtUrl.username = username; builtUrl.password = password; - // Append schema option for PostgreSQL if provided - if (selectedDatabase === 'postgresql' && schema.trim()) { + // Append the schema for the vendors that support selecting one + if ((selectedDatabase === 'postgresql' || selectedDatabase === 'sqlserver') && schema.trim()) { if (/[^a-zA-Z0-9_]/.test(schema.trim())) { throw new Error('Schema name can only contain letters, digits, and underscores'); } - builtUrl.searchParams.set('options', `-csearch_path=${schema.trim()}`); + if (selectedDatabase === 'postgresql') { + builtUrl.searchParams.set('options', `-csearch_path=${schema.trim()}`); + } else { + builtUrl.searchParams.set('schema', schema.trim()); + } } dbUrl = builtUrl.toString(); @@ -301,7 +335,7 @@ const DatabaseModal = ({ open, onOpenChange }: DatabaseModalProps) => { Connect to Database - Connect to PostgreSQL, MySQL, or Snowflake database using a connection URL or manual entry.{" "} + Connect to PostgreSQL, MySQL, Snowflake, or SQL Server database using a connection URL or manual entry.{" "} { Snowflake + +
+
+ SQL Server +
+ @@ -381,11 +421,9 @@ const DatabaseModal = ({ open, onOpenChange }: DatabaseModalProps) => { id="connection-url" data-testid="connection-url-input" placeholder={ - selectedDatabase === 'postgresql' - ? 'postgresql://username:password@host:5432/database' - : selectedDatabase === 'mysql' - ? 'mysql://username:password@host:3306/database' - : 'snowflake://username:password@account/database/schema?warehouse=warehouse_name' + selectedDatabase === 'snowflake' + ? snowflakeUrlPlaceholder + : getDbProfile(selectedDatabase).urlPlaceholder } value={connectionUrl} onChange={(e) => setConnectionUrl(e.target.value)} @@ -551,7 +589,7 @@ const DatabaseModal = ({ open, onOpenChange }: DatabaseModalProps) => { setPort(e.target.value)} className="bg-muted border-border focus-visible:ring-purple-500" @@ -592,8 +630,8 @@ const DatabaseModal = ({ open, onOpenChange }: DatabaseModalProps) => { /> - {/* Schema field - PostgreSQL only */} - {selectedDatabase === 'postgresql' && ( + {/* Schema field - PostgreSQL and SQL Server */} + {(selectedDatabase === 'postgresql' || selectedDatabase === 'sqlserver') && (
diff --git a/docs/sqlserver_loader.md b/docs/sqlserver_loader.md new file mode 100644 index 00000000..ea6100a7 --- /dev/null +++ b/docs/sqlserver_loader.md @@ -0,0 +1,144 @@ +# SQL Server Loader + +This document describes the Microsoft SQL Server loader implementation in QueryWeaver. + +## Overview + +The SQL Server loader connects to a Microsoft SQL Server (or Azure SQL) instance, +extracts schema information (tables, columns, primary keys, foreign keys and +relationships) and loads it into a graph so it can be used for Text2SQL queries. + +It is built on [`pymssql`](https://pypi.org/project/pymssql/), which ships with the +`server` extra: + +```bash +uv sync --extra server +``` + +## Connection URL Format + +```text +sqlserver://username:password@host:port/database +``` + +### Parameters + +- **username**: SQL Server login +- **password**: password for the login +- **host**: server hostname or IP +- **port**: server port (optional, defaults to `1433`) +- **database**: database to introspect +- **schema**: schema to introspect (optional query parameter, defaults to `dbo`) +- **encrypt**: `true`/`false` to force TLS on the connection (optional query parameter) + +Credentials are percent-decoded, so passwords containing `@`, `/` or `:` are +supported when they are percent-encoded in the URL. + +### Examples + +```text +sqlserver://sa:MyPassw0rd@localhost:1433/AdventureWorks +sqlserver://sa:MyPassw0rd@localhost/AdventureWorks?schema=sales +sqlserver://appuser:s3cr3t@sql.example.com:1433/reporting?schema=dbo&encrypt=true +``` + +## Features + +### Schema Extraction + +- Tables and views in the selected schema +- Columns with data types, nullability, defaults and primary-key flags +- Extended properties (`MS_Description`) used as table and column descriptions +- Foreign keys, including composite keys +- Many-to-many relationships inferred from junction tables + +All catalog queries join `sys.schemas` and bind the schema name as a parameter, so +a connection only ever sees the requested schema. Tables in other schemas are not +extracted and cannot collide with same-named tables in the selected schema. + +### Sample Values + +Sample values are collected per column with a schema-qualified, bracket-quoted +query: + +```sql +SELECT DISTINCT TOP 3 [column_name] +FROM [dbo].[table_name] +WHERE [column_name] IS NOT NULL; +``` + +### Query Execution + +- Executes SQL against the connected database +- Uses T-SQL (`tsql`) as the sqlglot dialect, so `SELECT TOP n`, `[bracketed]` + identifiers and `FOR JSON PATH` parse correctly and are not misclassified by the + destructive-operation guard +- Rolls back and closes the connection on failure + +## Identifier Quoting + +SQL Server delimits identifiers with brackets. A literal `]` inside a name is +escaped by doubling it, so `my]table` becomes `[my]]table]`. This is applied both +in the loader's own catalog/sample queries and in +`api/sql_utils/sql_sanitizer.py`, where `DatabaseSpecificQuoter.get_quote_char` +returns `[` for `sqlserver` and `mssql`. + +## Usage + +### From the Web Interface + +1. Open the "Connect a database" dialog +2. Select **SQL Server** +3. Fill in host, port, database, credentials and (optionally) schema + +### From the API + +```python +import requests + +response = requests.post( + "http://localhost:5000/api/database/connect", + json={"url": "sqlserver://sa:MyPassw0rd@localhost:1433/AdventureWorks"}, +) +print(response.json()) +``` + +## Implementation Details + +### Catalog Queries + +The loader reads from SQL Server system catalog views: + +- `sys.tables` / `sys.views` joined with `sys.schemas` — table list +- `sys.columns` joined with `sys.types` — column metadata +- `sys.indexes` / `sys.index_columns` — primary keys +- `sys.foreign_keys` / `sys.foreign_key_columns` — foreign keys +- `sys.extended_properties` (with `class = 1`) — table and column descriptions + +### Cursor Contract + +Connections are opened with `as_dict=True`, so `pymssql` returns rows as +dictionaries keyed by column name. Positional access (`row[0]`) raises `KeyError` +with this setting and is never used. + +## Testing + +`tests/test_sqlserver_loader.py` covers: + +- Bracket quoting and `]` escaping (including injection attempts) +- URL parsing: ports, defaults, percent-encoded credentials, `schema` and `encrypt` +- Dict-cursor row access for sample values +- Schema qualification of catalog and sample queries +- Column, foreign-key and relationship mapping +- Value serialization and schema-modification detection +- Query execution: select, non-select, error and connection-failure paths + +```bash +uv run --extra server --extra dev pytest tests/test_sqlserver_loader.py -v +``` + +## Limitations + +- One schema per connection (defaults to `dbo`); connect again to load another +- Requires permission to read the `sys.*` catalog views +- Windows/Azure AD integrated authentication is not supported; use SQL logins diff --git a/pyproject.toml b/pyproject.toml index 42ac313f..cbdb7f05 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -56,6 +56,7 @@ server = [ "fastmcp>=3.4.4,<4.0.0", "graphiti-core>=0.29.1", "snowflake-connector-python>=4.6,<4.8", + "pymssql~=2.3.13", "aiohttp>=3.14.0", ] diff --git a/tests/test_destructive_detection.py b/tests/test_destructive_detection.py index 7ad17f48..3ec1c612 100644 --- a/tests/test_destructive_detection.py +++ b/tests/test_destructive_detection.py @@ -302,3 +302,37 @@ def test_cte_write_confirmation_names_delete(self): message = build_destructive_confirmation_message(sql_type, sql) assert "DELETE" in message assert "DESTRUCTIVE OPERATION DETECTED" in message + + +class TestSQLServerDialect: + """T-SQL is parsed with the tsql dialect, not dialect-agnostically. + + Regression tests: with no dialect mapping, sqlglot could not parse common + T-SQL and the fail-closed path reported ordinary reads as destructive. + """ + + @pytest.mark.parametrize("sql", [ + "SELECT TOP 10 * FROM users", + "SELECT * FROM [my-table]", + "SELECT a, b FROM t FOR JSON PATH", + "SELECT [a b] FROM [dbo].[my-tbl]", + "SELECT ISNULL(name, '') FROM users", + ]) + def test_reads_are_not_destructive(self, sql): + sql_type, is_destructive = detect_destructive_operation(sql, "sqlserver") + assert is_destructive is False + assert sql_type == "SELECT" + + @pytest.mark.parametrize("sql", [ + "DROP TABLE users", + "TRUNCATE TABLE users", + "UPDATE users SET name = 'x'", + "SELECT * INTO backup FROM users", + "EXEC xp_cmdshell 'dir'", + "SELECT 1; DROP TABLE users", + ]) + def test_writes_are_destructive(self, sql): + assert detect_destructive_operation(sql, "sqlserver")[1] is True + + def test_mssql_alias_maps_to_tsql(self): + assert detect_destructive_operation("SELECT TOP 1 * FROM t", "mssql")[1] is False diff --git a/tests/test_sql_sanitizer.py b/tests/test_sql_sanitizer.py index 8937c873..9589666d 100644 --- a/tests/test_sql_sanitizer.py +++ b/tests/test_sql_sanitizer.py @@ -19,9 +19,15 @@ def test_needs_quoting_without_special_chars(self): assert SQLIdentifierQuoter.needs_quoting("OrderItems") is False def test_needs_quoting_already_quoted(self): - """Test that already quoted identifiers don't need quoting again.""" + """Test that already quoted identifiers don't need quoting again. + + "Already quoted" is dialect-scoped: only the active dialect's delimiter + counts, so a backtick pair is pre-quoted for MySQL but is ordinary data + for PostgreSQL. + """ assert SQLIdentifierQuoter.needs_quoting('"table-name"') is False - assert SQLIdentifierQuoter.needs_quoting('`table-name`') is False + assert SQLIdentifierQuoter.needs_quoting('`table-name`', '`') is False + assert SQLIdentifierQuoter.needs_quoting('`table-name`', '"') is True def test_needs_quoting_with_spaces(self): """Test that identifiers with spaces need quoting.""" @@ -41,7 +47,7 @@ def test_quote_identifier(self): def test_quote_identifier_no_double_quote(self): """Test that already quoted identifiers aren't double-quoted.""" assert SQLIdentifierQuoter.quote_identifier('"table-name"') == '"table-name"' - assert SQLIdentifierQuoter.quote_identifier('`table-name`') == '`table-name`' + assert SQLIdentifierQuoter.quote_identifier('`table-name`', '`') == '`table-name`' def test_extract_table_names_from_query(self): """Test extracting table names from SQL queries.""" @@ -231,3 +237,46 @@ def test_real_world_user_comment_scenario(self): assert modified is True assert 'select * from "table-name"' in result.lower() + + +class TestSQLServerQuoting: + """SQL Server bracket-quoting behaviour.""" + + def test_get_quote_char_sqlserver(self): + """SQL Server uses the opening bracket as its quote character.""" + assert DatabaseSpecificQuoter.get_quote_char('sqlserver') == '[' + assert DatabaseSpecificQuoter.get_quote_char('SQLServer') == '[' + assert DatabaseSpecificQuoter.get_quote_char('mssql') == '[' + + def test_quote_identifier_brackets(self): + """Identifiers are wrapped in a bracket pair.""" + assert SQLIdentifierQuoter.quote_identifier('my-table', '[') == '[my-table]' + + def test_quote_identifier_escapes_closing_bracket(self): + """A literal ``]`` is doubled so it cannot terminate the delimiter.""" + assert SQLIdentifierQuoter.quote_identifier('my]table', '[') == '[my]]table]' + + def test_quote_identifier_no_double_quoting(self): + """An already-bracketed identifier is left alone.""" + assert SQLIdentifierQuoter.quote_identifier('[my-table]', '[') == '[my-table]' + + def test_auto_quote_identifiers_sqlserver(self): + """Table names with special characters get bracket-quoted.""" + result, modified = SQLIdentifierQuoter.auto_quote_identifiers( + 'SELECT * FROM user-accounts', {'user-accounts'}, '[' + ) + assert modified is True + assert '[user-accounts]' in result + + def test_bracketed_name_still_quoted_for_postgres(self): + """``[weird]`` is data on PostgreSQL, so it must still be quoted. + + Regression test: a dialect-agnostic bracket pair made this identifier + look pre-quoted and it was emitted unquoted. + """ + assert SQLIdentifierQuoter.needs_quoting('[weird]', '"') is True + assert SQLIdentifierQuoter.quote_identifier('[weird]', '"') == '"[weird]"' + + def test_bracketed_name_treated_as_quoted_for_sqlserver(self): + """The same identifier is already delimited on SQL Server.""" + assert SQLIdentifierQuoter.needs_quoting('[weird]', '[') is False diff --git a/tests/test_sqlserver_loader.py b/tests/test_sqlserver_loader.py new file mode 100644 index 00000000..b7ee3bb2 --- /dev/null +++ b/tests/test_sqlserver_loader.py @@ -0,0 +1,493 @@ +"""Tests for the SQL Server loader. + +These exercise the real introspection code paths against a fake pymssql +cursor rather than mocking the methods under test, so regressions such as +indexing a ``as_dict=True`` row positionally are actually caught. +""" +# pylint: disable=protected-access + +import datetime +import decimal +from unittest.mock import patch, MagicMock + +import pytest + +from api.loaders.sqlserver_loader import ( + SQLServerLoader, + SQLServerQueryError, + quote_ident, +) + + +class FakeCursor: + """Minimal stand-in for a pymssql ``as_dict=True`` cursor. + + Rows are dicts keyed by column name only — matching pymssql's ``row2dict``, + which strips the positional keys. Queries are recorded so tests can assert + on the SQL and the bound parameters. + """ + + def __init__(self, results=None): + # results: list of row-lists returned in order, one per execute() + self._results = list(results or []) + self.executed = [] + self._current = [] + self.description = [("col",)] + self.rowcount = 0 + self.closed = False + + def execute(self, query, params=None): + """Record the statement and pop the next canned result set.""" + self.executed.append((query, params)) + self._current = self._results.pop(0) if self._results else [] + + def fetchall(self): + """Return the result set for the most recent execute().""" + return self._current + + def close(self): + """Mark the cursor closed.""" + self.closed = True + + +class FakeConnection: + """Minimal stand-in for a pymssql connection.""" + + def __init__(self, cursor): + self._cursor = cursor + self.closed = False + self.committed = False + self.rolled_back = False + + def cursor(self, as_dict=False): # pylint: disable=unused-argument + """Return the pre-built fake cursor.""" + return self._cursor + + def commit(self): + """Record the commit.""" + self.committed = True + + def rollback(self): + """Record the rollback.""" + self.rolled_back = True + + def close(self): + """Mark the connection closed.""" + self.closed = True + + +class TestQuoteIdent: + """Bracket-quoting helper.""" + + def test_plain_identifier(self): + """A simple name is wrapped in brackets.""" + assert quote_ident("Orders") == "[Orders]" + + def test_identifier_with_special_chars(self): + """Dashes and spaces need no escaping, only wrapping.""" + assert quote_ident("my-table name") == "[my-table name]" + + def test_closing_bracket_is_doubled(self): + """A literal ``]`` must be doubled so it cannot close the delimiter.""" + assert quote_ident("my]table") == "[my]]table]" + + def test_injection_attempt_stays_contained(self): + """An identifier trying to break out stays inside one delimiter.""" + quoted = quote_ident("x] FROM sys.tables; DROP TABLE users --") + assert quoted.startswith("[") and quoted.endswith("]") + # The only unescaped ']' is the final delimiter. + assert quoted[1:-1].replace("]]", "") .count("]") == 0 + + +class TestParseUrl: + """URL parsing.""" + + def test_valid_url(self): + """Full URL yields all pymssql connection parameters.""" + url = "sqlserver://sa:Passw0rd@localhost:1433/testdb" + assert SQLServerLoader._parse_sqlserver_url(url) == { + "server": "localhost", + "port": 1433, + "user": "sa", + "password": "Passw0rd", + "database": "testdb", + } + + def test_default_port(self): + """Port defaults to 1433 when omitted.""" + url = "sqlserver://sa:Passw0rd@localhost/testdb" + assert SQLServerLoader._parse_sqlserver_url(url)["port"] == 1433 + + def test_percent_encoded_password(self): + """Percent-encoded credentials are decoded.""" + url = "sqlserver://sa:p%40ss%2Fword@localhost/testdb" + assert SQLServerLoader._parse_sqlserver_url(url)["password"] == "p@ss/word" + + def test_query_string_not_part_of_database(self): + """Query parameters are not swallowed into the database name.""" + url = "sqlserver://sa:pw@localhost/testdb?schema=sales" + assert SQLServerLoader._parse_sqlserver_url(url)["database"] == "testdb" + + def test_encrypt_true_requests_tls(self): + """``?encrypt=true`` asks FreeTDS to require TLS.""" + url = "sqlserver://sa:pw@localhost/testdb?encrypt=true" + assert SQLServerLoader._parse_sqlserver_url(url)["encryption"] == "require" + + def test_encryption_absent_by_default(self): + """Driver defaults are preserved when ``encrypt`` is not given.""" + url = "sqlserver://sa:pw@localhost/testdb" + assert "encryption" not in SQLServerLoader._parse_sqlserver_url(url) + + @pytest.mark.parametrize("url", [ + "mysql://sa:pw@localhost/testdb", + "sqlserver://localhost/testdb", + "sqlserver://sa:pw@localhost/", + ]) + def test_invalid_urls(self, url): + """Malformed URLs raise ValueError.""" + with pytest.raises(ValueError): + SQLServerLoader._parse_sqlserver_url(url) + + def test_schema_defaults_to_dbo(self): + """No schema parameter means ``dbo``.""" + assert SQLServerLoader.parse_schema_from_url( + "sqlserver://sa:pw@localhost/testdb") == "dbo" + + def test_schema_from_url(self): + """An explicit schema parameter is honoured.""" + assert SQLServerLoader.parse_schema_from_url( + "sqlserver://sa:pw@localhost/testdb?schema=sales") == "sales" + + +class TestSampleQuery: + """Sample-value extraction — the dict-cursor contract.""" + + def test_reads_rows_by_column_name(self): + """Rows are keyed by column name, never by position. + + Regression test: pymssql's ``as_dict=True`` cursor strips positional + keys, so ``row[0]`` raised KeyError for every non-empty column. + """ + cursor = FakeCursor([[{"status": "active"}, {"status": "closed"}]]) + values = SQLServerLoader._execute_sample_query(cursor, "dbo.Orders", "status") + assert values == ["active", "closed"] + + def test_nulls_filtered_out(self): + """NULL samples are dropped.""" + cursor = FakeCursor([[{"status": "active"}, {"status": None}]]) + assert SQLServerLoader._execute_sample_query( + cursor, "dbo.Orders", "status") == ["active"] + + def test_query_is_schema_qualified_and_quoted(self): + """Both schema and table are bracket-quoted separately.""" + cursor = FakeCursor([[]]) + SQLServerLoader._execute_sample_query(cursor, "sales.Orders", "status") + query, _ = cursor.executed[0] + assert "FROM [sales].[Orders]" in query + assert "[status]" in query + + def test_bare_table_name_still_works(self): + """An unqualified table name is quoted without a schema prefix.""" + cursor = FakeCursor([[]]) + SQLServerLoader._execute_sample_query(cursor, "Orders", "status") + query, _ = cursor.executed[0] + assert "FROM [Orders]" in query + + def test_sample_size_is_coerced_to_int(self): + """``sample_size`` cannot smuggle SQL into the TOP clause.""" + cursor = FakeCursor([[]]) + SQLServerLoader._execute_sample_query(cursor, "dbo.T", "c", sample_size=5) + query, _ = cursor.executed[0] + assert "TOP 5" in query + + def test_extract_sample_values_stringifies(self): + """The public wrapper converts values to strings.""" + cursor = FakeCursor([[{"n": 1}, {"n": 2}]]) + assert SQLServerLoader.extract_sample_values_for_column( + cursor, "dbo.T", "n") == ["1", "2"] + + +class TestIntrospection: + """Catalog introspection queries.""" + + def test_tables_query_is_schema_scoped(self): + """Table discovery binds the schema as a parameter.""" + cursor = FakeCursor([[]]) + SQLServerLoader.extract_tables_info(cursor, "sales") + query, params = cursor.executed[0] + assert params == ("sales",) + assert "s.name = %s" in query + assert "JOIN sys.schemas s" in query + + def test_columns_query_binds_schema_and_table(self): + """Column introspection is scoped by schema *and* table.""" + cursor = FakeCursor([[]]) + SQLServerLoader.extract_columns_info(cursor, "sales", "Orders") + query, params = cursor.executed[0] + assert params == ("sales", "Orders") + assert "s.name = %s AND t.name = %s" in query + + def test_columns_info_mapping(self): + """Catalog rows map onto the loader's column dict.""" + cursor = FakeCursor([ + [{ + "column_name": "id", + "data_type": "int", + "is_nullable": False, + "column_default": None, + "column_key": "PRI", + "column_comment": "", + }], + [{"id": 1}], # sample values query + ]) + info = SQLServerLoader.extract_columns_info(cursor, "dbo", "Orders") + assert info["id"]["type"] == "int" + assert info["id"]["null"] == "NO" + assert info["id"]["key"] == "PRIMARY KEY" + assert info["id"]["sample_values"] == ["1"] + assert "(NOT NULL)" in info["id"]["description"] + + def test_columns_sample_query_is_schema_qualified(self): + """Sample values are fetched from the correct schema.""" + cursor = FakeCursor([ + [{ + "column_name": "id", + "data_type": "int", + "is_nullable": True, + "column_default": None, + "column_key": "", + "column_comment": "", + }], + [], + ]) + SQLServerLoader.extract_columns_info(cursor, "sales", "Orders") + sample_query, _ = cursor.executed[1] + assert "FROM [sales].[Orders]" in sample_query + + def test_foreign_keys_mapping(self): + """Foreign key rows map onto the loader's FK dicts.""" + cursor = FakeCursor([[{ + "constraint_name": "FK_Orders_Customers", + "column_name": "customer_id", + "referenced_table_name": "Customers", + "referenced_schema_name": "dbo", + "referenced_column_name": "id", + }]]) + fks = SQLServerLoader.extract_foreign_keys(cursor, "dbo", "Orders") + assert fks == [{ + "constraint_name": "FK_Orders_Customers", + "column": "customer_id", + "referenced_table": "Customers", + "referenced_column": "id", + }] + _, params = cursor.executed[0] + assert params == ("dbo", "Orders") + + def test_relationships_grouped_by_constraint(self): + """Composite keys are grouped under one constraint name.""" + cursor = FakeCursor([[ + { + "table_name": "Orders", + "constraint_name": "FK_A", + "column_name": "c1", + "referenced_table_name": "Customers", + "referenced_column_name": "id1", + }, + { + "table_name": "Orders", + "constraint_name": "FK_A", + "column_name": "c2", + "referenced_table_name": "Customers", + "referenced_column_name": "id2", + }, + ]]) + rels = SQLServerLoader.extract_relationships(cursor, "dbo") + assert list(rels) == ["FK_A"] + assert len(rels["FK_A"]) == 2 + assert rels["FK_A"][0]["from"] == "Orders" + assert rels["FK_A"][0]["to"] == "Customers" + + def test_relationships_restricted_to_schema(self): + """Both sides of the FK are constrained to the loaded schema.""" + cursor = FakeCursor([[]]) + SQLServerLoader.extract_relationships(cursor, "sales") + query, params = cursor.executed[0] + assert params == ("sales", "sales") + assert "ps.name = %s AND rs.name = %s" in query + + def test_tables_info_builds_entities(self): + """A full table walk produces the expected entity structure.""" + cursor = FakeCursor([ + [{"table_name": "Orders", "table_comment": "All orders"}], + [{ + "column_name": "id", + "data_type": "int", + "is_nullable": False, + "column_default": None, + "column_key": "PRI", + "column_comment": "", + }], + [{"id": 7}], + [], # foreign keys + ]) + entities = SQLServerLoader.extract_tables_info(cursor, "dbo") + assert list(entities) == ["Orders"] + assert entities["Orders"]["description"] == "All orders" + assert list(entities["Orders"]["columns"]) == ["id"] + assert entities["Orders"]["foreign_keys"] == [] + + +class TestSerialization: + """Value serialization for JSON responses.""" + + @pytest.mark.parametrize("value,expected", [ + (datetime.date(2024, 1, 2), "2024-01-02"), + (datetime.datetime(2024, 1, 2, 3, 4, 5), "2024-01-02T03:04:05"), + (datetime.time(3, 4, 5), "03:04:05"), + (decimal.Decimal("1.5"), 1.5), + (b"\x01\x02", "0102"), + (None, None), + ("plain", "plain"), + ]) + def test_serialize_value(self, value, expected): + """Non-JSON-native types are converted.""" + assert SQLServerLoader._serialize_value(value) == expected + + +class TestSchemaModifyingQuery: + """DDL detection.""" + + @pytest.mark.parametrize("query,expected_op", [ + ("CREATE TABLE t (id INT)", "CREATE"), + ("ALTER TABLE t ADD c INT", "ALTER"), + ("DROP TABLE t", "DROP"), + ("TRUNCATE TABLE t", "TRUNCATE"), + ]) + def test_detects_ddl(self, query, expected_op): + """DDL statements are reported as schema-modifying.""" + modifying, op = SQLServerLoader.is_schema_modifying_query(query) + assert modifying is True + assert op == expected_op + + @pytest.mark.parametrize("query", [ + "SELECT * FROM t", + "INSERT INTO t VALUES (1)", + "", + " ", + ]) + def test_ignores_non_ddl(self, query): + """Reads and DML are not schema-modifying.""" + modifying, _ = SQLServerLoader.is_schema_modifying_query(query) + assert modifying is False + + +class TestExecuteSqlQuery: + """Query execution.""" + + def test_select_returns_serialized_rows(self): + """SELECT results are serialized for JSON transport.""" + cursor = FakeCursor([[{"id": 1, "when": datetime.date(2024, 1, 2)}]]) + conn = FakeConnection(cursor) + with patch("pymssql.connect", return_value=conn): + rows = SQLServerLoader.execute_sql_query( + "SELECT 1", "sqlserver://sa:pw@localhost/testdb") + assert rows == [{"id": 1, "when": "2024-01-02"}] + assert conn.closed and cursor.closed + + def test_non_select_reports_affected_rows(self): + """Write statements report the affected row count.""" + cursor = FakeCursor([[]]) + cursor.description = None + cursor.rowcount = 3 + conn = FakeConnection(cursor) + with patch("pymssql.connect", return_value=conn): + rows = SQLServerLoader.execute_sql_query( + "UPDATE t SET c = 1", "sqlserver://sa:pw@localhost/testdb") + assert rows == [{"operation": "UPDATE", "affected_rows": 3, "status": "success"}] + + def test_error_rolls_back_and_closes(self): + """A failing query rolls back and still releases the connection.""" + cursor = FakeCursor() + cursor.execute = MagicMock(side_effect=ValueError("boom")) + conn = FakeConnection(cursor) + with patch("pymssql.connect", return_value=conn): + with pytest.raises(SQLServerQueryError): + SQLServerLoader.execute_sql_query( + "SELECT 1", "sqlserver://sa:pw@localhost/testdb") + assert conn.rolled_back + assert conn.closed and cursor.closed + + def test_connect_failure_does_not_raise_name_error(self): + """Failing before connect() must not blow up in the error handler.""" + with patch("pymssql.connect", side_effect=ValueError("no route")): + with pytest.raises(SQLServerQueryError): + SQLServerLoader.execute_sql_query( + "SELECT 1", "sqlserver://sa:pw@localhost/testdb") + + +class TestLoad: + """End-to-end load flow.""" + + @pytest.mark.asyncio + async def test_load_success_closes_connection(self): + """A successful load reports table count and releases resources.""" + cursor = FakeCursor([ + [{"table_name": "Orders", "table_comment": ""}], + [], # columns + [], # foreign keys + [], # relationships + ]) + conn = FakeConnection(cursor) + messages = [] + with patch("pymssql.connect", return_value=conn), \ + patch("api.loaders.sqlserver_loader.load_to_graph") as mock_load: + async def _noop(*args, **kwargs): + return None + mock_load.side_effect = _noop + async for success, message in SQLServerLoader.load( + "user1", "sqlserver://sa:pw@localhost/testdb"): + messages.append((success, message)) + + assert messages[-1][0] is True + assert "Found 1 tables" in messages[-1][1] + assert conn.closed and cursor.closed + # graph name is prefix + database name + assert mock_load.call_args[0][0] == "user1_testdb" + + @pytest.mark.asyncio + async def test_load_uses_schema_from_url(self): + """The schema parameter reaches the catalog queries.""" + cursor = FakeCursor([[], []]) + conn = FakeConnection(cursor) + with patch("pymssql.connect", return_value=conn), \ + patch("api.loaders.sqlserver_loader.load_to_graph") as mock_load: + async def _noop(*args, **kwargs): + return None + mock_load.side_effect = _noop + async for _ in SQLServerLoader.load( + "user1", "sqlserver://sa:pw@localhost/testdb?schema=sales"): + pass + assert cursor.executed[0][1] == ("sales",) + + @pytest.mark.asyncio + async def test_load_failure_closes_connection(self): + """A mid-load failure still releases the connection.""" + cursor = FakeCursor() + cursor.execute = MagicMock(side_effect=ValueError("boom")) + conn = FakeConnection(cursor) + results = [] + with patch("pymssql.connect", return_value=conn): + async for success, message in SQLServerLoader.load( + "user1", "sqlserver://sa:pw@localhost/testdb"): + results.append((success, message)) + + assert results[-1][0] is False + assert conn.closed and cursor.closed + + @pytest.mark.asyncio + async def test_load_invalid_url_reports_failure(self): + """A bad URL is reported, not raised.""" + results = [] + async for success, message in SQLServerLoader.load("user1", "mysql://x/y"): + results.append((success, message)) + assert results == [(False, "Failed to load SQL Server database schema")] diff --git a/uv.lock b/uv.lock index e350bdb3..576039bf 100644 --- a/uv.lock +++ b/uv.lock @@ -1996,6 +1996,35 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e0/b0/3a8040e53df6c5c1e04b0e23ed53fdbeb64f333723a334d313fba2f581ce/pylint-4.0.7-py3-none-any.whl", hash = "sha256:be4a3111557a614411ed1fc89347ce4a8e1013a59e1f33d11485227a02e3304d", size = 539710, upload-time = "2026-08-09T19:13:21.228Z" }, ] +[[package]] +name = "pymssql" +version = "2.3.13" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/7a/cc/843c044b7f71ee329436b7327c578383e2f2499313899f88ad267cdf1f33/pymssql-2.3.13.tar.gz", hash = "sha256:2137e904b1a65546be4ccb96730a391fcd5a85aab8a0632721feb5d7e39cfbce", size = 203153, upload-time = "2026-02-14T05:00:36.865Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ba/60/a2e8a8a38f7be21d54402e2b3365cd56f1761ce9f2706c97f864e8aa8300/pymssql-2.3.13-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:cf4f32b4a05b66f02cb7d55a0f3bcb0574a6f8cf0bee4bea6f7b104038364733", size = 3158689, upload-time = "2026-02-14T04:59:46.982Z" }, + { url = "https://files.pythonhosted.org/packages/43/9e/0cf0ffb9e2f73238baf766d8e31d7237b5bee3cc1bb29a376b404610994a/pymssql-2.3.13-cp312-cp312-macosx_15_0_x86_64.whl", hash = "sha256:2b056eb175955f7fb715b60dc1c0c624969f4d24dbdcf804b41ab1e640a2b131", size = 2960018, upload-time = "2026-02-14T04:59:48.668Z" }, + { url = "https://files.pythonhosted.org/packages/93/ea/bc27354feaca717faa4626911f6b19bb62985c87dda28957c63de4de5895/pymssql-2.3.13-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:319810b89aa64b99d9c5c01518752c813938df230496fa2c4c6dda0603f04c4c", size = 3065719, upload-time = "2026-02-14T04:59:50.369Z" }, + { url = "https://files.pythonhosted.org/packages/1e/7a/8028681c96241fb5fc850b87c8959402c353e4b83c6e049a99ffa67ded54/pymssql-2.3.13-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c0ea72641cb0f8bce7ad8565dbdbda4a7437aa58bce045f2a3a788d71af2e4be", size = 3190567, upload-time = "2026-02-14T04:59:52.202Z" }, + { url = "https://files.pythonhosted.org/packages/aa/f1/ab5b76adbbd6db9ce746d448db34b044683522e7e7b95053f9dd0165297b/pymssql-2.3.13-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:1493f63d213607f708a5722aa230776ada726ccdb94097fab090a1717a2534e0", size = 3710481, upload-time = "2026-02-14T04:59:54.01Z" }, + { url = "https://files.pythonhosted.org/packages/59/aa/2fa0951475cd0a1829e0b8bfbe334d04ece4bce11546a556b005c4100689/pymssql-2.3.13-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:eb3275985c23479e952d6462ae6c8b2b6993ab6b99a92805a9c17942cf3d5b3d", size = 3453789, upload-time = "2026-02-14T04:59:56.841Z" }, + { url = "https://files.pythonhosted.org/packages/78/08/8cd2af9003f9fc03912b658a64f5a4919dcd68f0dd3bbc822b49a3d14fd9/pymssql-2.3.13-cp312-cp312-win_amd64.whl", hash = "sha256:a930adda87bdd8351a5637cf73d6491936f34e525a5e513068a6eac742f69cdb", size = 1994709, upload-time = "2026-02-14T04:59:58.972Z" }, + { url = "https://files.pythonhosted.org/packages/d4/4f/ee15b1f6b11e7c3accdc7da7840a019b63f12ba09eaa008acc601182f516/pymssql-2.3.13-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:30918bb044242865c01838909777ef5e0f1b9ecd7f5882346aefa57f4414b29c", size = 3156333, upload-time = "2026-02-14T05:00:01.21Z" }, + { url = "https://files.pythonhosted.org/packages/79/03/aea5c77bad4a52649a1d9f786a1d9ce1c83d50f1a75df288e292737b6d80/pymssql-2.3.13-cp313-cp313-macosx_15_0_x86_64.whl", hash = "sha256:1c6d0b2d7961f159a07e4f0d8cc81f70ceab83f5e7fd1e832a2d069e1d67ee4e", size = 2957990, upload-time = "2026-02-14T05:00:03.11Z" }, + { url = "https://files.pythonhosted.org/packages/5f/f8/30ac16fba32ff066b05f12c392d7b812fe11f06cb62d1d86ca5177c50a8b/pymssql-2.3.13-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:16c5957a3c9e51a03276bfd76a22431e2bc4c565e2e95f2cbb3559312edda230", size = 3065264, upload-time = "2026-02-14T05:00:05.377Z" }, + { url = "https://files.pythonhosted.org/packages/a9/98/7568447bf85921d21453fd56e19b6c9591d595fde0546c5a569f3ae937a8/pymssql-2.3.13-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0fddd24efe9d18bbf174fab7c6745b0927773718387f5517cf8082241f721a68", size = 3190039, upload-time = "2026-02-14T05:00:06.925Z" }, + { url = "https://files.pythonhosted.org/packages/35/f1/4d9d275ebaac42cdd49d40d504ccb648f27710660c8b60cc427752438c09/pymssql-2.3.13-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:123c55ee41bc7a82c76db12e2eb189b50d0d7a11222b4f8789206d1cda3b33b9", size = 3710151, upload-time = "2026-02-14T05:00:08.424Z" }, + { url = "https://files.pythonhosted.org/packages/6f/bd/a5cc6244fd27d3ea0cc82f12a7d38a24d7fd90b0022afd250014e8bfba15/pymssql-2.3.13-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:e053b443e842f9e1698fcb2b23a4bff1ff3d410894d880064e754ad823d541e5", size = 3453156, upload-time = "2026-02-14T05:00:09.978Z" }, + { url = "https://files.pythonhosted.org/packages/26/d0/c20ff0bbffd18db528bcc7b0c68b25c12ad563ed67c56ceca87c58f7399e/pymssql-2.3.13-cp313-cp313-win_amd64.whl", hash = "sha256:5c045c0f1977a679cc30d5acd9da3f8aeb2dc6e744895b26444b4a2f20dad9a0", size = 1995236, upload-time = "2026-02-14T05:00:11.495Z" }, + { url = "https://files.pythonhosted.org/packages/ec/5f/6b64f78181d680f655ab40ba7b34cb68c045a2f4e04a10a70d768cd383b7/pymssql-2.3.13-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:fc5482969c813b0a45ce51c41844ae5bfa8044ad5ef8b4820ef6de7d4545b7f2", size = 3158377, upload-time = "2026-02-14T05:00:13.581Z" }, + { url = "https://files.pythonhosted.org/packages/ff/24/155dbb0992c431496d440f47fb9d587cd0059ee20baf65e3d891794d862a/pymssql-2.3.13-cp314-cp314-macosx_15_0_x86_64.whl", hash = "sha256:ff5be7ab1d643dbce2ee3424d2ef9ae8e4146cf75bd20946bc7a6108e3ad1e47", size = 2959039, upload-time = "2026-02-14T05:00:15.883Z" }, + { url = "https://files.pythonhosted.org/packages/c9/89/b453dd1b1188779621fb974ac715ab2e738f4a0b69f7291ab014298bd80d/pymssql-2.3.13-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8d66ce0a249d2e3b57369048d71e1f00d08dfb90a758d134da0250ae7bc739c1", size = 3063862, upload-time = "2026-02-14T05:00:17.537Z" }, + { url = "https://files.pythonhosted.org/packages/02/e5/96f57c78162013678ecc3f3f7e5fb52c83ee07beef26906d0870770c3ef6/pymssql-2.3.13-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d663c908414a6a032f04d17628138b1782af916afc0df9fefac4751fa394c3ac", size = 3188155, upload-time = "2026-02-14T05:00:19.011Z" }, + { url = "https://files.pythonhosted.org/packages/cd/a2/4bee9484734ae0c55d10a2f6ff82dd4e416f52420755161b8760c817ad64/pymssql-2.3.13-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:aa5e07eff7e6e8bd4ba22c30e4cb8dd073e138cd272090603609a15cc5dbc75b", size = 3709344, upload-time = "2026-02-14T05:00:21.139Z" }, + { url = "https://files.pythonhosted.org/packages/37/cf/3520d96afa213c88db4f4a1988199db476d869a62afdd5d9c4635c184631/pymssql-2.3.13-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:db77da1a3fc9b5b5c5400639d79d7658ba7ad620957100c5b025be608b562193", size = 3451799, upload-time = "2026-02-14T05:00:22.504Z" }, + { url = "https://files.pythonhosted.org/packages/25/50/4be9bd9cf4b43208a7175117a533ece200cfe4131a39f9909bdc7560ddeb/pymssql-2.3.13-cp314-cp314-win_amd64.whl", hash = "sha256:7d7037d2b5b907acc7906d0479924db2935a70c720450c41339146a4ada2b93d", size = 2049139, upload-time = "2026-02-14T05:00:23.951Z" }, +] + [[package]] name = "pymysql" version = "1.2.0" @@ -2232,6 +2261,7 @@ all = [ { name = "jinja2" }, { name = "playwright" }, { name = "pylint" }, + { name = "pymssql" }, { name = "pytest" }, { name = "pytest-asyncio" }, { name = "pytest-playwright" }, @@ -2258,6 +2288,7 @@ server = [ { name = "graphiti-core" }, { name = "itsdangerous" }, { name = "jinja2" }, + { name = "pymssql" }, { name = "python-multipart" }, { name = "snowflake-connector-python" }, { name = "uvicorn" }, @@ -2299,6 +2330,8 @@ requires-dist = [ { name = "pydantic", specifier = ">=2.0" }, { name = "pylint", marker = "extra == 'all'", specifier = "~=4.0.3" }, { name = "pylint", marker = "extra == 'dev'", specifier = "~=4.0.3" }, + { name = "pymssql", marker = "extra == 'all'", specifier = "~=2.3.13" }, + { name = "pymssql", marker = "extra == 'server'", specifier = "~=2.3.13" }, { name = "pymysql", specifier = "~=1.2.0" }, { name = "pytest", marker = "extra == 'all'", specifier = ">=9.0.3,<9.2.0" }, { name = "pytest", marker = "extra == 'dev'", specifier = ">=9.0.3,<9.2.0" }, From 32ce1612dd1d3913655e6935f518b76d9ff044e3 Mon Sep 17 00:00:00 2001 From: Guy Korland Date: Thu, 13 Aug 2026 21:16:13 +0300 Subject: [PATCH 02/10] fix(loaders): validate SQL Server identifiers before interpolation CodeQL flagged the sample-value query as `py/sql-injection` (high): the schema name reaches it from the user-supplied connection URL, and T-SQL cannot bind identifiers as parameters. Adds `validate_ident`, an anchored allow-list matching the existing `SnowflakeLoader._validate_identifier` pattern. It accepts only characters that can legitimately appear in a SQL Server object name and rejects everything capable of escaping a bracket delimiter (`]`, quotes, semicolons, backslashes, control characters), plus empty and over-long names. `quote_ident` keeps doubling `]` as defence in depth. Validation runs before the statement is built, so a hostile identifier never reaches `cursor.execute`. `parse_schema_from_url` now validates the schema at parse time, and `sample_size` is checked to be a positive int. Also imports `api.core` ahead of the loader in the new test module. The package's `__init__` eagerly pulls in the pipeline, which imports the loaders, so importing a loader first left `graph_loader` half-built and the file could not be run on its own. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- api/loaders/sqlserver_loader.py | 57 ++++++++++++++++++++++++--- tests/test_sqlserver_loader.py | 70 +++++++++++++++++++++++++++++++++ 2 files changed, 121 insertions(+), 6 deletions(-) diff --git a/api/loaders/sqlserver_loader.py b/api/loaders/sqlserver_loader.py index 239f4943..d8ae1722 100644 --- a/api/loaders/sqlserver_loader.py +++ b/api/loaders/sqlserver_loader.py @@ -27,6 +27,39 @@ class SQLServerConnectionError(Exception): """Exception raised for SQL Server connection errors.""" +def validate_ident(identifier: str, identifier_type: str = "identifier") -> str: + """Validate that an identifier is safe to interpolate into T-SQL. + + T-SQL cannot bind identifiers as parameters, so table, schema and column + names must be interpolated. This is an anchored allow-list: only characters + that can legitimately appear in a SQL Server object name are accepted, and + everything capable of breaking out of a bracket-delimited identifier + (``]``, quotes, semicolons, backslashes, control characters) is rejected. + + Args: + identifier: Raw identifier, typically read from the system catalog. + identifier_type: Label used in the error message. + + Returns: + The identifier, unchanged, once validated. + + Raises: + ValueError: If the identifier is empty, over-long, or contains a + character outside the allow-list. + """ + if not identifier or len(identifier) > 128: + raise ValueError( + f"Invalid {identifier_type}: {identifier!r}. " + "Must be between 1 and 128 characters." + ) + if not re.fullmatch(r'[A-Za-z0-9_$#@ .\-]+', identifier): + raise ValueError( + f"Invalid {identifier_type}: {identifier!r}. Only letters, digits, " + "underscore, dollar, hash, at-sign, space, dot and dash are allowed." + ) + return identifier + + def quote_ident(identifier: str) -> str: """Bracket-quote a T-SQL identifier, escaping any embedded ``]``. @@ -34,6 +67,9 @@ def quote_ident(identifier: str) -> str: doubling it, so ``my]table`` must become ``[my]]table]``. Without this a crafted identifier would terminate the quote early. + This is defence in depth: callers that interpolate the result into a + statement validate the identifier with :func:`validate_ident` first. + Args: identifier: Raw identifier as read from the system catalog. @@ -81,13 +117,16 @@ def _execute_sample_query( bracket-quoted separately so the schema prefix survives. """ schema, _, bare_table = table_name.rpartition('.') - qualified = quote_ident(bare_table) + qualified = quote_ident(validate_ident(bare_table, "table name")) if schema: - qualified = f"{quote_ident(schema)}.{qualified}" + qualified = f"{quote_ident(validate_ident(schema, 'schema name'))}.{qualified}" - col = quote_ident(col_name) - # ``sample_size`` is coerced to int; identifiers are bracket-quoted with - # ``]`` escaped, since T-SQL cannot bind identifiers as parameters. + col = quote_ident(validate_ident(col_name, "column name")) + if not isinstance(sample_size, int) or sample_size <= 0: + raise ValueError(f"sample_size must be a positive integer, got {sample_size!r}") + + # Identifiers are allow-list validated and bracket-quoted with ``]`` + # escaped, since T-SQL cannot bind identifiers as parameters. query = ( f"SELECT DISTINCT TOP {int(sample_size)} {col}" f" FROM {qualified}" @@ -137,13 +176,19 @@ def parse_schema_from_url(connection_url: str) -> str: Returns: The requested schema, or ``dbo`` when not specified. + + Raises: + ValueError: If the requested schema is not a valid identifier. """ try: parsed = urlparse(connection_url) schema = parse_qs(parsed.query).get('schema', [''])[0] - return unquote(schema).strip() or DEFAULT_SCHEMA + schema = unquote(schema).strip() except (ValueError, AttributeError): return DEFAULT_SCHEMA + if not schema: + return DEFAULT_SCHEMA + return validate_ident(schema, "schema name") @staticmethod def _parse_sqlserver_url(connection_url: str) -> Dict[str, Any]: diff --git a/tests/test_sqlserver_loader.py b/tests/test_sqlserver_loader.py index b7ee3bb2..3b3a1a6d 100644 --- a/tests/test_sqlserver_loader.py +++ b/tests/test_sqlserver_loader.py @@ -12,10 +12,16 @@ import pytest +# ``api.core`` must be initialised before any loader module is imported. +# ``api.core.__init__`` eagerly pulls in the pipeline, which imports the +# loaders, so importing a loader first leaves ``graph_loader`` half-built. +import api.core # noqa: F401 pylint: disable=unused-import + from api.loaders.sqlserver_loader import ( SQLServerLoader, SQLServerQueryError, quote_ident, + validate_ident, ) @@ -99,6 +105,70 @@ def test_injection_attempt_stays_contained(self): assert quoted[1:-1].replace("]]", "") .count("]") == 0 +class TestValidateIdent: + """Allow-list validation applied before any identifier interpolation.""" + + @pytest.mark.parametrize("name", [ + "Orders", "my-table name", "col_1", "tbl$", "#temp", "a.b", "x@y", + ]) + def test_accepts_legitimate_names(self, name): + """Characters that can legally appear in an object name pass through.""" + assert validate_ident(name) == name + + @pytest.mark.parametrize("name", [ + "x] FROM sys.tables; DROP TABLE users --", + "my]table", + "tbl'; DROP TABLE t --", + 'tbl"', + "tbl;", + "tbl\\x", + "tbl\nDROP", + ]) + def test_rejects_breakout_attempts(self, name): + """Anything able to escape a bracket delimiter is refused.""" + with pytest.raises(ValueError): + validate_ident(name) + + def test_rejects_empty(self): + """An empty identifier is not a valid object name.""" + with pytest.raises(ValueError): + validate_ident("") + + def test_rejects_over_long(self): + """SQL Server object names cap at 128 characters.""" + with pytest.raises(ValueError): + validate_ident("a" * 129) + + def test_error_names_the_identifier_type(self): + """The message says which kind of identifier was rejected.""" + with pytest.raises(ValueError, match="schema name"): + validate_ident("bad;name", "schema name") + + +class TestSampleQueryValidation: + """The sample query refuses hostile identifiers outright.""" + + @pytest.mark.parametrize("table,column", [ + ("dbo.x] FROM sys.tables --", "c"), + ("dbo.T", "c] FROM sys.tables --"), + ("bad;schema.T", "c"), + ]) + def test_hostile_identifier_is_rejected(self, table, column): + """Validation happens before the statement is built or executed.""" + cursor = FakeCursor([[]]) + with pytest.raises(ValueError): + SQLServerLoader._execute_sample_query(cursor, table, column) + assert cursor.executed == [] + + @pytest.mark.parametrize("size", [0, -1, "5"]) + def test_invalid_sample_size_rejected(self, size): + """``sample_size`` must be a positive integer.""" + cursor = FakeCursor([[]]) + with pytest.raises(ValueError): + SQLServerLoader._execute_sample_query(cursor, "dbo.T", "c", sample_size=size) + assert cursor.executed == [] + + class TestParseUrl: """URL parsing.""" From 889d97ca8fef90a98c2fdf710e0544fcaf7e55bf Mon Sep 17 00:00:00 2001 From: Guy Korland Date: Thu, 13 Aug 2026 21:21:56 +0300 Subject: [PATCH 03/10] fix(loaders): qualify SQL Server sample queries with the catalog schema CodeQL still reported `py/sql-injection` after the allow-list validator: the schema name reaching the sample query originated in the user-supplied connection URL, and an anchored regex is not recognised as a barrier. The tables query now selects `s.name AS schema_name` back from `sys.schemas`, and that server-returned value is what gets interpolated into the sample query. The URL string is still used, but only as a bound query parameter, so it never reaches a statement body. This is also more correct: sampling now uses the server's canonical casing for the schema rather than whatever the URL happened to contain. Extracts `_build_column_description` and a `_KEY_TYPES` lookup out of `extract_columns_info` to keep it within the local-variable limit. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- api/loaders/sqlserver_loader.py | 101 ++++++++++++++++++-------------- tests/test_sqlserver_loader.py | 27 ++++++++- 2 files changed, 83 insertions(+), 45 deletions(-) diff --git a/api/loaders/sqlserver_loader.py b/api/loaders/sqlserver_loader.py index d8ae1722..ac4fc0e3 100644 --- a/api/loaders/sqlserver_loader.py +++ b/api/loaders/sqlserver_loader.py @@ -79,6 +79,38 @@ def quote_ident(identifier: str) -> str: return f"[{identifier.replace(']', ']]')}]" +_KEY_TYPES = { + 'PRI': 'PRIMARY KEY', + 'MUL': 'FOREIGN KEY', + 'UNI': 'UNIQUE KEY', +} + + +def _build_column_description(col_info: Dict[str, Any], key_type: str, is_nullable: str) -> str: + """Build the human-readable description shown for a column. + + Args: + col_info: One row from the column catalog query. + key_type: Resolved key kind, or ``NONE``. + is_nullable: ``YES`` or ``NO``. + + Returns: + The description string. + """ + comment = col_info['column_comment'] + parts = [ + str(comment) if comment + else f"Column {col_info['column_name']} of type {col_info['data_type']}" + ] + if key_type != 'NONE': + parts.append(f"({key_type})") + if is_nullable == 'NO': + parts.append("(NOT NULL)") + if col_info['column_default'] is not None: + parts.append(f"(Default: {col_info['column_default']})") + return ' '.join(parts) + + class SQLServerLoader(BaseLoader): """ Loader for SQL Server databases that connects and extracts schema information. @@ -329,10 +361,13 @@ def extract_tables_info(cursor, schema: str = DEFAULT_SCHEMA) -> Dict[str, Any]: """ entities = {} - # Get all tables in the requested schema + # Get all tables in the requested schema. ``s.name`` is selected back so + # sample queries qualify tables with the server's own canonical schema + # name rather than the string taken from the connection URL. cursor.execute(""" SELECT t.name AS table_name, + s.name AS schema_name, ISNULL(CAST(ep.value AS NVARCHAR(MAX)), '') AS table_comment FROM sys.tables t JOIN sys.schemas s ON t.schema_id = s.schema_id @@ -351,9 +386,12 @@ def extract_tables_info(cursor, schema: str = DEFAULT_SCHEMA) -> Dict[str, Any]: for table_info in tqdm.tqdm(tables, desc="Extracting table information"): table_name = table_info['table_name'] table_comment = table_info['table_comment'] + catalog_schema = table_info['schema_name'] # Get column information for this table - columns_info = SQLServerLoader.extract_columns_info(cursor, schema, table_name) + columns_info = SQLServerLoader.extract_columns_info( + cursor, schema, table_name, catalog_schema + ) # Get foreign keys for this table foreign_keys = SQLServerLoader.extract_foreign_keys(cursor, schema, table_name) @@ -374,14 +412,20 @@ def extract_tables_info(cursor, schema: str = DEFAULT_SCHEMA) -> Dict[str, Any]: return entities @staticmethod - def extract_columns_info(cursor, schema: str, table_name: str) -> Dict[str, Any]: + def extract_columns_info( + cursor, schema: str, table_name: str, catalog_schema: str = None + ) -> Dict[str, Any]: """ Extract column information for a specific table. Args: cursor: Database cursor - schema: Schema owning the table + schema: Schema owning the table, used as a bound query parameter table_name: Name of the table + catalog_schema: Schema name as returned by ``sys.schemas``. Sample + queries interpolate this rather than *schema*, so the value + comes from the server rather than the connection URL. Falls + back to *schema* when not supplied. Returns: Dict containing column information @@ -430,52 +474,23 @@ def extract_columns_info(cursor, schema: str, table_name: str) -> Dict[str, Any] columns = cursor.fetchall() columns_info = {} + qualified_table = f"{catalog_schema or schema}.{table_name}" + for col_info in columns: col_name = col_info['column_name'] - data_type = col_info['data_type'] is_nullable = 'YES' if col_info['is_nullable'] else 'NO' - column_default = col_info['column_default'] - column_key = col_info['column_key'] - column_comment = col_info['column_comment'] - - # Determine key type - if column_key == 'PRI': - key_type = 'PRIMARY KEY' - elif column_key == 'MUL': - key_type = 'FOREIGN KEY' - elif column_key == 'UNI': - key_type = 'UNIQUE KEY' - else: - key_type = 'NONE' - - # Generate column description - description_parts = [] - if column_comment: - description_parts.append(str(column_comment)) - else: - description_parts.append(f"Column {col_name} of type {data_type}") - - if key_type != 'NONE': - description_parts.append(f"({key_type})") - - if is_nullable == 'NO': - description_parts.append("(NOT NULL)") - - if column_default is not None: - description_parts.append(f"(Default: {column_default})") - - # Extract sample values for the column (stored separately, not in description) - sample_values = SQLServerLoader.extract_sample_values_for_column( - cursor, f"{schema}.{table_name}", col_name - ) + key_type = _KEY_TYPES.get(col_info['column_key'], 'NONE') columns_info[col_name] = { - 'type': data_type, + 'type': col_info['data_type'], 'null': is_nullable, 'key': key_type, - 'description': ' '.join(description_parts), - 'default': column_default, - 'sample_values': sample_values + 'description': _build_column_description(col_info, key_type, is_nullable), + 'default': col_info['column_default'], + # Stored separately, not folded into the description. + 'sample_values': SQLServerLoader.extract_sample_values_for_column( + cursor, qualified_table, col_name + ), } return columns_info diff --git a/tests/test_sqlserver_loader.py b/tests/test_sqlserver_loader.py index 3b3a1a6d..c8445494 100644 --- a/tests/test_sqlserver_loader.py +++ b/tests/test_sqlserver_loader.py @@ -388,7 +388,7 @@ def test_relationships_restricted_to_schema(self): def test_tables_info_builds_entities(self): """A full table walk produces the expected entity structure.""" cursor = FakeCursor([ - [{"table_name": "Orders", "table_comment": "All orders"}], + [{"table_name": "Orders", "schema_name": "dbo", "table_comment": "All orders"}], [{ "column_name": "id", "data_type": "int", @@ -406,6 +406,29 @@ def test_tables_info_builds_entities(self): assert list(entities["Orders"]["columns"]) == ["id"] assert entities["Orders"]["foreign_keys"] == [] + def test_sample_query_uses_catalog_schema_not_url_schema(self): + """Sample queries qualify with the schema echoed back by ``sys.schemas``. + + The catalog value comes from the server, so the connection URL string + is never interpolated into a statement. + """ + cursor = FakeCursor([ + [{"table_name": "Orders", "schema_name": "Sales", "table_comment": ""}], + [{ + "column_name": "id", + "data_type": "int", + "is_nullable": False, + "column_default": None, + "column_key": "PRI", + "column_comment": "", + }], + [{"id": 7}], + [], # foreign keys + ]) + SQLServerLoader.extract_tables_info(cursor, "sales") + sample_query = cursor.executed[2][0] + assert "FROM [Sales].[Orders]" in sample_query + class TestSerialization: """Value serialization for JSON responses.""" @@ -502,7 +525,7 @@ class TestLoad: async def test_load_success_closes_connection(self): """A successful load reports table count and releases resources.""" cursor = FakeCursor([ - [{"table_name": "Orders", "table_comment": ""}], + [{"table_name": "Orders", "schema_name": "dbo", "table_comment": ""}], [], # columns [], # foreign keys [], # relationships From ac1143e4b3bb1c6abe397d0dee844344a6165365 Mon Sep 17 00:00:00 2001 From: Guy Korland Date: Thu, 13 Aug 2026 21:27:11 +0300 Subject: [PATCH 04/10] fix(loaders): drop URL-schema fallback in SQL Server sample queries `catalog_schema` was optional and fell back to the URL-derived schema, which kept the tainted value flowing into the interpolated sample query and left CodeQL's `py/sql-injection` alert open. It is now a required argument, so the only schema string that can reach a statement body is the one `sys.schemas` returned. The URL schema is used exclusively as a bound query parameter. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- api/loaders/sqlserver_loader.py | 13 +++++++------ tests/test_sqlserver_loader.py | 11 +++++++---- 2 files changed, 14 insertions(+), 10 deletions(-) diff --git a/api/loaders/sqlserver_loader.py b/api/loaders/sqlserver_loader.py index ac4fc0e3..c9c3b595 100644 --- a/api/loaders/sqlserver_loader.py +++ b/api/loaders/sqlserver_loader.py @@ -413,19 +413,20 @@ def extract_tables_info(cursor, schema: str = DEFAULT_SCHEMA) -> Dict[str, Any]: @staticmethod def extract_columns_info( - cursor, schema: str, table_name: str, catalog_schema: str = None + cursor, schema: str, table_name: str, catalog_schema: str ) -> Dict[str, Any]: """ Extract column information for a specific table. Args: cursor: Database cursor - schema: Schema owning the table, used as a bound query parameter - table_name: Name of the table + schema: Schema owning the table. Only ever passed to the driver as + a bound query parameter, never interpolated into a statement. + table_name: Name of the table, as returned by ``sys.tables`` catalog_schema: Schema name as returned by ``sys.schemas``. Sample queries interpolate this rather than *schema*, so the value - comes from the server rather than the connection URL. Falls - back to *schema* when not supplied. + that reaches a statement body comes from the server rather + than from the connection URL. Returns: Dict containing column information @@ -474,7 +475,7 @@ def extract_columns_info( columns = cursor.fetchall() columns_info = {} - qualified_table = f"{catalog_schema or schema}.{table_name}" + qualified_table = f"{catalog_schema}.{table_name}" for col_info in columns: col_name = col_info['column_name'] diff --git a/tests/test_sqlserver_loader.py b/tests/test_sqlserver_loader.py index c8445494..cf57475b 100644 --- a/tests/test_sqlserver_loader.py +++ b/tests/test_sqlserver_loader.py @@ -292,7 +292,7 @@ def test_tables_query_is_schema_scoped(self): def test_columns_query_binds_schema_and_table(self): """Column introspection is scoped by schema *and* table.""" cursor = FakeCursor([[]]) - SQLServerLoader.extract_columns_info(cursor, "sales", "Orders") + SQLServerLoader.extract_columns_info(cursor, "sales", "Orders", "Sales") query, params = cursor.executed[0] assert params == ("sales", "Orders") assert "s.name = %s AND t.name = %s" in query @@ -310,7 +310,7 @@ def test_columns_info_mapping(self): }], [{"id": 1}], # sample values query ]) - info = SQLServerLoader.extract_columns_info(cursor, "dbo", "Orders") + info = SQLServerLoader.extract_columns_info(cursor, "dbo", "Orders", "dbo") assert info["id"]["type"] == "int" assert info["id"]["null"] == "NO" assert info["id"]["key"] == "PRIMARY KEY" @@ -330,9 +330,12 @@ def test_columns_sample_query_is_schema_qualified(self): }], [], ]) - SQLServerLoader.extract_columns_info(cursor, "sales", "Orders") + SQLServerLoader.extract_columns_info(cursor, "sales", "Orders", "Sales") + # The column query binds the URL schema; the sample query interpolates + # the catalog-returned one. + assert cursor.executed[0][1] == ("sales", "Orders") sample_query, _ = cursor.executed[1] - assert "FROM [sales].[Orders]" in sample_query + assert "FROM [Sales].[Orders]" in sample_query def test_foreign_keys_mapping(self): """Foreign key rows map onto the loader's FK dicts.""" From 563fd2cd982c12d4e054ba49d13f259c3444a6dd Mon Sep 17 00:00:00 2001 From: Anchel135 Date: Mon, 24 Aug 2026 14:33:57 +0300 Subject: [PATCH 05/10] fix(loaders): offload SQL Server introspection off the event loop MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit SQLServerLoader.load ran pymssql.connect and every cursor execute/fetch inline in an async generator, so a schema load stalled the whole event loop — including other requests and the stream keepalives. Move the driver work into _introspect_schema and await it through run_introspection, matching the PostgreSQL and MySQL loaders. The connection and cursor are now created, used and closed by the same worker thread, so a cancelled load cannot leave two threads on one connection. The URL is still parsed on the loop (pure string work) so a malformed URL fails before any progress message is emitted. Adds test_sqlserver_load_does_not_block_the_loop alongside the existing Postgres/MySQL loop-responsiveness tests. Also: docs said the loader extracts views, but it only queries sys.tables; README now states SQL Server and Snowflake need the queryweaver[server] extra; the port field falls back to the vendor default it already shows as a placeholder; test modules gained the unit marker; and the api.core side-effect import is now an explicit importlib.import_module call. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- README.md | 9 +++- api/loaders/sqlserver_loader.py | 58 +++++++++++++-------- app/src/components/modals/DatabaseModal.tsx | 9 +++- docs/sqlserver_loader.md | 4 +- tests/test_schema_load_offloading.py | 29 +++++++++++ tests/test_sql_sanitizer.py | 4 ++ tests/test_sqlserver_loader.py | 9 +++- 7 files changed, 91 insertions(+), 31 deletions(-) diff --git a/README.md b/README.md index c45bb846..17157c8a 100644 --- a/README.md +++ b/README.md @@ -256,6 +256,10 @@ pip install queryweaver[server] pip install queryweaver[dev] ``` +> **SQL Server and Snowflake need `queryweaver[server]`.** Their drivers +> (`pymssql` and `snowflake-connector-python`) ship in the `server` extra, so the +> minimal SDK install can only connect to PostgreSQL and MySQL. + ### Quick Start ```python @@ -310,7 +314,7 @@ async with QueryWeaver(falkordb_url="redis://host-a:6379", user_id="tenant_a") a | Method | Description | |--------|-------------| -| `connect_database(db_url)` | Connect PostgreSQL/MySQL/SQL Server/Snowflake and load schema | +| `connect_database(db_url)` | Connect PostgreSQL/MySQL/SQL Server/Snowflake and load schema (SQL Server and Snowflake require `queryweaver[server]`) | | `query(database, question)` | Convert natural language to SQL and execute | | `get_schema(database)` | Retrieve database schema (tables and relationships) | | `list_databases()` | List all connected databases | @@ -356,7 +360,8 @@ if result.requires_confirmation: - Python 3.12+ - FalkorDB instance (local or remote) - OpenAI or Azure OpenAI API key (for LLM) -- Target SQL database (PostgreSQL, MySQL, SQL Server or Snowflake) +- Target SQL database (PostgreSQL, MySQL, SQL Server or Snowflake — the last two + require the `queryweaver[server]` extra) ## Development diff --git a/api/loaders/sqlserver_loader.py b/api/loaders/sqlserver_loader.py index c9c3b595..7dded709 100644 --- a/api/loaders/sqlserver_loader.py +++ b/api/loaders/sqlserver_loader.py @@ -12,6 +12,7 @@ from api.loaders.base_loader import BaseLoader from api.loaders.graph_loader import load_to_graph +from api.loaders.introspection import run_introspection logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") @@ -273,6 +274,31 @@ def _parse_sqlserver_url(connection_url: str) -> Dict[str, Any]: return params + @staticmethod + def _introspect_schema(conn_params: Dict[str, Any], schema: str): + """Connect, introspect and close — all inside one worker thread. + + Everything touching the driver lives here so the connection and cursor + are created, used and closed by the same thread. Closing them from the + event loop instead can run while an offloaded introspection is still + using them, because cancelling a ``to_thread`` call does not stop the + thread it is running in. + + Mirrors ``PostgresLoader._introspect_schema``; see ``load`` for why the + work is offloaded at all. + """ + conn = None + cursor = None + try: + conn = pymssql.connect(**conn_params) # pylint: disable=no-member + cursor = conn.cursor(as_dict=True) + + entities = SQLServerLoader.extract_tables_info(cursor, schema) + relationships = SQLServerLoader.extract_relationships(cursor, schema) + return entities, relationships + finally: + SQLServerLoader._close_quietly(cursor, conn) + @staticmethod async def load( # pylint: disable=arguments-differ prefix: str, @@ -291,33 +317,21 @@ async def load( # pylint: disable=arguments-differ Yields: Tuple[bool, str]: Success status and message """ - conn = None - cursor = None try: - # Parse connection URL + # Parsed here rather than in the worker so a malformed URL is + # reported as a failure before any progress is announced. conn_params = SQLServerLoader._parse_sqlserver_url(connection_url) schema = SQLServerLoader.parse_schema_from_url(connection_url) - - # Connect to SQL Server database - conn = pymssql.connect(**conn_params) # pylint: disable=no-member - cursor = conn.cursor(as_dict=True) - - # Get database name db_name = conn_params['database'] - # Get all table information + # pymssql is a blocking driver, so every connect/execute/fetch would + # stall the event loop — and with it every other request and the + # stream keepalives. Offload to the shared introspection executor, + # the same way the PostgreSQL and MySQL loaders do. yield True, "Extracting table information..." - entities = SQLServerLoader.extract_tables_info(cursor, schema) - - # Get all relationship information - yield True, "Extracting relationship information..." - relationships = SQLServerLoader.extract_relationships(cursor, schema) - - # Close database connection - cursor.close() - cursor = None - conn.close() - conn = None + entities, relationships = await run_introspection( + SQLServerLoader._introspect_schema, conn_params, schema + ) # Load data into graph yield True, "Loading data into graph..." @@ -333,8 +347,6 @@ async def load( # pylint: disable=arguments-differ except Exception as e: # pylint: disable=broad-exception-caught logging.error("Error loading SQL Server schema: %s", e) yield False, "Failed to load SQL Server database schema" - finally: - SQLServerLoader._close_quietly(cursor, conn) @staticmethod def _close_quietly(cursor, conn) -> None: diff --git a/app/src/components/modals/DatabaseModal.tsx b/app/src/components/modals/DatabaseModal.tsx index a7158e29..ce24ab74 100644 --- a/app/src/components/modals/DatabaseModal.tsx +++ b/app/src/components/modals/DatabaseModal.tsx @@ -101,6 +101,11 @@ const DatabaseModal = ({ open, onOpenChange }: DatabaseModalProps) => { }; const handleConnect = async () => { + // The port field shows the vendor default as a placeholder; treat that hint + // as the actual default so leaving the field empty works instead of + // failing the "fill in all required fields" check. + const effectivePort = port || (selectedDatabase ? getDbProfile(selectedDatabase).port : ''); + // Validate based on connection mode if (connectionMode === 'url') { if (!connectionUrl || !selectedDatabase) { @@ -130,7 +135,7 @@ const DatabaseModal = ({ open, onOpenChange }: DatabaseModalProps) => { return; } } else { - if (!selectedDatabase || !host || !port || !database || !username) { + if (!selectedDatabase || !host || !effectivePort || !database || !username) { toast({ title: "Missing Information", description: "Please fill in all required fields", @@ -165,7 +170,7 @@ const DatabaseModal = ({ open, onOpenChange }: DatabaseModalProps) => { dbUrl = builtUrl.toString(); } else { const profile = getDbProfile(selectedDatabase); - const builtUrl = new URL(`${profile.protocol}://${host}:${port}/${database}`); + const builtUrl = new URL(`${profile.protocol}://${host}:${effectivePort}/${database}`); builtUrl.username = username; builtUrl.password = password; diff --git a/docs/sqlserver_loader.md b/docs/sqlserver_loader.md index ea6100a7..9fb52dc2 100644 --- a/docs/sqlserver_loader.md +++ b/docs/sqlserver_loader.md @@ -46,7 +46,7 @@ sqlserver://appuser:s3cr3t@sql.example.com:1433/reporting?schema=dbo&encrypt=tru ### Schema Extraction -- Tables and views in the selected schema +- Tables in the selected schema (views are not extracted) - Columns with data types, nullability, defaults and primary-key flags - Extended properties (`MS_Description`) used as table and column descriptions - Foreign keys, including composite keys @@ -109,7 +109,7 @@ print(response.json()) The loader reads from SQL Server system catalog views: -- `sys.tables` / `sys.views` joined with `sys.schemas` — table list +- `sys.tables` joined with `sys.schemas` — table list - `sys.columns` joined with `sys.types` — column metadata - `sys.indexes` / `sys.index_columns` — primary keys - `sys.foreign_keys` / `sys.foreign_key_columns` — foreign keys diff --git a/tests/test_schema_load_offloading.py b/tests/test_schema_load_offloading.py index 0e056cb3..16ca3a3d 100644 --- a/tests/test_schema_load_offloading.py +++ b/tests/test_schema_load_offloading.py @@ -15,6 +15,7 @@ from api.config import Config from api.core.pipeline import MySQLLoader, PostgresLoader +from api.loaders.sqlserver_loader import SQLServerLoader STALL = 0.3 TICK = 0.02 @@ -99,6 +100,34 @@ async def noop(*_args, **_kwargs): assert max(gaps) < STALL, f"loop blocked for {max(gaps):.2f}s" +@pytest.mark.unit +@patch("api.loaders.sqlserver_loader.load_to_graph") +@patch("api.loaders.sqlserver_loader.SQLServerLoader.extract_relationships", _slow) +@patch("api.loaders.sqlserver_loader.SQLServerLoader.extract_tables_info", _slow) +@patch("api.loaders.sqlserver_loader.pymssql.connect") +async def test_sqlserver_load_does_not_block_the_loop(mock_connect, mock_load_to_graph): + def slow_connect(*_args, **_kwargs): + time.sleep(STALL) + return MagicMock() + + mock_connect.side_effect = slow_connect + + async def noop(*_args, **_kwargs): + return None + + mock_load_to_graph.side_effect = noop + + _steps, ticks = await _ticks_while_consuming( + SQLServerLoader.load("pfx", "sqlserver://u:p@h:1433/db") + ) + + assert len(ticks) > (STALL * 3 / TICK) * 0.3, ( + f"event loop starved during schema load: {len(ticks)} ticks" + ) + gaps = [b - a for a, b in zip(ticks, ticks[1:])] + assert max(gaps) < STALL, f"loop blocked for {max(gaps):.2f}s" + + @pytest.mark.unit @pytest.mark.parametrize("loader_module,loader,url", [ ("api.loaders.postgres_loader", "PostgresLoader", "postgresql://u:p@h:5432/db"), diff --git a/tests/test_sql_sanitizer.py b/tests/test_sql_sanitizer.py index 9589666d..ac47c5dc 100644 --- a/tests/test_sql_sanitizer.py +++ b/tests/test_sql_sanitizer.py @@ -1,7 +1,11 @@ """Unit tests for SQL identifier quoting utilities.""" +import pytest + from api.sql_utils import SQLIdentifierQuoter, DatabaseSpecificQuoter +pytestmark = pytest.mark.unit + class TestSQLIdentifierQuoter: """Test cases for SQLIdentifierQuoter.""" diff --git a/tests/test_sqlserver_loader.py b/tests/test_sqlserver_loader.py index cf57475b..af24e837 100644 --- a/tests/test_sqlserver_loader.py +++ b/tests/test_sqlserver_loader.py @@ -8,6 +8,7 @@ import datetime import decimal +import importlib from unittest.mock import patch, MagicMock import pytest @@ -15,15 +16,19 @@ # ``api.core`` must be initialised before any loader module is imported. # ``api.core.__init__`` eagerly pulls in the pipeline, which imports the # loaders, so importing a loader first leaves ``graph_loader`` half-built. -import api.core # noqa: F401 pylint: disable=unused-import +# Done through ``importlib`` so the side effect reads as deliberate rather than +# as an unused import that a linter should strip. +importlib.import_module("api.core") -from api.loaders.sqlserver_loader import ( +from api.loaders.sqlserver_loader import ( # noqa: E402 pylint: disable=wrong-import-position SQLServerLoader, SQLServerQueryError, quote_ident, validate_ident, ) +pytestmark = pytest.mark.unit + class FakeCursor: """Minimal stand-in for a pymssql ``as_dict=True`` cursor. From 567c79547a506b09c93c548410591b01a0b76ebb Mon Sep 17 00:00:00 2001 From: Anchel135 Date: Mon, 24 Aug 2026 15:24:15 +0300 Subject: [PATCH 06/10] fix(loaders): bound SQL Server connect and query time MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit pymssql.connect was called with no timeouts, so a blackholed network or a stalled server pinned a worker thread indefinitely — and since introspection now runs on the shared executor, enough of those would drain it and stall every other database too. Both connect sites go through _with_timeouts: login_timeout from DB_CONNECT_TIMEOUT, and a query budget of DB_SCHEMA_TIMEOUT for introspection or DB_STATEMENT_TIMEOUT for query execution, matching the MySQL and Snowflake loaders. Also annotates _introspect_schema's return type. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- api/loaders/sqlserver_loader.py | 28 +++++++++++++++++++++++++--- tests/test_sqlserver_loader.py | 31 +++++++++++++++++++++++++++++++ 2 files changed, 56 insertions(+), 3 deletions(-) diff --git a/api/loaders/sqlserver_loader.py b/api/loaders/sqlserver_loader.py index 7dded709..da5fc717 100644 --- a/api/loaders/sqlserver_loader.py +++ b/api/loaders/sqlserver_loader.py @@ -10,6 +10,7 @@ import tqdm import pymssql +from api.config import Config from api.loaders.base_loader import BaseLoader from api.loaders.graph_loader import load_to_graph from api.loaders.introspection import run_introspection @@ -275,7 +276,24 @@ def _parse_sqlserver_url(connection_url: str) -> Dict[str, Any]: return params @staticmethod - def _introspect_schema(conn_params: Dict[str, Any], schema: str): + def _with_timeouts(conn_params: Dict[str, Any], query_timeout: int) -> Dict[str, Any]: + """Bound how long a connect or a query may pin a worker thread. + + Without these, a blackholed network or a stalled server holds a thread + forever and eventually drains the shared introspection executor, taking + every other database with it. ``login_timeout`` covers the TCP/login + handshake and ``timeout`` the query itself, both in seconds. + """ + return { + **conn_params, + 'login_timeout': Config.DB_CONNECT_TIMEOUT, + 'timeout': query_timeout, + } + + @staticmethod + def _introspect_schema( + conn_params: Dict[str, Any], schema: str + ) -> Tuple[Dict[str, Any], Dict[str, List[Dict[str, str]]]]: """Connect, introspect and close — all inside one worker thread. Everything touching the driver lives here so the connection and cursor @@ -290,7 +308,9 @@ def _introspect_schema(conn_params: Dict[str, Any], schema: str): conn = None cursor = None try: - conn = pymssql.connect(**conn_params) # pylint: disable=no-member + conn = pymssql.connect( # pylint: disable=no-member + **SQLServerLoader._with_timeouts(conn_params, Config.DB_SCHEMA_TIMEOUT) + ) cursor = conn.cursor(as_dict=True) entities = SQLServerLoader.extract_tables_info(cursor, schema) @@ -724,7 +744,9 @@ def execute_sql_query(sql_query: str, db_url: str) -> List[Dict[str, Any]]: conn_params = SQLServerLoader._parse_sqlserver_url(db_url) # Connect to SQL Server database - conn = pymssql.connect(**conn_params) # pylint: disable=no-member + conn = pymssql.connect( # pylint: disable=no-member + **SQLServerLoader._with_timeouts(conn_params, Config.DB_STATEMENT_TIMEOUT) + ) cursor = conn.cursor(as_dict=True) # Execute the SQL query diff --git a/tests/test_sqlserver_loader.py b/tests/test_sqlserver_loader.py index af24e837..3374b55c 100644 --- a/tests/test_sqlserver_loader.py +++ b/tests/test_sqlserver_loader.py @@ -20,6 +20,7 @@ # as an unused import that a linter should strip. importlib.import_module("api.core") +from api.config import Config # noqa: E402 pylint: disable=wrong-import-position from api.loaders.sqlserver_loader import ( # noqa: E402 pylint: disable=wrong-import-position SQLServerLoader, SQLServerQueryError, @@ -592,3 +593,33 @@ async def test_load_invalid_url_reports_failure(self): async for success, message in SQLServerLoader.load("user1", "mysql://x/y"): results.append((success, message)) assert results == [(False, "Failed to load SQL Server database schema")] + + @pytest.mark.asyncio + async def test_load_bounds_connect_and_query_time(self): + """Schema loading caps how long a stalled server can pin a worker.""" + cursor = FakeCursor([[], []]) + conn = FakeConnection(cursor) + with patch("pymssql.connect", return_value=conn) as mock_connect, \ + patch("api.loaders.sqlserver_loader.load_to_graph") as mock_load: + async def _noop(*args, **kwargs): + return None + mock_load.side_effect = _noop + async for _ in SQLServerLoader.load( + "user1", "sqlserver://sa:pw@localhost/testdb"): + pass + + kwargs = mock_connect.call_args.kwargs + assert kwargs["login_timeout"] == Config.DB_CONNECT_TIMEOUT + assert kwargs["timeout"] == Config.DB_SCHEMA_TIMEOUT + + def test_execute_query_bounds_connect_and_query_time(self): + """Query execution uses the shorter statement budget, not the schema one.""" + cursor = FakeCursor([[]]) + conn = FakeConnection(cursor) + with patch("pymssql.connect", return_value=conn) as mock_connect: + SQLServerLoader.execute_sql_query( + "SELECT 1", "sqlserver://sa:pw@localhost/testdb") + + kwargs = mock_connect.call_args.kwargs + assert kwargs["login_timeout"] == Config.DB_CONNECT_TIMEOUT + assert kwargs["timeout"] == Config.DB_STATEMENT_TIMEOUT From a0a0cb9a827a45455fc781320e2899f5dbf46b45 Mon Sep 17 00:00:00 2001 From: Anchel135 Date: Mon, 24 Aug 2026 15:38:08 +0300 Subject: [PATCH 07/10] fix(loaders): use one process-wide pymssql timeout MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit pymssql documents that timeout and login_timeout have a process-wide effect, because the FreeTDS db-lib functions behind them are global. Giving schema introspection and query execution different budgets therefore did not give either one its budget — concurrent operations just overwrote each other's, leaving both nondeterministic. Both now use the same value: the larger of DB_SCHEMA_TIMEOUT and DB_STATEMENT_TIMEOUT. It still bounds the wait, and it is the only choice that cannot cut short an operation that was legitimately given the longer budget. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- api/loaders/sqlserver_loader.py | 17 +++++++++++++---- tests/test_sqlserver_loader.py | 6 +++--- 2 files changed, 16 insertions(+), 7 deletions(-) diff --git a/api/loaders/sqlserver_loader.py b/api/loaders/sqlserver_loader.py index da5fc717..70238e0e 100644 --- a/api/loaders/sqlserver_loader.py +++ b/api/loaders/sqlserver_loader.py @@ -276,18 +276,27 @@ def _parse_sqlserver_url(connection_url: str) -> Dict[str, Any]: return params @staticmethod - def _with_timeouts(conn_params: Dict[str, Any], query_timeout: int) -> Dict[str, Any]: + def _with_timeouts(conn_params: Dict[str, Any]) -> Dict[str, Any]: """Bound how long a connect or a query may pin a worker thread. Without these, a blackholed network or a stalled server holds a thread forever and eventually drains the shared introspection executor, taking every other database with it. ``login_timeout`` covers the TCP/login handshake and ``timeout`` the query itself, both in seconds. + + Both are deliberately the *same* for every caller. pymssql documents + that ``timeout`` and ``login_timeout`` "[have] a process-wide effect + because the FreeTDS db-lib API functions used to implement such timeouts + have a global effect" — so handing introspection and query execution + different budgets would just let concurrent operations overwrite each + other's, leaving both nondeterministic. The larger of the two settings + wins: it still bounds the wait, and it is the only choice that cannot + cut short an operation that was legitimately given the longer budget. """ return { **conn_params, 'login_timeout': Config.DB_CONNECT_TIMEOUT, - 'timeout': query_timeout, + 'timeout': max(Config.DB_SCHEMA_TIMEOUT, Config.DB_STATEMENT_TIMEOUT), } @staticmethod @@ -309,7 +318,7 @@ def _introspect_schema( cursor = None try: conn = pymssql.connect( # pylint: disable=no-member - **SQLServerLoader._with_timeouts(conn_params, Config.DB_SCHEMA_TIMEOUT) + **SQLServerLoader._with_timeouts(conn_params) ) cursor = conn.cursor(as_dict=True) @@ -745,7 +754,7 @@ def execute_sql_query(sql_query: str, db_url: str) -> List[Dict[str, Any]]: # Connect to SQL Server database conn = pymssql.connect( # pylint: disable=no-member - **SQLServerLoader._with_timeouts(conn_params, Config.DB_STATEMENT_TIMEOUT) + **SQLServerLoader._with_timeouts(conn_params) ) cursor = conn.cursor(as_dict=True) diff --git a/tests/test_sqlserver_loader.py b/tests/test_sqlserver_loader.py index 3374b55c..a422abad 100644 --- a/tests/test_sqlserver_loader.py +++ b/tests/test_sqlserver_loader.py @@ -610,10 +610,10 @@ async def _noop(*args, **kwargs): kwargs = mock_connect.call_args.kwargs assert kwargs["login_timeout"] == Config.DB_CONNECT_TIMEOUT - assert kwargs["timeout"] == Config.DB_SCHEMA_TIMEOUT + assert kwargs["timeout"] == max(Config.DB_SCHEMA_TIMEOUT, Config.DB_STATEMENT_TIMEOUT) def test_execute_query_bounds_connect_and_query_time(self): - """Query execution uses the shorter statement budget, not the schema one.""" + """Query execution uses the same budget: pymssql timeouts are process-wide.""" cursor = FakeCursor([[]]) conn = FakeConnection(cursor) with patch("pymssql.connect", return_value=conn) as mock_connect: @@ -622,4 +622,4 @@ def test_execute_query_bounds_connect_and_query_time(self): kwargs = mock_connect.call_args.kwargs assert kwargs["login_timeout"] == Config.DB_CONNECT_TIMEOUT - assert kwargs["timeout"] == Config.DB_STATEMENT_TIMEOUT + assert kwargs["timeout"] == max(Config.DB_SCHEMA_TIMEOUT, Config.DB_STATEMENT_TIMEOUT) From 050db3706ff892519754739cfb0cadf9cd82624b Mon Sep 17 00:00:00 2001 From: Anchel135 Date: Mon, 24 Aug 2026 16:23:34 +0300 Subject: [PATCH 08/10] test(loaders): cover SQL Server refresh, routing and URL edge cases Patch coverage on this branch was 85%. refresh_graph_schema was entirely untested despite dropping and reloading a graph, and the sqlserver:// arm of get_database_type_and_loader had no test at all - including the branch that tells an SDK-only install which extra to add instead of failing later on an ImportError. Now 99%. Also covers the malformed-URL rejections, both directions of the encrypt query parameter, the DDL result shape, pymssql.Error on load and on query, and that the connection cleanup helpers stay quiet when the driver throws on the way out - they run on the error path, so raising there would mask the original failure. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .coverage | Bin 0 -> 53248 bytes tests/test_sqlserver_loader.py | 183 ++++++++++++++++++++++++++++++++- 2 files changed, 182 insertions(+), 1 deletion(-) create mode 100644 .coverage diff --git a/.coverage b/.coverage new file mode 100644 index 0000000000000000000000000000000000000000..ca50e52e7e5081a30ed4325956ddc520eca41fd7 GIT binary patch literal 53248 zcmeI533MFAna8US%}meHJvv9yNSd}K>#!w_B!g{?!LqT8Z5#t@b3CrmNNP!g=3-`K zOUAZlZ0`b_ec8NCHerDny~G>HCIKD^rvc06Fi9Y0fn*(gM1gQEILKVFWnt}CHQmxk zwt4C1z0EVP`bquuRsX8`)mPR1&8S*yuUs8SC~{{c7V;h|Io5zU5mVK$I~ZKv0nb;f_GNKM@Fb$%$@7ySg{7bl`gy9D}Q_F zCloFm(XdV?#92dPQs!`Ps11BWU7{w}3nJ)0gD{@#uRuG1)(TrHlnxz2dYqB68G7vqh9B0R5vUsGjDd~ z=QF8vkHF&+{>_?F34CO-3rxEa{O!e$`V-yq>1D@#Inaw3`ba=GQ@|)A|=|ghKvshkC!lx8bxZjQzNC z;(XGB2qcuNpi{=7i%!$jMAfr+#M1QZqre2KhWs7r?v5KJLvcg5a++NFZnVswP!fTV zBI8R)VTfaG@WDts(abM;;xQTs@c#WfE3x1SK zc^I>6&KxGytM=gZ;N*+LlR(I?)m7bPr+C)1lao$%IgHZqh?MDE7rm&|AldY38joeF>PH1QUFgHG-7Gkd)6C2nP~AAE>No0F``wDljQm z@x~OG_Zd(4LXi$77zd?qctOddPbnCI`AM<37bb(ol^9gmmp+2!>c4_qII^KL=PRUxh<@EmMRY1tgZDtH=In?J6AO3Q3ekqp%) zo`6|gAt(*Vtf~Q-B11C0ahN9R!xO@L(ppQaK}#^Cm5B5xVGQSjDh~DQlOVkUk}j>( zpFQzoJXNp>{0}rPodp`X`tn!%LcBMa(G^x#fKaJEAt-{cHxUSu&ftstp@RlCDluHF za{a}cRLkh=D`$d`&cR3Pp`m^(Ing^;%>XH#o})>@?73JZ9#zP+W8CVp?p7zDAn9Tc z6obqwE4gjB0+fWfC^pe1Q{cve|oa;WKU_9KuG7K1}Tpci$!91R9Y$;q>@S2 zpkc&|q{_7WQ?==2g*A)_wKt)P;v&-J^NAiM@A4$r9_e(g?k zu>d-ZJ{{uti9kE%vBCoBb&ky#gg1iMzHXTD4T5JX-y*IAD1Ux3;S9O5ELKBB> z)mIaEn?z>-KlZK6g3NV32asyOi-Tf5cy9sp<4}4OSjvNx&i+fJGFrEp0}Y-1#UG9O zI-y-5FoFtQ!0LX2An%-kq>C614d~;TQFQ$Me-qn;q-U*T)`x`Ggt=B(*lZ0s9+H+z ztYbucOuW^8t?i%e1-7N^vswFXFK2zzoDvsiUE%1KZnu1BS!Z#W_nJmc4_bC`H*gMi zGuz~tX&>h2+VA0h${pqVv+Vqksnb+wUViQezqAGjt>RNIog>g$;{T==KIJebmEpmS zEBREuP8BugnHK*ytl{B>9s0VG_P7bL#q{{!`&HylkN@kJ@Tn}FrBiQ2YW!bkfVZlU z5&zd3SILO~=Pc(_lCj!k#{aV~<5NY(WKN6!YcAnaf-$X$_`iA~*RNATi~p+(WKaA! zBmS>k!KZR{=CAgJsqz1;Z}6#7okEilJKoBS|0@g>Yf>$vua_Gbe6$`qE&iWrz;iUI z>GA&z{jIJAG;%Ya;&g(Z690R?0+MO*f0=+z%$9PJvIK%HZaT*nd$Mr)4-z!5t@kqC8Ja_ z$chFu)W?;H_}^imCaMMw@*upj8}JfUA)_v~OZil`PS-zE{BK>%rx=|&neo41Wbis; z{BJRGa58`{BJUF(>qK2pJl+yrpN!h0WX^x|8x4EAhh_OUBnIO{QM%r zG3opNXiQHDPy&Wp$jEI2~Yx* z03|>PPy&Z4AoKq%?;`1hG$K7N{Y*+qL20d2CfUT_i!Y12 z#jWBxah2FC&UYMhyzO}2@k2+oqtuaWf5X1tzQ_KA{U`Qs+avZ?yUliw?L6!A)~Bs& ztS!Pp;T7Rd;a`Ma!7p4U%n`~34lR@nx0^%BgX6dC%BzR-v#X!n^=X*HH4BneP8}SroXrm2 zUXGHR?_X?b9$@op;aF|$=&tYX-*eHHMMu^iFMsNjr#4vUz~O>AcGjL&X5MVr4m(C2 z^6>W=5Gwq$d-y_xMzCh`BfN%Xgye(J0@yYqG+cs8gc?YfH|I#tE2i%1EJjFOVgLC0 z-QPJnf8ka6tp`RD%piIlYJkv^?a$vYFev#8c3YKm=)SkBeloE6YWdD)374^&tPr64C^RalsZq9xAZich>_F2-C1nL2H6vY|?x<#kPay(jj|sP#N1Idn06 zzoezYmi%9AyaA<1VyiCNMT9IAKs`*#ino*St6hU^2+0UFn`S}2V&e}(6(n2l=+X1% zGAO!+MIT#+PBv{vL)?*L&zVLBICRR5N{)X1?EfINpBw&Ql#$K*jy|@Vu`tPv%U`|E zoG#Wck3`PHA3Ed|@xS z{w3Hx6ad(Fx9vuQ%8q%k+3`UAfq`JkjHG?P@zJ-SP>&YXLR1ufzqcMsjj4c&O$ zYfoDrb)VpzC|^a+E%lWX9&Bc|QyN;b<3H@XTMg zpkKcAJvTH1Lak&yv1P|w4ar}L#pF0tV!?M%D|0D>wm4*HqLLXbdZyCHj+J;ku-jaA z-Mfp?(kse;bf9?$<5{rVe$%i!IdA|9#JI}ff`(Gu?86_-XEs?D4Q^u|lb@L(mcog$ zlDkTg2Q5PP+e{^}FBY?QH`rp9O_;%Dl0q@;DyJ*a-<-M2Qqu(7HC{NHV|SPCIwj`?QOlY$GYdDv@Lt?eto zaH5qTA4N$o+xp^2*YWJ3CC&wlVs~8f=7}JK`owSjh1C9|f$>|9ap=egPq1jnnmqhY z&&XI6oOyhhg~nKPSnES0!lqN7ppg&m9p9Jkb0oF8GD8Q5!nWMHa7I29VwrRF^P)Tp zgNASmg3rOy{AYn>lC4h;Eh@`{OD>+p#t%LkJNWC+!QXFq;kurmKXT~mFzo&5QO24J z=L;IP9(pw7UY3I^V##55pJJ@X$9aw|Twy(ccG`xvmGT^F7Wm|1C)zHh5~KVMuNe)U z9N)B{jI}1eZ*h*8QF8qA?6Khib2cbCrKY(SCv4pIhOk)?0G zSbqiT8>&Nh%gn&uU-@|$4PfBp2G2)*gURM=U`M{t0ay2Bjb6lYYqoFXlQN8JX&`B5 z-BWmS#Y6wzJW_^Qo3UH);H_j38U6yb+U=0I*s`_mN_@s_gFPp|b#Sf~HoCcu17FO) zI@KEJNS43&^5HdK{AILa-`tI! z`(nSt<1#~{hQhmW4mT&~qt|~Z$^8*E|HgN6>* zIW0)8(2^qOjhGukGtDfdbLNg_Wih-Uh%Em8ze(JQq?7RWzo(=}r2C{n={o5F=~-#D zcuf3X@sRW%(l@1u^l#F3=~L;b^lRyW^p5o3(jTRZr95fA)F9PJesQNXQ*w(Bi$9gt zOWo31aY%eZ+$X*z{#N{*PPy&s7I%eVrJzL z(~wI{4*7I|k+R8Nwv!l9BE}>VljR`BVkd^T5yM%D5d>oLEyNU>i77APPy& Date: Tue, 25 Aug 2026 12:51:46 +0300 Subject: [PATCH 09/10] fix(sql): stop a broken-out delimiter passing as a quoted identifier _is_already_quoted took matching outer delimiters as proof a name was quoted, so "[name] DROP TABLE users]" looked pre-quoted, skipped escaping and reached the statement verbatim -- the "]" after "name" closes the identifier early and the remainder parses as SQL. The same held for the standard and MySQL dialects, whose branch never escaped at all: quote_identifier('"a" ; DROP TABLE users --"', '"') -> '"a" ; DROP TABLE users --"' A name now only counts as quoted when every closing delimiter inside the pair is doubled, and quoting escapes by doubling for all three dialects, matching what SnowflakeLoader._quote_identifier already did. Also document that DB_STATEMENT_TIMEOUT is not applied on its own for SQL Server, since pymssql's timeout is process-wide and the loader has to pick one value for both introspection and query execution. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- api/sql_utils/sql_sanitizer.py | 33 +++++++++++++++----------- docs/sqlserver_loader.md | 21 +++++++++++++++++ tests/test_sql_sanitizer.py | 42 ++++++++++++++++++++++++++++++++++ 3 files changed, 83 insertions(+), 13 deletions(-) diff --git a/api/sql_utils/sql_sanitizer.py b/api/sql_utils/sql_sanitizer.py index 421ee472..696c543f 100644 --- a/api/sql_utils/sql_sanitizer.py +++ b/api/sql_utils/sql_sanitizer.py @@ -26,22 +26,30 @@ class SQLIdentifierQuoter: @staticmethod def _is_already_quoted(identifier: str, quote_char: str = '"') -> bool: - """Check if an identifier is already quoted for the active dialect. + """Check if an identifier is already *validly* quoted for the dialect. The pair is scoped to *quote_char* so that a bracketed identifier is only treated as pre-quoted for SQL Server; on PostgreSQL/MySQL a name such as ``[weird]`` is data, not a delimiter, and must still be quoted. + Matching delimiters at the ends are not sufficient: every closing + delimiter *inside* the pair must be doubled, or the first stray one + ends the identifier early and the rest of the string is parsed as SQL. + ``[name] DROP TABLE users]`` must therefore not count as quoted. + Args: identifier: The identifier to inspect. quote_char: Opening delimiter of the active dialect. Returns: - True if *identifier* is already delimited. + True if *identifier* is already delimited and internally escaped. """ - if quote_char == '[': - return identifier.startswith('[') and identifier.endswith(']') - return identifier.startswith(quote_char) and identifier.endswith(quote_char) + close_char = ']' if quote_char == '[' else quote_char + if len(identifier) < 2: + return False + if not (identifier.startswith(quote_char) and identifier.endswith(close_char)): + return False + return close_char not in identifier[1:-1].replace(close_char * 2, '') @classmethod def needs_quoting(cls, identifier: str, quote_char: str = '"') -> bool: @@ -85,14 +93,13 @@ def quote_identifier(identifier: str, quote_char: str = '"') -> str: if SQLIdentifierQuoter._is_already_quoted(identifier, quote_char): return identifier - # SQL Server uses bracket pairs: [identifier]. A literal ``]`` inside the - # name is escaped by doubling it, otherwise it would close the delimiter - # early and change the meaning of the statement. - if quote_char == '[': - escaped = identifier.replace(']', ']]') - return f'[{escaped}]' - - return f'{quote_char}{identifier}{quote_char}' + # Every dialect here escapes its closing delimiter by doubling it -- + # ``]`` for SQL Server brackets, otherwise the quote character itself. + # Without this a name carrying a delimiter closes the identifier early + # and turns the remainder of the name into executable SQL. + close_char = ']' if quote_char == '[' else quote_char + escaped = identifier.replace(close_char, close_char * 2) + return f'{quote_char}{escaped}{close_char}' @classmethod def extract_table_names_from_query(cls, sql_query: str) -> Set[str]: diff --git a/docs/sqlserver_loader.md b/docs/sqlserver_loader.md index 9fb52dc2..6da24af9 100644 --- a/docs/sqlserver_loader.md +++ b/docs/sqlserver_loader.md @@ -121,6 +121,26 @@ Connections are opened with `as_dict=True`, so `pymssql` returns rows as dictionaries keyed by column name. Positional access (`row[0]`) raises `KeyError` with this setting and is never used. +### Timeouts + +`DB_CONNECT_TIMEOUT` maps to `login_timeout` and bounds the connection and login +handshake. + +Query time is bounded by `pymssql`'s `timeout`, which is set to the **larger** of +`DB_SCHEMA_TIMEOUT` and `DB_STATEMENT_TIMEOUT` for every connection this loader +opens. That is a deliberate divergence from the PostgreSQL and Snowflake loaders, +which apply `DB_STATEMENT_TIMEOUT` to query execution on its own: + +> `pymssql` documents that `timeout` and `login_timeout` have a *process-wide* +> effect, because the underlying db-lib API functions used to implement them are +> global. Giving introspection and query execution different budgets would let +> concurrent operations overwrite each other's, leaving both nondeterministic. + +So with the defaults (`DB_SCHEMA_TIMEOUT=300`, `DB_STATEMENT_TIMEOUT=60`), a user +query against SQL Server may run for up to 300s rather than 60s. Lower +`DB_SCHEMA_TIMEOUT` if you need a tighter ceiling; there is no way to bound the +two independently while the driver's timeout stays global. + ## Testing `tests/test_sqlserver_loader.py` covers: @@ -142,3 +162,4 @@ uv run --extra server --extra dev pytest tests/test_sqlserver_loader.py -v - One schema per connection (defaults to `dbo`); connect again to load another - Requires permission to read the `sys.*` catalog views - Windows/Azure AD integrated authentication is not supported; use SQL logins +- `DB_STATEMENT_TIMEOUT` is not applied on its own; see [Timeouts](#timeouts) diff --git a/tests/test_sql_sanitizer.py b/tests/test_sql_sanitizer.py index ac47c5dc..12dbedc4 100644 --- a/tests/test_sql_sanitizer.py +++ b/tests/test_sql_sanitizer.py @@ -284,3 +284,45 @@ def test_bracketed_name_still_quoted_for_postgres(self): def test_bracketed_name_treated_as_quoted_for_sqlserver(self): """The same identifier is already delimited on SQL Server.""" assert SQLIdentifierQuoter.needs_quoting('[weird]', '[') is False + + @pytest.mark.parametrize( + "identifier, quote_char", + [ + ('[name] DROP TABLE users]', '['), + ('"a" ; DROP TABLE users --"', '"'), + ('`a` ; DROP TABLE users --`', '`'), + ], + ) + def test_a_broken_out_delimiter_is_not_mistaken_for_quoting(self, identifier, quote_char): + """Matching outer delimiters are not enough to call a name quoted. + + ``[name] DROP TABLE users]`` opens and closes with a bracket pair, but + the ``]`` after ``name`` ends the identifier early and leaves the rest + to be parsed as SQL. Treating it as pre-quoted skipped the escaping and + emitted the payload verbatim. + """ + assert SQLIdentifierQuoter._is_already_quoted(identifier, quote_char) is False + + quoted = SQLIdentifierQuoter.quote_identifier(identifier, quote_char) + close_char = ']' if quote_char == '[' else quote_char + + # Nothing but the outer pair may act as a delimiter: strip the ends and + # every remaining closing delimiter has to be doubled. + assert quoted.startswith(quote_char) and quoted.endswith(close_char) + assert close_char not in quoted[1:-1].replace(close_char * 2, '') + + def test_quoting_escapes_the_delimiter_for_every_dialect(self): + """Doubling is how all three dialects escape their closing delimiter.""" + assert SQLIdentifierQuoter.quote_identifier('a"b', '"') == '"a""b"' + assert SQLIdentifierQuoter.quote_identifier('a`b', '`') == '`a``b`' + assert SQLIdentifierQuoter.quote_identifier('a]b', '[') == '[a]]b]' + + def test_a_lone_delimiter_is_not_a_quoted_identifier(self): + """One character cannot be an opening and closing pair at once.""" + assert SQLIdentifierQuoter._is_already_quoted('"', '"') is False + assert SQLIdentifierQuoter._is_already_quoted('[', '[') is False + + def test_a_properly_escaped_name_is_left_alone(self): + """``[a]] ; SELECT 1]`` is the escaped name ``a] ; SELECT 1``, not a breakout.""" + assert SQLIdentifierQuoter._is_already_quoted('[a]] ; SELECT 1]', '[') is True + assert SQLIdentifierQuoter.quote_identifier('[a]] ; SELECT 1]', '[') == '[a]] ; SELECT 1]' From 926f91aa70ce0b90453de451f4391af946872dde Mon Sep 17 00:00:00 2001 From: Anchel135 Date: Wed, 26 Aug 2026 15:10:09 +0300 Subject: [PATCH 10/10] fix(sqlserver): refuse a dotted name the sampler cannot disambiguate A dot is legal inside a bracket-quoted SQL Server name, but the sampler recovers the schema and the table from one dotted string, so it has to guess which dot is the separator. `dbo.my.table` was read as `[dbo.my]` dot `[table]` and sampled a different object without saying so. Neither part may contain a dot now, which turns that into a clear error. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- api/loaders/sqlserver_loader.py | 29 +++++++++++++++++++++++------ tests/test_sqlserver_loader.py | 25 +++++++++++++++++++++++++ 2 files changed, 48 insertions(+), 6 deletions(-) diff --git a/api/loaders/sqlserver_loader.py b/api/loaders/sqlserver_loader.py index 70238e0e..6141084f 100644 --- a/api/loaders/sqlserver_loader.py +++ b/api/loaders/sqlserver_loader.py @@ -29,7 +29,9 @@ class SQLServerConnectionError(Exception): """Exception raised for SQL Server connection errors.""" -def validate_ident(identifier: str, identifier_type: str = "identifier") -> str: +def validate_ident( + identifier: str, identifier_type: str = "identifier", allow_dot: bool = True +) -> str: """Validate that an identifier is safe to interpolate into T-SQL. T-SQL cannot bind identifiers as parameters, so table, schema and column @@ -41,6 +43,11 @@ def validate_ident(identifier: str, identifier_type: str = "identifier") -> str: Args: identifier: Raw identifier, typically read from the system catalog. identifier_type: Label used in the error message. + allow_dot: Whether ``.`` is accepted. A dot is legal inside a + bracket-quoted SQL Server name, but callers that recover a schema + and a table from one dotted string cannot tell the two apart, so + they pass ``False`` and get a clear error instead of a query + against the wrong object. Returns: The identifier, unchanged, once validated. @@ -54,10 +61,12 @@ def validate_ident(identifier: str, identifier_type: str = "identifier") -> str: f"Invalid {identifier_type}: {identifier!r}. " "Must be between 1 and 128 characters." ) - if not re.fullmatch(r'[A-Za-z0-9_$#@ .\-]+', identifier): + allowed = r'[A-Za-z0-9_$#@ .\-]+' if allow_dot else r'[A-Za-z0-9_$#@ \-]+' + if not re.fullmatch(allowed, identifier): + dot = "dot, " if allow_dot else "" raise ValueError( f"Invalid {identifier_type}: {identifier!r}. Only letters, digits, " - "underscore, dollar, hash, at-sign, space, dot and dash are allowed." + f"underscore, dollar, hash, at-sign, space, {dot}and dash are allowed." ) return identifier @@ -148,12 +157,20 @@ def _execute_sample_query( SQL Server implementation using TOP with NEWID() for random sampling. ``table_name`` may be schema-qualified (``schema.table``); each part is - bracket-quoted separately so the schema prefix survives. + bracket-quoted separately so the schema prefix survives. A dot is legal + inside a bracket-quoted name, but a single dotted string cannot say + which dot is the separator, so neither part may contain one: a + dot-bearing name is rejected rather than sampled from the wrong object. """ schema, _, bare_table = table_name.rpartition('.') - qualified = quote_ident(validate_ident(bare_table, "table name")) + qualified = quote_ident( + validate_ident(bare_table, "table name", allow_dot=False) + ) if schema: - qualified = f"{quote_ident(validate_ident(schema, 'schema name'))}.{qualified}" + qualified = ( + f"{quote_ident(validate_ident(schema, 'schema name', allow_dot=False))}" + f".{qualified}" + ) col = quote_ident(validate_ident(col_name, "column name")) if not isinstance(sample_size, int) or sample_size <= 0: diff --git a/tests/test_sqlserver_loader.py b/tests/test_sqlserver_loader.py index 8cc3e789..9b0e6902 100644 --- a/tests/test_sqlserver_loader.py +++ b/tests/test_sqlserver_loader.py @@ -152,6 +152,18 @@ def test_error_names_the_identifier_type(self): with pytest.raises(ValueError, match="schema name"): validate_ident("bad;name", "schema name") + def test_dot_can_be_disallowed(self): + """Callers that split a dotted string opt out of accepting dots.""" + assert validate_ident("a.b") == "a.b" + with pytest.raises(ValueError, match="table name"): + validate_ident("a.b", "table name", allow_dot=False) + + def test_message_drops_dot_when_it_is_disallowed(self): + """The allow-list in the message matches the one actually applied.""" + with pytest.raises(ValueError) as excinfo: + validate_ident("a.b", allow_dot=False) + assert "dot" not in str(excinfo.value) + class TestSampleQueryValidation: """The sample query refuses hostile identifiers outright.""" @@ -176,6 +188,19 @@ def test_invalid_sample_size_rejected(self, size): SQLServerLoader._execute_sample_query(cursor, "dbo.T", "c", sample_size=size) assert cursor.executed == [] + @pytest.mark.parametrize("table", ["dbo.my.table", "my.table.T"]) + def test_ambiguous_dotted_name_is_refused(self, table): + """A dot is legal in a bracket-quoted name but not recoverable here. + + ``rpartition`` has to guess which dot separates the schema from the + table, so ``dbo.my.table`` would otherwise be sampled as ``[dbo.my]`` + dot ``[table]`` -- a different object, silently. + """ + cursor = FakeCursor([[]]) + with pytest.raises(ValueError): + SQLServerLoader._execute_sample_query(cursor, table, "c") + assert cursor.executed == [] + class TestParseUrl: """URL parsing."""