diff --git a/.github/wordlist.txt b/.github/wordlist.txt index 393bfb34..9840414c 100644 --- a/.github/wordlist.txt +++ b/.github/wordlist.txt @@ -124,6 +124,14 @@ SDK Dependabot PyPI pypi +pymssql +sqlserver +SQLServerLoader +dbo +tsql +hostname +TLS +sqlglot signup SMTP outbox diff --git a/README.md b/README.md index 3e7e7e46..04974a1c 100644 --- a/README.md +++ b/README.md @@ -335,6 +335,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 @@ -345,7 +349,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" @@ -389,7 +393,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 (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 | @@ -435,7 +439,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 or MySQL) +- Target SQL database (PostgreSQL, MySQL, SQL Server or Snowflake — the last two + require the `queryweaver[server]` extra) ## Development diff --git a/api/core/pipeline.py b/api/core/pipeline.py index 2aca3998..f7731826 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,7 @@ def truncate_for_log(query: str, max_length: int = 200) -> str: "postgres": "postgres", "mysql": "mysql", "snowflake": "snowflake", + "sqlserver": "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..81c71c6a --- /dev/null +++ b/api/loaders/sqlserver_loader.py @@ -0,0 +1,892 @@ +"""SQL Server loader for loading database schemas into FalkorDB graphs.""" + +import datetime +import decimal +import logging +import re +import uuid +from typing import AsyncGenerator, Dict, Any, List, Tuple +from urllib.parse import urlparse, parse_qs, unquote + +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 + +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.""" + + +class SQLServerConnectionError(Exception): + """Exception raised for SQL Server connection errors.""" + + +def validate_ident( + identifier: str, identifier_type: str = "identifier", allow_dot: bool = True +) -> str: + """Validate that a *user-supplied* identifier is safe to interpolate into T-SQL. + + 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 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 + 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. + + 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." + ) + 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, " + f"underscore, dollar, hash, at-sign, space, {dot}and dash are allowed." + ) + 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 ``]``. + + 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. + + This is defence in depth: callers that interpolate the result into a + statement validate the identifier with :func:`validate_ident` or + :func:`validate_catalog_ident` first. + + Args: + identifier: Raw identifier as read from the system catalog. + + Returns: + The bracket-quoted identifier. + """ + return f"[{identifier.replace(']', ']]')}]" + + +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. + """ + + # 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. 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_catalog_ident(bare_table, "table name", allow_dot=False) + ) + if schema: + qualified = ( + f"{quote_ident(validate_catalog_ident(schema, 'schema name', allow_dot=False))}" + f".{qualified}" + ) + + 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 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 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) + + # The cursor is opened with ``as_dict=True`` so rows are keyed by column + # name only — pymssql's ``row2dict`` strips positional keys. + # + # Serialized here because the base wrapper keeps a sample only when it + # is already a str/int/float, so the driver's own objects — uuid.UUID + # for ``uniqueidentifier``, datetime, Decimal, bytes — would otherwise + # be dropped and the column would silently describe itself with no + # examples at all. + sample_results = cursor.fetchall() + return [ + SQLServerLoader._serialize_value(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): + """ + 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, 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() + 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. + + Raises: + ValueError: If the requested schema is not a valid identifier. + """ + try: + parsed = urlparse(connection_url) + # 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: + return DEFAULT_SCHEMA + return validate_ident(schema, "schema name") + + @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 + 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': max(Config.DB_SCHEMA_TIMEOUT, Config.DB_STATEMENT_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 + 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( # pylint: disable=no-member + **SQLServerLoader._with_timeouts(conn_params) + ) + cursor = conn.cursor(as_dict=True) + + 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) + + @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 + """ + try: + # 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) + db_name = conn_params['database'] + + # 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, relationships = await run_introspection( + SQLServerLoader._introspect_schema, conn_params, schema + ) + + # 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" + + @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, + 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 + # 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 + 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'] + catalog_schema = table_info['schema_name'] + + # Get column information for this table + columns_info = SQLServerLoader.extract_columns_info( + cursor, schema, table_name, catalog_schema + ) + + # 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_by_table.get(table_name, []), + 'col_descriptions': col_descriptions + } + + return entities + + @staticmethod + def extract_columns_info( + 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. 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 + that reaches a statement body comes from the server rather + than from the connection URL. + + 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 '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 + 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 ( + 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 + ) 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 = {} + + qualified_table = f"{catalog_schema}.{table_name}" + + for col_info in columns: + col_name = col_info['column_name'] + is_nullable = 'YES' if col_info['is_nullable'] else 'NO' + key_type = col_info['column_key'] + + columns_info[col_name] = { + 'type': col_info['data_type'], + 'null': is_nullable, + 'key': key_type, + '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 + + @staticmethod + def extract_foreign_keys( + cursor, schema: str = DEFAULT_SCHEMA + ) -> List[Dict[str, str]]: + """ + Extract every foreign key in a schema, in a single query. + + 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 foreign keys from (default: ``dbo``) + + Returns: + One dict per foreign-key column, in constraint order. + """ + 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)) + + 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()] + + @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 + + @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 + 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 + + # ``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() + + # 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( # pylint: disable=no-member + **SQLServerLoader._with_timeouts(conn_params) + ) + 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..b05f35d8 100644 --- a/api/sql_utils/sql_sanitizer.py +++ b/api/sql_utils/sql_sanitizer.py @@ -24,20 +24,47 @@ 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 *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 and internally escaped. + """ + 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) -> 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,22 +78,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 - 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]: @@ -130,7 +163,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 +200,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() == 'sqlserver': + 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..ce24ab74 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(""); @@ -71,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) { @@ -100,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", @@ -134,17 +169,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}:${effectivePort}/${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 +340,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 +426,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 +594,7 @@ const DatabaseModal = ({ open, onOpenChange }: DatabaseModalProps) => { setPort(e.target.value)} className="bg-muted border-border focus-visible:ring-purple-500" @@ -592,8 +635,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..b6a68b75 --- /dev/null +++ b/docs/sqlserver_loader.md @@ -0,0 +1,188 @@ +# 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 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 — both sides must live in the selected schema + +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. The `DISTINCT` sits in a derived table because SQL Server rejects +`SELECT DISTINCT … ORDER BY NEWID()`: + +```sql +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 +- 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`. + +## 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` 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. + +### 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: + +- 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 +``` + +`tests/test_sqlserver_integration.py` runs the same introspection against a real +server, because a fake cursor cannot tell you whether the SQL parses or what the +driver decodes a column into. It builds a deliberately awkward schema — non-ASCII +table and column names, a `]` and a `.` in a table name, `xml`/`text`/`geography` +columns, a `uniqueidentifier`, a composite key and a cross-schema foreign key — +and asserts the tables load, the non-comparable types degrade to no samples, and +the cross-schema key is excluded from both entities and relationships. + +It skips unless `SQLSERVER_TEST_URL` is set: + +```bash +docker run -e ACCEPT_EULA=Y -e MSSQL_SA_PASSWORD='Str0ng!Passw0rd' \ + -p 1433:1433 -d mcr.microsoft.com/mssql/server:2022-latest + +SQLSERVER_TEST_URL='sqlserver://sa:Str0ng!Passw0rd@localhost:1433/master' \ + uv run --extra server --extra dev pytest tests/test_sqlserver_integration.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 +- `DB_STATEMENT_TIMEOUT` is not applied on its own; see [Timeouts](#timeouts) diff --git a/pyproject.toml b/pyproject.toml index ff0317ff..bf138ea3 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..f85c13dc 100644 --- a/tests/test_destructive_detection.py +++ b/tests/test_destructive_detection.py @@ -302,3 +302,34 @@ 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 diff --git a/tests/test_schema_load_offloading.py b/tests/test_schema_load_offloading.py index 0e056cb3..ebba8c25 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_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): + 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 8937c873..03d81a3b 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.""" @@ -19,9 +23,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 +51,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 +241,87 @@ 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') == '[' + + 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 + + @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]' diff --git a/tests/test_sqlserver_integration.py b/tests/test_sqlserver_integration.py new file mode 100644 index 00000000..01346671 --- /dev/null +++ b/tests/test_sqlserver_integration.py @@ -0,0 +1,265 @@ +"""Introspection run against a real SQL Server. + +The unit tests drive a fake cursor, so they pin what the loader *does* with a +result set but never send a statement to a server. Every blocker found in review +lived in the part they cannot reach: ``SELECT DISTINCT ... ORDER BY NEWID()`` is +rejected by the parser, ``uniqueidentifier`` only becomes a ``uuid.UUID`` once a +driver decodes one, and an identifier the allow-list refused only fails when a +real catalog hands it back. + +So this file builds a deliberately awkward schema -- non-ASCII names, a bracket +and a dot in a table name, types SQL Server will not compare, a GUID column, a +composite key and a cross-schema reference -- and asserts on what +``_introspect_schema`` returns. It is one pass over one fixture, not a second +copy of the unit tests: what is being checked is that the SQL is valid and that +the awkward cases degrade the way they are supposed to. + +Set ``SQLSERVER_TEST_URL`` to run it, e.g. against + + docker run -e ACCEPT_EULA=Y -e MSSQL_SA_PASSWORD=Str0ng!Passw0rd \ + -p 1433:1433 -d mcr.microsoft.com/mssql/server:2022-latest + + SQLSERVER_TEST_URL='sqlserver://sa:Str0ng!Passw0rd@localhost:1433/master' + +Without it, the module skips. +""" +# pylint: disable=protected-access + +import json +import os +import uuid +import importlib + +import pytest + +pymssql = pytest.importorskip("pymssql") + +# See ``tests/test_sqlserver_loader.py``: ``api.core.__init__`` eagerly pulls in +# the pipeline, which imports the loaders, so importing a loader first leaves +# ``graph_loader`` half-built. +importlib.import_module("api.core") + +from api.loaders.sqlserver_loader import SQLServerLoader # noqa: E402 pylint: disable=wrong-import-position + +pytestmark = [pytest.mark.integration] + +# Names the old ASCII allow-list rejected outright, plus the two punctuation +# cases that have to survive quoting rather than validation. +GERMAN = "Kunden_Ä" +HEBREW = "לקוחות" +CJK = "顧客" +BRACKETED = "my]table" +DOTTED = "weird.name" + + +def _q(identifier: str) -> 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 new file mode 100644 index 00000000..d4787013 --- /dev/null +++ b/tests/test_sqlserver_loader.py @@ -0,0 +1,966 @@ +"""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 +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. +# ``api.core.__init__`` eagerly pulls in the pipeline, which imports the +# loaders, so importing a loader first leaves ``graph_loader`` half-built. +# 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.config import Config # noqa: E402 pylint: disable=wrong-import-position +from api.core.errors import InvalidArgumentError # noqa: E402 pylint: disable=wrong-import-position +from api.core.pipeline import get_database_type_and_loader # noqa: E402 pylint: disable=wrong-import-position +from api.loaders.sqlserver_loader import ( # noqa: E402 pylint: disable=wrong-import-position + SQLServerLoader, + SQLServerQueryError, + quote_ident, + validate_catalog_ident, + validate_ident, +) + +pytestmark = pytest.mark.unit + + +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 TestValidateIdent: + """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", + ]) + 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") + + 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 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 contains hostile identifiers rather than trusting them.""" + + @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_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([[]]) + 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): + """``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 == [] + + @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.""" + + 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" + + 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.""" + + 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 + + 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): + """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"] + + 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"), + ]) + 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.""" + + 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", "Sales") + 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": "PRIMARY KEY", + "column_comment": "", + }], + [{"id": 1}], # sample values query + ]) + 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" + 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": "NONE", + "column_comment": "", + }], + [], + ]) + 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 + + 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") + 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] + 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.""" + rels = SQLServerLoader.build_relationships([ + { + "table": "Orders", + "constraint_name": "FK_A", + "column": "c1", + "referenced_table": "Customers", + "referenced_column": "id1", + }, + { + "table": "Orders", + "constraint_name": "FK_A", + "column": "c2", + "referenced_table": "Customers", + "referenced_column": "id2", + }, + ]) + 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_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.""" + cursor = FakeCursor([ + [{"table_name": "Orders", "schema_name": "dbo", "table_comment": "All orders"}], + [{ + "column_name": "id", + "data_type": "int", + "is_nullable": False, + "column_default": None, + "column_key": "PRIMARY KEY", + "column_comment": "", + }], + [{"id": 7}], + ]) + 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"] == [] + + 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": "PRIMARY KEY", + "column_comment": "", + }], + [{"id": 7}], + ]) + 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.""" + + @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"), + # 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"), + ]) + 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([ + [], # foreign keys + [{"table_name": "Orders", "schema_name": "dbo", "table_comment": ""}], + [], # columns + ]) + 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", "sales") + assert cursor.executed[1][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")] + + @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"] == max(Config.DB_SCHEMA_TIMEOUT, Config.DB_STATEMENT_TIMEOUT) + + def test_execute_query_bounds_connect_and_query_time(self): + """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: + 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"] == max(Config.DB_SCHEMA_TIMEOUT, Config.DB_STATEMENT_TIMEOUT) + + +class TestLoaderRouting: + """`sqlserver://` URLs must resolve to this loader, and only with the extra.""" + + def test_server_install_gets_the_loader(self): + db_type, loader = get_database_type_and_loader( + "sqlserver://sa:pw@localhost/testdb") + assert (db_type, loader) == ("sqlserver", SQLServerLoader) + + def test_sdk_only_install_is_told_which_extra_to_add(self): + # A clean error beats a deferred ImportError on pymssql. + with pytest.raises(InvalidArgumentError, match=r"\[server\] extra"): + get_database_type_and_loader( + "sqlserver://sa:pw@localhost/testdb", sdk_only=True) + + +class TestRefreshGraphSchema: + """Reloading the graph after a DDL statement.""" + + @staticmethod + def _graph_and_db(): + graph = MagicMock() + graph.delete = AsyncMock() + db = MagicMock() + db.select_graph.return_value = graph + return graph, db + + @pytest.mark.asyncio + async def test_reload_drops_the_graph_and_reports_success(self): + graph, db = self._graph_and_db() + + async def _load(prefix, _url, db=None): # pylint: disable=unused-argument + assert prefix == "user1" + yield True, "reloaded" + + with patch("api.core.db_resolver.resolve_db", return_value=db), \ + patch.object(SQLServerLoader, "load", _load): + ok, message = await SQLServerLoader.refresh_graph_schema( + "user1_testdb", "sqlserver://sa:pw@localhost/testdb", db=db) + + graph.delete.assert_awaited_once() + assert (ok, message) == (True, "reloaded") + + @pytest.mark.asyncio + async def test_a_failed_reload_is_reported(self): + _graph, db = self._graph_and_db() + + async def _load(_prefix, _url, db=None): # pylint: disable=unused-argument + yield False, "nope" + + with patch("api.core.db_resolver.resolve_db", return_value=db), \ + patch.object(SQLServerLoader, "load", _load): + ok, message = await SQLServerLoader.refresh_graph_schema( + "user1_testdb", "sqlserver://sa:pw@localhost/testdb", db=db) + + assert ok is False + assert message == "Failed to reload schema" + + @pytest.mark.asyncio + async def test_a_graph_id_without_a_prefix_is_used_whole(self): + _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( + "testdb", "sqlserver://sa:pw@localhost/testdb", db=db) + + 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")): + ok, message = await SQLServerLoader.refresh_graph_schema( + "user1_testdb", "sqlserver://sa:pw@localhost/testdb") + + assert (ok, message) == (False, "Error refreshing graph schema") + + +class TestConnectionCleanupIsQuiet: + """Cleanup runs on the error path, so it must not raise a second failure.""" + + def test_close_swallows_driver_errors(self): + cursor = MagicMock() + cursor.close.side_effect = RuntimeError("already gone") + conn = MagicMock() + conn.close.side_effect = RuntimeError("already gone") + + SQLServerLoader._close_quietly(cursor, conn) # must not raise + + def test_rollback_swallows_driver_errors(self): + conn = MagicMock() + conn.rollback.side_effect = RuntimeError("no transaction") + + SQLServerLoader._rollback_quietly(conn) # must not raise + + +class TestUrlEdgeCases: + """Malformed and opt-out variants of the connection URL.""" + + @pytest.mark.parametrize("url,message", [ + ("sqlserver://sa:pw@/testdb", "must include a host"), + ("sqlserver://sa:pw@localhost/", "must include database name"), + ("sqlserver://localhost/testdb", "must include username"), + ("mysql://sa:pw@localhost/testdb", "Invalid SQL Server URL format"), + ]) + def test_malformed_urls_are_rejected(self, url, message): + with pytest.raises(ValueError, match=message): + SQLServerLoader._parse_sqlserver_url(url) + + @pytest.mark.parametrize("value,expected", [ + ("true", "require"), ("1", "require"), ("yes", "require"), ("require", "require"), + ("false", "off"), ("0", "off"), ("no", "off"), ("off", "off"), + ]) + def test_encryption_is_opt_in_both_ways(self, value, expected): + params = SQLServerLoader._parse_sqlserver_url( + f"sqlserver://sa:pw@localhost/testdb?encrypt={value}") + assert params["encryption"] == expected + + def test_driver_defaults_are_left_alone_when_unspecified(self): + params = SQLServerLoader._parse_sqlserver_url("sqlserver://sa:pw@localhost/testdb") + assert "encryption" not in params + + @pytest.mark.parametrize("url", [None, 12345]) + def test_an_unparseable_url_falls_back_to_the_default_schema(self, url): + # urlparse raises AttributeError/ValueError rather than returning empty. + assert SQLServerLoader.parse_schema_from_url(url) == "dbo" + + +class TestDdlResultShape: + """DDL reports no row count, because there is none to report.""" + + def test_ddl_reports_the_operation_only(self): + cursor = FakeCursor([[]]) + cursor.description = None + conn = FakeConnection(cursor) + + with patch("pymssql.connect", return_value=conn): + rows = SQLServerLoader.execute_sql_query( + "CREATE TABLE t (id INT)", "sqlserver://sa:pw@localhost/testdb") + + assert rows == [{"operation": "CREATE", "status": "success"}] + assert conn.committed + + +class TestDriverErrorsDuringLoad: + """A pymssql failure is reported as a load failure, not raised.""" + + @pytest.mark.asyncio + async def test_driver_error_is_reported(self): + import pymssql # pylint: disable=import-outside-toplevel + + results = [] + with patch("pymssql.connect", side_effect=pymssql.OperationalError("refused")): + async for success, message in SQLServerLoader.load( + "user1", "sqlserver://sa:pw@localhost/testdb"): + results.append((success, message)) + + assert results[-1][0] is False + + @pytest.mark.asyncio + async def test_driver_error_during_a_query_is_wrapped(self): + import pymssql # pylint: disable=import-outside-toplevel + + cursor = FakeCursor() + cursor.execute = MagicMock(side_effect=pymssql.OperationalError("deadlock")) + 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 diff --git a/uv.lock b/uv.lock index 1b27b6f6..d693b973 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" },