From d7876255a9d1e275defe02368a8fece5c9ea641a Mon Sep 17 00:00:00 2001 From: Guy Korland Date: Thu, 13 Aug 2026 21:09:22 +0300 Subject: [PATCH 01/15] 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/15] 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/15] 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/15] 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/15] 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/15] 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/15] 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/15] 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/15] 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/15] 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.""" From ec7cc04db0139e8aaa21587f893c5938d9b0807c Mon Sep 17 00:00:00 2001 From: Anchel135 Date: Mon, 7 Sep 2026 10:20:38 +0300 Subject: [PATCH 11/15] fix(sqlserver): scope foreign keys to the loaded schema extract_foreign_keys pinned the parent table to the requested schema but left the referenced side unconstrained, so a cross-schema key could point at a table that was never loaded. Constrain rs.name too, matching extract_relationships. Also drop the redundant unquote() on the schema URL parameter: parse_qs already percent-decodes, and decoding twice let double-encoded values past validate_ident. --- api/loaders/sqlserver_loader.py | 13 ++++++++----- docs/sqlserver_loader.md | 2 +- tests/test_sqlserver_loader.py | 13 ++++++++++--- 3 files changed, 19 insertions(+), 9 deletions(-) diff --git a/api/loaders/sqlserver_loader.py b/api/loaders/sqlserver_loader.py index 6141084f..cfe33ae4 100644 --- a/api/loaders/sqlserver_loader.py +++ b/api/loaders/sqlserver_loader.py @@ -233,8 +233,9 @@ def parse_schema_from_url(connection_url: str) -> str: """ try: parsed = urlparse(connection_url) - schema = parse_qs(parsed.query).get('schema', [''])[0] - schema = unquote(schema).strip() + # parse_qs already percent-decodes; decoding again would accept + # double-encoded values that validate_ident should reject. + schema = parse_qs(parsed.query).get('schema', [''])[0].strip() except (ValueError, AttributeError): return DEFAULT_SCHEMA if not schema: @@ -559,6 +560,9 @@ def extract_foreign_keys(cursor, schema: str, table_name: str) -> List[Dict[str, """ Extract foreign key information for a specific table. + Only foreign keys whose referenced table also lives in *schema* are + returned, so they never point at a table outside the loaded schema. + Args: cursor: Database cursor schema: Schema owning the table @@ -572,7 +576,6 @@ def extract_foreign_keys(cursor, schema: str, table_name: str) -> List[Dict[str, 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 @@ -589,9 +592,9 @@ def extract_foreign_keys(cursor, schema: str, table_name: str) -> List[Dict[str, 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 + WHERE ps.name = %s AND rs.name = %s AND pt.name = %s ORDER BY fk.name; - """, (schema, table_name)) + """, (schema, schema, table_name)) foreign_keys = [] for fk_info in cursor.fetchall(): diff --git a/docs/sqlserver_loader.md b/docs/sqlserver_loader.md index 6da24af9..fb569edc 100644 --- a/docs/sqlserver_loader.md +++ b/docs/sqlserver_loader.md @@ -49,7 +49,7 @@ sqlserver://appuser:s3cr3t@sql.example.com:1433/reporting?schema=dbo&encrypt=tru - 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 +- Foreign keys, including composite keys — both sides must live in the selected schema - Many-to-many relationships inferred from junction tables All catalog queries join `sys.schemas` and bind the schema name as a parameter, so diff --git a/tests/test_sqlserver_loader.py b/tests/test_sqlserver_loader.py index 9b0e6902..7ad3cc76 100644 --- a/tests/test_sqlserver_loader.py +++ b/tests/test_sqlserver_loader.py @@ -261,6 +261,12 @@ def test_schema_from_url(self): assert SQLServerLoader.parse_schema_from_url( "sqlserver://sa:pw@localhost/testdb?schema=sales") == "sales" + def test_double_encoded_schema_is_rejected(self): + """The parameter is decoded once, so ``%2520`` stays a literal ``%20``.""" + with pytest.raises(ValueError): + SQLServerLoader.parse_schema_from_url( + "sqlserver://sa:pw@localhost/testdb?schema=sa%2520les") + class TestSampleQuery: """Sample-value extraction — the dict-cursor contract.""" @@ -376,7 +382,6 @@ def test_foreign_keys_mapping(self): "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") @@ -386,8 +391,10 @@ def test_foreign_keys_mapping(self): "referenced_table": "Customers", "referenced_column": "id", }] - _, params = cursor.executed[0] - assert params == ("dbo", "Orders") + query, params = cursor.executed[0] + # Both sides of the key are pinned to the loaded schema. + assert "rs.name = %s" in query + assert params == ("dbo", "dbo", "Orders") def test_relationships_grouped_by_constraint(self): """Composite keys are grouped under one constraint name.""" From b6b8334c9c344a159c7a55915e8027b621ae4b7c Mon Sep 17 00:00:00 2001 From: Anchel135 Date: Mon, 7 Sep 2026 10:50:02 +0300 Subject: [PATCH 12/15] fix(sqlserver): make schema introspection valid, resilient and cheaper `SELECT DISTINCT ... ORDER BY NEWID()` is not valid T-SQL (error 145), so sampling failed on the first column of the first table and took the whole schema load with it. The DISTINCT now happens in a derived table. Sampling is also best-effort: a column type SQL Server cannot compare (`xml`, `text`, `image`, spatial) or a name outside the old ASCII allow-list used to abort the load. Catalog identifiers now go through a validator that only rejects what `quote_ident` cannot make safe, and a failed sample costs that one column instead of the run. `uniqueidentifier` values reach the result stream as strings rather than breaking `json.dumps`. Foreign keys are read once for the whole schema instead of once per table, and feed both the column key kinds and the relationships. The `fk`/`uc` joins are deduplicated so a composite key cannot multiply column rows. Drops the unreachable `mssql` alias and corrects the docs to describe what the loader actually does. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- api/core/pipeline.py | 1 - api/loaders/sqlserver_loader.py | 274 ++++++++++++++++----------- api/sql_utils/sql_sanitizer.py | 2 +- docs/sqlserver_loader.md | 17 +- tests/test_destructive_detection.py | 3 - tests/test_schema_load_offloading.py | 2 +- tests/test_sql_sanitizer.py | 1 - tests/test_sqlserver_loader.py | 168 +++++++++++----- 8 files changed, 297 insertions(+), 171 deletions(-) diff --git a/api/core/pipeline.py b/api/core/pipeline.py index 5ba8d426..f7731826 100644 --- a/api/core/pipeline.py +++ b/api/core/pipeline.py @@ -218,7 +218,6 @@ def truncate_for_log(query: str, max_length: int = 200) -> str: "mysql": "mysql", "snowflake": "snowflake", "sqlserver": "tsql", - "mssql": "tsql", } # sqlglot expression class names that represent a write, DDL, privilege change, diff --git a/api/loaders/sqlserver_loader.py b/api/loaders/sqlserver_loader.py index cfe33ae4..fc4fa720 100644 --- a/api/loaders/sqlserver_loader.py +++ b/api/loaders/sqlserver_loader.py @@ -4,6 +4,7 @@ import decimal import logging import re +import uuid from typing import AsyncGenerator, Dict, Any, List, Tuple from urllib.parse import urlparse, parse_qs, unquote @@ -20,6 +21,8 @@ DEFAULT_SCHEMA = "dbo" DEFAULT_PORT = 1433 +_CONTROL_CHARS = re.compile(r'[\x00-\x1f\x7f]') + class SQLServerQueryError(Exception): """Exception raised for SQL Server query execution errors.""" @@ -32,16 +35,18 @@ class SQLServerConnectionError(Exception): 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. + """Validate that a *user-supplied* 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. + T-SQL cannot bind identifiers as parameters, so the schema taken from the + connection URL 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. Names read back from the system catalog go through + :func:`validate_catalog_ident` instead, which is deliberately laxer. Args: - identifier: Raw identifier, typically read from the system catalog. + identifier: Raw identifier supplied by the user. 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 @@ -71,6 +76,40 @@ def validate_ident( return identifier +def validate_catalog_ident( + identifier: str, identifier_type: str = "identifier", allow_dot: bool = True +) -> str: + """Validate an identifier that the server itself returned from the catalog. + + SQL Server allows almost anything inside a delimited identifier, and Hebrew, + German and CJK table names are ordinary in the databases this loader reads. + Holding catalog names to the ASCII allow-list above would fail the whole + schema load over one such name, so this only rejects what bracket-quoting + cannot survive: control characters, and — when *allow_dot* is False — a dot + that would make a schema-qualified name ambiguous. + + Raises: + ValueError: If the identifier is empty, over-long, contains a control + character, or contains a dot that the caller cannot allow. + """ + if not identifier or len(identifier) > 128: + raise ValueError( + f"Invalid {identifier_type}: {identifier!r}. " + "Must be between 1 and 128 characters." + ) + if _CONTROL_CHARS.search(identifier): + raise ValueError( + f"Invalid {identifier_type}: {identifier!r}. " + "Control characters are not allowed." + ) + if not allow_dot and '.' in identifier: + raise ValueError( + f"Invalid {identifier_type}: {identifier!r}. A dot would make a " + "schema-qualified name ambiguous." + ) + return identifier + + def quote_ident(identifier: str) -> str: """Bracket-quote a T-SQL identifier, escaping any embedded ``]``. @@ -79,7 +118,8 @@ def quote_ident(identifier: str) -> str: 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. + statement validate the identifier with :func:`validate_ident` or + :func:`validate_catalog_ident` first. Args: identifier: Raw identifier as read from the system catalog. @@ -90,13 +130,6 @@ 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. @@ -164,24 +197,28 @@ def _execute_sample_query( """ schema, _, bare_table = table_name.rpartition('.') qualified = quote_ident( - validate_ident(bare_table, "table name", allow_dot=False) + validate_catalog_ident(bare_table, "table name", allow_dot=False) ) if schema: qualified = ( - f"{quote_ident(validate_ident(schema, 'schema name', allow_dot=False))}" + f"{quote_ident(validate_catalog_ident(schema, 'schema name', allow_dot=False))}" f".{qualified}" ) - col = quote_ident(validate_ident(col_name, "column name")) + col = quote_ident(validate_catalog_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. + # Identifiers are validated and bracket-quoted with ``]`` escaped, since + # T-SQL cannot bind identifiers as parameters. + # + # The DISTINCT sits in a derived table because SQL Server rejects + # ``SELECT DISTINCT ... ORDER BY NEWID()`` outright: "ORDER BY items + # must appear in the select list if SELECT DISTINCT is specified". query = ( - f"SELECT DISTINCT TOP {int(sample_size)} {col}" - f" FROM {qualified}" - f" WHERE {col} IS NOT NULL" + f"SELECT TOP {int(sample_size)} {col}" + f" FROM (SELECT DISTINCT {col} FROM {qualified}" + f" WHERE {col} IS NOT NULL) AS sampled" f" ORDER BY NEWID()" ) cursor.execute(query) @@ -191,6 +228,30 @@ def _execute_sample_query( sample_results = cursor.fetchall() return [row[col_name] for row in sample_results if row[col_name] is not None] + @classmethod + def extract_sample_values_for_column( + cls, cursor, table_name: str, col_name: str, sample_size: int = 3 + ) -> List[Any]: + """Sample *col_name*, returning ``[]`` rather than failing the whole load. + + Sampling is best-effort decoration on top of the catalog data, but it is + also the only part of introspection that touches user tables, so it is + where the surprises live: ``DISTINCT`` is not defined for ``xml``, + ``text``, ``image`` or the spatial types, and a name that survives the + catalog can still be one this loader will not interpolate. Either would + otherwise abort a schema that is fine apart from one column. + """ + try: + return super().extract_sample_values_for_column( + cursor, table_name, col_name, sample_size + ) + except (ValueError, pymssql.Error) as exc: + # %r so a control character in a catalog name cannot forge a log line. + logging.warning( + "Skipping sample values for %r.%r: %s", table_name, col_name, exc + ) + return [] + @staticmethod def _serialize_value(value): """ @@ -202,16 +263,16 @@ def _serialize_value(value): Returns: JSON serializable version of the value """ - if isinstance(value, (datetime.date, datetime.datetime)): - return value.isoformat() - if isinstance(value, datetime.time): + if isinstance(value, (datetime.date, datetime.datetime, datetime.time)): return value.isoformat() if isinstance(value, decimal.Decimal): return float(value) + # pymssql decodes ``uniqueidentifier`` to uuid.UUID, which the result + # stream's plain json.dumps cannot encode. + if isinstance(value, uuid.UUID): + return str(value) if isinstance(value, bytes): return value.hex() - if value is None: - return None return value @staticmethod @@ -340,8 +401,11 @@ def _introspect_schema( ) cursor = conn.cursor(as_dict=True) - entities = SQLServerLoader.extract_tables_info(cursor, schema) - relationships = SQLServerLoader.extract_relationships(cursor, schema) + foreign_keys = SQLServerLoader.extract_foreign_keys(cursor, schema) + entities = SQLServerLoader.extract_tables_info( + cursor, schema, SQLServerLoader.group_foreign_keys(foreign_keys) + ) + relationships = SQLServerLoader.build_relationships(foreign_keys) return entities, relationships finally: SQLServerLoader._close_quietly(cursor, conn) @@ -407,18 +471,27 @@ def _close_quietly(cursor, conn) -> None: 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]: + def extract_tables_info( + cursor, + schema: str = DEFAULT_SCHEMA, + foreign_keys_by_table: Dict[str, List[Dict[str, str]]] | None = None, + ) -> 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``) + foreign_keys_by_table: Foreign keys for the whole schema, keyed by + owning table, as produced by :meth:`group_foreign_keys`. Passed + in rather than queried per table, which would be one round trip + per table on a large schema. Returns: Dict containing table information """ entities = {} + foreign_keys_by_table = foreign_keys_by_table or {} # Get all tables in the requested schema. ``s.name`` is selected back so # sample queries qualify tables with the server's own canonical schema @@ -452,9 +525,6 @@ def extract_tables_info(cursor, schema: str = DEFAULT_SCHEMA) -> Dict[str, Any]: cursor, schema, table_name, catalog_schema ) - # 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}" @@ -464,7 +534,7 @@ def extract_tables_info(cursor, schema: str = DEFAULT_SCHEMA) -> Dict[str, Any]: entities[table_name] = { 'description': table_description, 'columns': columns_info, - 'foreign_keys': foreign_keys, + 'foreign_keys': foreign_keys_by_table.get(table_name, []), 'col_descriptions': col_descriptions } @@ -497,10 +567,10 @@ def extract_columns_info( 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 '' + WHEN pk.column_id IS NOT NULL THEN 'PRIMARY KEY' + WHEN fk.parent_column_id IS NOT NULL THEN 'FOREIGN KEY' + WHEN uc.column_id IS NOT NULL THEN 'UNIQUE KEY' + ELSE 'NONE' END AS column_key, ISNULL(CAST(ep.value AS NVARCHAR(MAX)), '') AS column_comment FROM sys.columns c @@ -514,10 +584,12 @@ def extract_columns_info( 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 + SELECT DISTINCT fkc.parent_object_id, fkc.parent_column_id + FROM sys.foreign_key_columns fkc + ) fk ON fk.parent_object_id = c.object_id AND fk.parent_column_id = c.column_id + LEFT JOIN ( + SELECT DISTINCT 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 @@ -539,7 +611,7 @@ def extract_columns_info( for col_info in columns: col_name = col_info['column_name'] is_nullable = 'YES' if col_info['is_nullable'] else 'NO' - key_type = _KEY_TYPES.get(col_info['column_key'], 'NONE') + key_type = col_info['column_key'] columns_info[col_name] = { 'type': col_info['data_type'], @@ -556,74 +628,24 @@ def extract_columns_info( 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. - - Only foreign keys whose referenced table also lives in *schema* are - returned, so they never point at a table outside the loaded schema. - - 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, - 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 rs.name = %s AND pt.name = %s - ORDER BY fk.name; - """, (schema, 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( + def extract_foreign_keys( cursor, schema: str = DEFAULT_SCHEMA - ) -> Dict[str, List[Dict[str, str]]]: + ) -> List[Dict[str, str]]: """ - Extract all relationship information from a schema. + Extract every foreign key in a schema, in a single query. - Only foreign keys whose parent *and* referenced tables both live in - *schema* are returned, so relationships always point at entities that - were actually loaded. + Both the parent and the referenced table must live in *schema*, so a key + never names a table that was not loaded. The per-entity view and the + relationship map are both derived from these rows — see + :meth:`group_foreign_keys` and :meth:`build_relationships` — rather than + from a second query per table. Args: cursor: Database cursor - schema: Schema to extract relationships from (default: ``dbo``) + schema: Schema to extract foreign keys from (default: ``dbo``) Returns: - Dict containing relationship information + One dict per foreign-key column, in constraint order. """ cursor.execute(""" SELECT @@ -651,21 +673,45 @@ def extract_relationships( 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'] + return [{ + 'table': fk_info['table_name'], + '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'], + } for fk_info in cursor.fetchall()] - if constraint_name not in relationships: - relationships[constraint_name] = [] + @staticmethod + def group_foreign_keys( + foreign_keys: List[Dict[str, str]] + ) -> Dict[str, List[Dict[str, str]]]: + """Group :meth:`extract_foreign_keys` rows by the table that owns them.""" + by_table: Dict[str, List[Dict[str, str]]] = {} + for fk in foreign_keys: + by_table.setdefault(fk['table'], []).append( + {key: value for key, value in fk.items() if key != 'table'} + ) + return by_table - 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'], + @staticmethod + def build_relationships( + foreign_keys: List[Dict[str, str]] + ) -> Dict[str, List[Dict[str, str]]]: + """Group :meth:`extract_foreign_keys` rows by constraint. + + A composite key contributes one row per column, all under the one + constraint name. + """ + relationships: Dict[str, List[Dict[str, str]]] = {} + for fk in foreign_keys: + constraint_name = fk['constraint_name'] + relationships.setdefault(constraint_name, []).append({ + 'from': fk['table'], + 'to': fk['referenced_table'], + 'source_column': fk['column'], + 'target_column': fk['referenced_column'], 'note': f'Foreign key constraint: {constraint_name}' }) - return relationships @staticmethod diff --git a/api/sql_utils/sql_sanitizer.py b/api/sql_utils/sql_sanitizer.py index 696c543f..b05f35d8 100644 --- a/api/sql_utils/sql_sanitizer.py +++ b/api/sql_utils/sql_sanitizer.py @@ -200,7 +200,7 @@ def get_quote_char(db_type: str) -> str: """ if db_type.lower() in ['mysql', 'mariadb']: return '`' - if db_type.lower() in ['sqlserver', 'mssql']: + if db_type.lower() == 'sqlserver': return '[' # PostgreSQL, SQLite use double quotes (standard SQL) return '"' diff --git a/docs/sqlserver_loader.md b/docs/sqlserver_loader.md index fb569edc..9a39091b 100644 --- a/docs/sqlserver_loader.md +++ b/docs/sqlserver_loader.md @@ -50,7 +50,6 @@ sqlserver://appuser:s3cr3t@sql.example.com:1433/reporting?schema=dbo&encrypt=tru - 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 — both sides must live in the selected schema -- 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 @@ -59,14 +58,20 @@ 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: +query. The `DISTINCT` sits in a derived table because SQL Server rejects +`SELECT DISTINCT … ORDER BY NEWID()`: ```sql -SELECT DISTINCT TOP 3 [column_name] -FROM [dbo].[table_name] -WHERE [column_name] IS NOT NULL; +SELECT TOP 3 [column_name] +FROM (SELECT DISTINCT [column_name] FROM [dbo].[table_name] + WHERE [column_name] IS NOT NULL) AS sampled +ORDER BY NEWID(); ``` +Sampling is best-effort: a column whose type has no `DISTINCT` (`xml`, `text`, +`image`, the spatial types) is logged and left without samples rather than +failing the schema load. + ### Query Execution - Executes SQL against the connected database @@ -81,7 +86,7 @@ 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`. +returns `[` for `sqlserver`. ## Usage diff --git a/tests/test_destructive_detection.py b/tests/test_destructive_detection.py index 3ec1c612..f85c13dc 100644 --- a/tests/test_destructive_detection.py +++ b/tests/test_destructive_detection.py @@ -333,6 +333,3 @@ def test_reads_are_not_destructive(self, sql): ]) 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_schema_load_offloading.py b/tests/test_schema_load_offloading.py index 16ca3a3d..ebba8c25 100644 --- a/tests/test_schema_load_offloading.py +++ b/tests/test_schema_load_offloading.py @@ -102,7 +102,7 @@ async def noop(*_args, **_kwargs): @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_foreign_keys", _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): diff --git a/tests/test_sql_sanitizer.py b/tests/test_sql_sanitizer.py index 12dbedc4..03d81a3b 100644 --- a/tests/test_sql_sanitizer.py +++ b/tests/test_sql_sanitizer.py @@ -250,7 +250,6 @@ 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.""" diff --git a/tests/test_sqlserver_loader.py b/tests/test_sqlserver_loader.py index 7ad3cc76..f2a96f34 100644 --- a/tests/test_sqlserver_loader.py +++ b/tests/test_sqlserver_loader.py @@ -9,8 +9,11 @@ import datetime import decimal import importlib +import re +import uuid from unittest.mock import AsyncMock, patch, MagicMock +import pymssql import pytest # ``api.core`` must be initialised before any loader module is imported. @@ -27,6 +30,7 @@ SQLServerLoader, SQLServerQueryError, quote_ident, + validate_catalog_ident, validate_ident, ) @@ -114,7 +118,7 @@ def test_injection_attempt_stays_contained(self): class TestValidateIdent: - """Allow-list validation applied before any identifier interpolation.""" + """Allow-list validation applied to the schema taken from the connection URL.""" @pytest.mark.parametrize("name", [ "Orders", "my-table name", "col_1", "tbl$", "#temp", "a.b", "x@y", @@ -165,20 +169,54 @@ def test_message_drops_dot_when_it_is_disallowed(self): assert "dot" not in str(excinfo.value) +class TestValidateCatalogIdent: + """Catalog names get a laxer check than user-supplied ones.""" + + @pytest.mark.parametrize("name", [ + "Kunden_Ä", "לקוחות", "顧客", "my]table", "tbl'; DROP TABLE t --", + ]) + def test_accepts_names_the_server_accepts(self, name): + """Non-ASCII names are ordinary; bracket-quoting is what makes them safe.""" + assert validate_catalog_ident(name) == name + + @pytest.mark.parametrize("name", ["tbl\nDROP", "tbl\x00", "tbl\x7f"]) + def test_rejects_control_characters(self, name): + """Control characters have no business in an object name.""" + with pytest.raises(ValueError): + validate_catalog_ident(name) + + def test_rejects_empty_and_over_long(self): + """The same length bounds as the strict check.""" + with pytest.raises(ValueError): + validate_catalog_ident("") + with pytest.raises(ValueError): + validate_catalog_ident("a" * 129) + + def test_dot_can_be_disallowed(self): + """Callers that split a dotted string opt out of accepting dots.""" + assert validate_catalog_ident("a.b") == "a.b" + with pytest.raises(ValueError, match="table name"): + validate_catalog_ident("a.b", "table name", allow_dot=False) + + class TestSampleQueryValidation: - """The sample query refuses hostile identifiers outright.""" + """The sample query contains hostile identifiers rather than trusting them.""" - @pytest.mark.parametrize("table,column", [ - ("dbo.x] FROM sys.tables --", "c"), - ("dbo.T", "c] FROM sys.tables --"), - ("bad;schema.T", "c"), + @pytest.mark.parametrize("table,column,quoted", [ + ("dbo.x] FROM t --", "c", "[dbo].[x]] FROM t --]"), + ("dbo.T", "c] FROM sys.tables --", "[c]] FROM sys.tables --]"), + ("bad;schema.T", "c", "[bad;schema].[T]"), ]) - def test_hostile_identifier_is_rejected(self, table, column): - """Validation happens before the statement is built or executed.""" + def test_hostile_identifier_stays_quoted(self, table, column, quoted): + """A catalog name is bracket-quoted with ``]`` doubled, never rejected. + + SQL Server permits these names, so refusing them would fail the whole + schema load; quoting is what keeps them inert. + """ cursor = FakeCursor([[]]) - with pytest.raises(ValueError): - SQLServerLoader._execute_sample_query(cursor, table, column) - assert cursor.executed == [] + SQLServerLoader._execute_sample_query(cursor, table, column) + query, _ = cursor.executed[0] + assert quoted in query @pytest.mark.parametrize("size", [0, -1, "5"]) def test_invalid_sample_size_rejected(self, size): @@ -293,6 +331,20 @@ def test_query_is_schema_qualified_and_quoted(self): SQLServerLoader._execute_sample_query(cursor, "sales.Orders", "status") query, _ = cursor.executed[0] assert "FROM [sales].[Orders]" in query + + def test_distinct_is_isolated_in_a_derived_table(self): + """SQL Server rejects ``SELECT DISTINCT … ORDER BY NEWID()`` outright. + + Regression test for error 145, "ORDER BY items must appear in the select + list if SELECT DISTINCT is specified", which failed every schema load on + the first table with a column. + """ + cursor = FakeCursor([[]]) + SQLServerLoader._execute_sample_query(cursor, "dbo.Orders", "status") + query, _ = cursor.executed[0] + assert re.search(r"SELECT\s+DISTINCT\s+TOP", query) is None + assert ") AS sampled" in query + assert query.index("AS sampled") < query.index("ORDER BY NEWID()") assert "[status]" in query def test_bare_table_name_still_works(self): @@ -315,6 +367,17 @@ def test_extract_sample_values_stringifies(self): assert SQLServerLoader.extract_sample_values_for_column( cursor, "dbo.T", "n") == ["1", "2"] + @pytest.mark.parametrize("error", [ + pymssql.Error("DISTINCT is not defined for xml"), + ValueError("Invalid table name"), + ]) + def test_a_failed_sample_costs_only_that_column(self, error): + """One non-comparable type or odd name must not fail the whole load.""" + cursor = FakeCursor([[]]) + cursor.execute = MagicMock(side_effect=error) + assert SQLServerLoader.extract_sample_values_for_column( + cursor, "dbo.T", "payload") == [] + class TestIntrospection: """Catalog introspection queries.""" @@ -344,7 +407,7 @@ def test_columns_info_mapping(self): "data_type": "int", "is_nullable": False, "column_default": None, - "column_key": "PRI", + "column_key": "PRIMARY KEY", "column_comment": "", }], [{"id": 1}], # sample values query @@ -364,7 +427,7 @@ def test_columns_sample_query_is_schema_qualified(self): "data_type": "int", "is_nullable": True, "column_default": None, - "column_key": "", + "column_key": "NONE", "column_comment": "", }], [], @@ -379,54 +442,70 @@ def test_columns_sample_query_is_schema_qualified(self): def test_foreign_keys_mapping(self): """Foreign key rows map onto the loader's FK dicts.""" cursor = FakeCursor([[{ + "table_name": "Orders", "constraint_name": "FK_Orders_Customers", "column_name": "customer_id", "referenced_table_name": "Customers", "referenced_column_name": "id", }]]) - fks = SQLServerLoader.extract_foreign_keys(cursor, "dbo", "Orders") - assert fks == [{ + fks = SQLServerLoader.extract_foreign_keys(cursor, "dbo") + assert SQLServerLoader.group_foreign_keys(fks) == {"Orders": [{ "constraint_name": "FK_Orders_Customers", "column": "customer_id", "referenced_table": "Customers", "referenced_column": "id", - }] + }]} + + def test_foreign_keys_are_fetched_once_for_the_whole_schema(self): + """One query for every key, with both sides pinned to the loaded schema.""" + cursor = FakeCursor([[]]) + SQLServerLoader.extract_foreign_keys(cursor, "sales") query, params = cursor.executed[0] - # Both sides of the key are pinned to the loaded schema. - assert "rs.name = %s" in query - assert params == ("dbo", "dbo", "Orders") + assert len(cursor.executed) == 1 + assert params == ("sales", "sales") + assert "ps.name = %s AND rs.name = %s" in query def test_relationships_grouped_by_constraint(self): """Composite keys are grouped under one constraint name.""" - cursor = FakeCursor([[ + rels = SQLServerLoader.build_relationships([ { - "table_name": "Orders", + "table": "Orders", "constraint_name": "FK_A", - "column_name": "c1", - "referenced_table_name": "Customers", - "referenced_column_name": "id1", + "column": "c1", + "referenced_table": "Customers", + "referenced_column": "id1", }, { - "table_name": "Orders", + "table": "Orders", "constraint_name": "FK_A", - "column_name": "c2", - "referenced_table_name": "Customers", - "referenced_column_name": "id2", + "column": "c2", + "referenced_table": "Customers", + "referenced_column": "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_entities_and_relationships_share_one_fk_query(self): + """The per-table list and the relationship map come from the same rows.""" + fks = [{ + "table": "Orders", + "constraint_name": "FK_A", + "column": "customer_id", + "referenced_table": "Customers", + "referenced_column": "id", + }] + cursor = FakeCursor([ + [{"table_name": "Orders", "schema_name": "dbo", "table_comment": ""}], + [], # columns + ]) + entities = SQLServerLoader.extract_tables_info( + cursor, "dbo", SQLServerLoader.group_foreign_keys(fks) + ) + assert entities["Orders"]["foreign_keys"][0]["referenced_table"] == "Customers" + assert list(SQLServerLoader.build_relationships(fks)) == ["FK_A"] def test_tables_info_builds_entities(self): """A full table walk produces the expected entity structure.""" @@ -437,11 +516,10 @@ def test_tables_info_builds_entities(self): "data_type": "int", "is_nullable": False, "column_default": None, - "column_key": "PRI", + "column_key": "PRIMARY KEY", "column_comment": "", }], [{"id": 7}], - [], # foreign keys ]) entities = SQLServerLoader.extract_tables_info(cursor, "dbo") assert list(entities) == ["Orders"] @@ -462,11 +540,10 @@ def test_sample_query_uses_catalog_schema_not_url_schema(self): "data_type": "int", "is_nullable": False, "column_default": None, - "column_key": "PRI", + "column_key": "PRIMARY KEY", "column_comment": "", }], [{"id": 7}], - [], # foreign keys ]) SQLServerLoader.extract_tables_info(cursor, "sales") sample_query = cursor.executed[2][0] @@ -482,6 +559,9 @@ class TestSerialization: (datetime.time(3, 4, 5), "03:04:05"), (decimal.Decimal("1.5"), 1.5), (b"\x01\x02", "0102"), + # pymssql decodes uniqueidentifier to uuid.UUID, which json.dumps rejects. + (uuid.UUID("3f2504e0-4f89-11d3-9a0c-0305e82c3301"), + "3f2504e0-4f89-11d3-9a0c-0305e82c3301"), (None, None), ("plain", "plain"), ]) @@ -568,10 +648,9 @@ class TestLoad: async def test_load_success_closes_connection(self): """A successful load reports table count and releases resources.""" cursor = FakeCursor([ + [], # foreign keys [{"table_name": "Orders", "schema_name": "dbo", "table_comment": ""}], [], # columns - [], # foreign keys - [], # relationships ]) conn = FakeConnection(cursor) messages = [] @@ -603,7 +682,8 @@ async def _noop(*args, **kwargs): async for _ in SQLServerLoader.load( "user1", "sqlserver://sa:pw@localhost/testdb?schema=sales"): pass - assert cursor.executed[0][1] == ("sales",) + assert cursor.executed[0][1] == ("sales", "sales") + assert cursor.executed[1][1] == ("sales",) @pytest.mark.asyncio async def test_load_failure_closes_connection(self): From e99cf1630744d67bd5e2470fb7aed3782d2d97b3 Mon Sep 17 00:00:00 2001 From: Anchel135 Date: Mon, 7 Sep 2026 10:58:17 +0300 Subject: [PATCH 13/15] refactor(sqlserver): stop reconfiguring global logging at import `logging.basicConfig` at module scope reaches the root logger, so importing the loader could override the format and level `app_factory` installs for the whole process. Records still propagate; only the configuration call goes. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- api/loaders/sqlserver_loader.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/api/loaders/sqlserver_loader.py b/api/loaders/sqlserver_loader.py index fc4fa720..c581a063 100644 --- a/api/loaders/sqlserver_loader.py +++ b/api/loaders/sqlserver_loader.py @@ -16,8 +16,6 @@ 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") - DEFAULT_SCHEMA = "dbo" DEFAULT_PORT = 1433 From 9d1dc67bd7b7c8d7b1c5f42ee5a4baa9e63e90fe Mon Sep 17 00:00:00 2001 From: Anchel135 Date: Mon, 7 Sep 2026 11:35:08 +0300 Subject: [PATCH 14/15] fix(sqlserver): keep the prefix intact when refreshing an underscored database `load` names the graph `f"{prefix}_{db_name}"`, so recovering the prefix by splitting the graph id on `_` picks the wrong boundary whenever the database name contains one: `user1_my_db` yielded prefix `user1_my`, and the refresh reloaded into `user1_my_my_db` after deleting the graph the user was looking at. Strip the exact `_{db_name}` suffix instead, parsed before the delete so a malformed URL cannot drop a graph it will not reload. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- api/loaders/sqlserver_loader.py | 17 ++++++++--------- tests/test_sqlserver_loader.py | 34 +++++++++++++++++++++++++++++++++ 2 files changed, 42 insertions(+), 9 deletions(-) diff --git a/api/loaders/sqlserver_loader.py b/api/loaders/sqlserver_loader.py index c581a063..51ec060a 100644 --- a/api/loaders/sqlserver_loader.py +++ b/api/loaders/sqlserver_loader.py @@ -761,20 +761,19 @@ async def refresh_graph_schema(graph_id: str, db_url: str, db=None) -> Tuple[boo from api.core.db_resolver import resolve_db # pylint: disable=import-outside-toplevel + # ``load`` names the graph f"{prefix}_{db_name}", so strip that exact + # suffix — splitting on "_" mistakes a database name that contains one + # for the prefix boundary. Parsed before the delete so a bad URL cannot + # drop the graph without reloading it. + db_name = SQLServerLoader._parse_sqlserver_url(db_url)['database'] + suffix = f"_{db_name}" + prefix = graph_id[:-len(suffix)] if graph_id.endswith(suffix) else graph_id + # 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): diff --git a/tests/test_sqlserver_loader.py b/tests/test_sqlserver_loader.py index f2a96f34..a6d81796 100644 --- a/tests/test_sqlserver_loader.py +++ b/tests/test_sqlserver_loader.py @@ -812,6 +812,40 @@ async def _load(prefix, _url, db=None): # pylint: disable=unused-argument assert seen["prefix"] == "testdb" + @pytest.mark.asyncio + async def test_an_underscored_database_name_keeps_the_whole_prefix(self): + """``load`` names the graph ``f"{prefix}_{db_name}"``. + + Regression test: splitting the graph id on ``_`` cut ``user1_my_db`` down + to prefix ``user1_my``, so the reload wrote to ``user1_my_my_db`` after + deleting the graph the user was actually looking at. + """ + _graph, db = self._graph_and_db() + seen = {} + + async def _load(prefix, _url, db=None): # pylint: disable=unused-argument + seen["prefix"] = prefix + yield True, "reloaded" + + with patch("api.core.db_resolver.resolve_db", return_value=db), \ + patch.object(SQLServerLoader, "load", _load): + await SQLServerLoader.refresh_graph_schema( + "user1_my_db", "sqlserver://sa:pw@localhost/my_db", db=db) + + assert seen["prefix"] == "user1" + + @pytest.mark.asyncio + async def test_a_malformed_url_leaves_the_graph_alone(self): + """The database name is needed before the delete, not after it.""" + graph, db = self._graph_and_db() + + with patch("api.core.db_resolver.resolve_db", return_value=db): + ok, _message = await SQLServerLoader.refresh_graph_schema( + "user1_testdb", "postgresql://sa:pw@localhost/testdb", db=db) + + assert ok is False + graph.delete.assert_not_awaited() + @pytest.mark.asyncio async def test_an_unreachable_graph_is_reported_not_raised(self): with patch("api.core.db_resolver.resolve_db", side_effect=ConnectionError("down")): From 041b9bf0ffa350730b8aa6e785d08433211b2565 Mon Sep 17 00:00:00 2001 From: Anchel135 Date: Tue, 8 Sep 2026 11:08:26 +0300 Subject: [PATCH 15/15] test(sqlserver): introspect a real server, and drop the stray .coverage The fake-cursor suite cannot tell you whether a statement parses or what the driver decodes a column into, which is where every blocker in review lived. So build a deliberately awkward schema on a real server -- non-ASCII names, a `]` and a `.` in a table name, xml/text/geography columns, a uniqueidentifier, a composite key and a cross-schema foreign key -- and assert on what `_introspect_schema` returns. Skipped unless SQLSERVER_TEST_URL is set. It paid for itself immediately: the base wrapper keeps a sample only when it is already a str/int/float, so pymssql's `uuid.UUID` was discarded and every GUID column described itself with no examples. Samples are now serialized before that filter, which also recovers datetime, Decimal and bytes columns. `.coverage` was force-added past .gitignore in 050db37; removed. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .coverage | Bin 53248 -> 0 bytes api/loaders/sqlserver_loader.py | 12 +- docs/sqlserver_loader.md | 18 ++ tests/test_sqlserver_integration.py | 265 ++++++++++++++++++++++++++++ tests/test_sqlserver_loader.py | 14 ++ 5 files changed, 308 insertions(+), 1 deletion(-) delete mode 100644 .coverage create mode 100644 tests/test_sqlserver_integration.py diff --git a/.coverage b/.coverage deleted file mode 100644 index ca50e52e7e5081a30ed4325956ddc520eca41fd7..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 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& str: + """Bracket-quote for the fixture DDL, the same way the loader does.""" + return f"[{identifier.replace(']', ']]')}]" + + +@pytest.fixture(name="schema", scope="module") +def _schema(): + """A throwaway schema, dropped afterwards so a rerun starts clean.""" + url = os.getenv("SQLSERVER_TEST_URL") + if not url: + pytest.skip("SQLSERVER_TEST_URL is not set") + + params = SQLServerLoader._parse_sqlserver_url(url) + try: + conn = pymssql.connect(**params) + except pymssql.Error as exc: + pytest.skip(f"SQL Server unreachable: {exc}") + + name = f"qw_{uuid.uuid4().hex[:12]}" + other = f"{name}_ext" + try: + with conn.cursor() as cur: + _build_fixture(cur, name, other) + conn.commit() + yield name + finally: + with conn.cursor() as cur: + for target in (name, other): + _drop_schema(cur, target) + conn.commit() + conn.close() + + +def _build_fixture(cur, schema: str, other: str) -> None: + """Create the awkward schema the assertions below rely on.""" + cur.execute(f"CREATE SCHEMA {_q(schema)}") + cur.execute(f"CREATE SCHEMA {_q(other)}") + + # Parent with a non-ASCII name, a non-ASCII column and a GUID. + cur.execute(f""" + CREATE TABLE {_q(schema)}.{_q(GERMAN)} ( + id INT NOT NULL PRIMARY KEY, + {_q("Straße")} NVARCHAR(50) NULL, + guid UNIQUEIDENTIFIER NULL + ) + """) + cur.execute(f""" + INSERT INTO {_q(schema)}.{_q(GERMAN)} (id, {_q("Straße")}, guid) + VALUES (1, N'Hauptstraße', '3F2504E0-4F89-11D3-9A0C-0305E82C3301') + """) + + # Child, also non-ASCII, referencing the parent. + cur.execute(f""" + CREATE TABLE {_q(schema)}.{_q(HEBREW)} ( + id INT NOT NULL PRIMARY KEY, + kunde_id INT NOT NULL + CONSTRAINT fk_hebrew_kunde REFERENCES {_q(schema)}.{_q(GERMAN)} (id) + ) + """) + + # Types SQL Server refuses to compare, alongside one that samples normally. + cur.execute(f""" + CREATE TABLE {_q(schema)}.{_q(CJK)} ( + id INT NOT NULL PRIMARY KEY, + label NVARCHAR(50) NULL, + payload XML NULL, + notes TEXT NULL, + spot GEOGRAPHY NULL + ) + """) + cur.execute(f""" + INSERT INTO {_q(schema)}.{_q(CJK)} (id, label, payload, notes, spot) + VALUES (1, N'顧客A', '', 'note', + geography::Point(47.6, -122.3, 4326)) + """) + + # Punctuation that has to survive quoting rather than validation. + cur.execute(f"CREATE TABLE {_q(schema)}.{_q(BRACKETED)} (id INT NOT NULL PRIMARY KEY)") + cur.execute(f"INSERT INTO {_q(schema)}.{_q(BRACKETED)} (id) VALUES (1)") + cur.execute(f"CREATE TABLE {_q(schema)}.{_q(DOTTED)} (id INT NOT NULL PRIMARY KEY)") + cur.execute(f"INSERT INTO {_q(schema)}.{_q(DOTTED)} (id) VALUES (1)") + + # Composite key: one constraint, two columns. + cur.execute(f""" + CREATE TABLE {_q(schema)}.[Region] ( + country CHAR(2) NOT NULL, + code INT NOT NULL, + CONSTRAINT pk_region PRIMARY KEY (country, code) + ) + """) + cur.execute(f""" + CREATE TABLE {_q(schema)}.[Store] ( + id INT NOT NULL PRIMARY KEY, + country CHAR(2) NOT NULL, + code INT NOT NULL, + CONSTRAINT fk_store_region FOREIGN KEY (country, code) + REFERENCES {_q(schema)}.[Region] (country, code) + ) + """) + + # Cross-schema reference: in range on the parent side, out of range on the + # referenced side, so the loader has to drop it rather than emit a + # relationship pointing at a table it never loaded. + cur.execute(f"CREATE TABLE {_q(other)}.[Outside] (id INT NOT NULL PRIMARY KEY)") + cur.execute(f""" + CREATE TABLE {_q(schema)}.[CrossRef] ( + id INT NOT NULL PRIMARY KEY, + outside_id INT NOT NULL + CONSTRAINT fk_crossref_outside REFERENCES {_q(other)}.[Outside] (id) + ) + """) + + +def _drop_schema(cur, schema: str) -> None: + """Drop every table in *schema*, then the schema itself.""" + cur.execute(""" + SELECT t.name FROM sys.tables t + JOIN sys.schemas s ON t.schema_id = s.schema_id + WHERE s.name = %s + """, (schema,)) + tables = [row[0] for row in cur.fetchall()] + # Foreign keys first; a table cannot be dropped while one points at it. + for table in tables: + cur.execute(""" + SELECT fk.name + FROM sys.foreign_keys fk + JOIN sys.tables t ON fk.parent_object_id = t.object_id + JOIN sys.schemas s ON t.schema_id = s.schema_id + WHERE s.name = %s AND t.name = %s + """, (schema, table)) + for (constraint,) in cur.fetchall(): + cur.execute( + f"ALTER TABLE {_q(schema)}.{_q(table)} DROP CONSTRAINT {_q(constraint)}" + ) + for table in tables: + cur.execute(f"DROP TABLE {_q(schema)}.{_q(table)}") + cur.execute(f"DROP SCHEMA {_q(schema)}") + + +@pytest.fixture(name="introspection", scope="module") +def _introspection(schema): + """One introspection pass, shared by every assertion below.""" + params = SQLServerLoader._parse_sqlserver_url(os.getenv("SQLSERVER_TEST_URL")) + return SQLServerLoader._introspect_schema(params, schema) + + +class TestIntrospectionAgainstRealServer: + """What the loader gets back from a server that actually parses the SQL.""" + + def test_every_table_loads(self, introspection): + """No statement in the walk is rejected, whatever the names look like. + + Regression test for error 145: the sample query paired ``DISTINCT`` with + ``ORDER BY NEWID()``, which SQL Server refuses, so the first column of + the first table took the whole load down. + """ + entities, _relationships = introspection + assert set(entities) == { + GERMAN, HEBREW, CJK, BRACKETED, DOTTED, "Region", "Store", "CrossRef", + } + + def test_non_comparable_types_cost_only_their_own_column(self, introspection): + """``xml``, ``text`` and ``geography`` cannot be DISTINCT-ed.""" + entities, _relationships = introspection + columns = entities[CJK]['columns'] + for col in ("payload", "notes", "spot"): + assert columns[col]['sample_values'] == [] + assert columns['label']['sample_values'] == ["顧客A"] + + def test_a_dotted_table_name_loads_without_samples(self, introspection): + """A dot makes the qualified name ambiguous, so sampling is skipped.""" + entities, _relationships = introspection + assert entities[DOTTED]['columns']['id']['sample_values'] == [] + + def test_a_bracketed_table_name_still_samples(self, introspection): + """``]`` is doubled by the quoter, so it needs no special handling.""" + entities, _relationships = introspection + assert entities[BRACKETED]['columns']['id']['sample_values'] == ["1"] + + def test_non_ascii_names_and_values_round_trip(self, introspection): + """The old allow-list rejected all three of these.""" + entities, _relationships = introspection + assert entities[GERMAN]['columns']["Straße"]['sample_values'] == ["Hauptstraße"] + + def test_guids_arrive_as_json_encodable_strings(self, introspection): + """pymssql decodes ``uniqueidentifier`` to ``uuid.UUID``.""" + entities, _relationships = introspection + samples = entities[GERMAN]['columns']['guid']['sample_values'] + assert samples == ["3f2504e0-4f89-11d3-9a0c-0305e82c3301"] + json.dumps(entities[GERMAN]['columns']['guid']) + + def test_a_composite_key_is_one_relationship_with_two_columns(self, introspection): + """Both column pairs land under the single constraint name.""" + _entities, relationships = introspection + assert sorted( + (edge['source_column'], edge['target_column']) + for edge in relationships['fk_store_region'] + ) == [("code", "code"), ("country", "country")] + + def test_a_cross_schema_key_is_left_out_of_both_structures(self, introspection): + """Otherwise the graph gains an edge to a table that was never loaded.""" + entities, relationships = introspection + assert "fk_crossref_outside" not in relationships + assert entities["CrossRef"]['foreign_keys'] == [] + + def test_key_kinds_come_back_in_the_loader_vocabulary(self, introspection): + """The CASE emits these directly; nothing translates MySQL codes.""" + entities, _relationships = introspection + assert entities[GERMAN]['columns']['id']['key'] == "PRIMARY KEY" + assert entities[HEBREW]['columns']['kunde_id']['key'] == "FOREIGN KEY" + assert entities[CJK]['columns']['label']['key'] == "NONE" diff --git a/tests/test_sqlserver_loader.py b/tests/test_sqlserver_loader.py index a6d81796..d4787013 100644 --- a/tests/test_sqlserver_loader.py +++ b/tests/test_sqlserver_loader.py @@ -367,6 +367,20 @@ def test_extract_sample_values_stringifies(self): assert SQLServerLoader.extract_sample_values_for_column( cursor, "dbo.T", "n") == ["1", "2"] + def test_driver_objects_are_serialized_before_the_base_filter(self): + """The base wrapper keeps a sample only if it is a str/int/float. + + Regression test: pymssql hands back ``uuid.UUID`` for + ``uniqueidentifier``, so every GUID column described itself with no + sample values at all. Caught by the integration suite, not by a fake + cursor that had only ever been fed strings. + """ + cursor = FakeCursor([[ + {"v": uuid.UUID("3f2504e0-4f89-11d3-9a0c-0305e82c3301")}, + ]]) + assert SQLServerLoader.extract_sample_values_for_column( + cursor, "dbo.T", "v") == ["3f2504e0-4f89-11d3-9a0c-0305e82c3301"] + @pytest.mark.parametrize("error", [ pymssql.Error("DISTINCT is not defined for xml"), ValueError("Invalid table name"),