diff --git a/docs/changelog.rst b/docs/changelog.rst index 3b6454409..1b40cca05 100644 --- a/docs/changelog.rst +++ b/docs/changelog.rst @@ -20,14 +20,18 @@ v0.56.0 **Added:** -* New sync and async schema checks can create missing tables and add columns. - Use ``ensure_schema_sync()`` or ``ensure_schema_async()`` for each driver mode. +* New ``SchemaTarget`` and ``SchemaEnsureResult`` types plus sync and async + schema checks can create missing tables and add columns. Use + ``ensure_schema_sync()`` or ``ensure_schema_async()`` for each driver mode. + ADK, Litestar session, and durable event stores expose ``manage_schema``, + ``create_schema``, and ``run_migrations`` controls for this lifecycle. * Oracle ADK, durable event, and Litestar session tables now share opt-in compression, partitioning, In-Memory, and table-option configuration. -* BigQuery session and queue partition options, CockroachDB session hash - sharding and row-level TTL, PostgreSQL table/autovacuum tuning, and opt-in - SQLite extension PRAGMA profiles are now available across the applicable - Litestar, Events, and ADK stores. +* BigQuery session and queue partition options and CockroachDB session hash + sharding and row-level TTL are now available. +* Applicable Litestar, Events, and ADK stores now expose PostgreSQL table and + autovacuum tuning, MySQL and MariaDB table/index options, Spanner sharding and + table/index options, and opt-in SQLite extension PRAGMA profiles. **Changed:** @@ -36,14 +40,20 @@ v0.56.0 * ADK, Litestar session, and durable event stores now derive additive schema currency from their canonical DDL. ADK no longer seeds or bumps a ``schema_version`` row for additive changes. -* Oracle extension-table optimizations are capability-gated through the - pool-scoped data dictionary and degrade to structured warnings when an - option is unavailable. +* Oracle server-version, JSON storage, and extension-table capability detection + now share the config/pool-scoped data dictionary cache. Requested storage + optimizations degrade to structured warnings when an option is unavailable. **Fixed:** +* PostgreSQL-family ADBC connections now bind top-level UUID parameters as + PostgreSQL ``uuid`` values across ordinary, batch, streaming, and Arrow + execution routes. (`#650 `_) * Psycopg sync and async transactions now restore the connection's original autocommit mode after SQLSpec-owned commit or rollback operations (`#648`_). +* Spanner data-dictionary queries now cast nullable metadata filters to their + concrete ``STRING`` or ``TIMESTAMP`` types, avoiding conflicting parameter + inference when optional filters are omitted. * Builder caching now reuses value-independent expression templates and binds each call's current parameters and statement configuration. This also isolates CTE bodies and returned ASTs instead of sharing mutable cached diff --git a/sqlspec/adapters/adbc/core.py b/sqlspec/adapters/adbc/core.py index 316c3d37b..e62b9c9fa 100644 --- a/sqlspec/adapters/adbc/core.py +++ b/sqlspec/adapters/adbc/core.py @@ -3,7 +3,13 @@ import datetime import decimal from collections.abc import Sized +from functools import lru_cache from typing import TYPE_CHECKING, Any, Final, cast +from uuid import UUID + +import sqlglot +from sqlglot import exp +from sqlglot.errors import ParseError from sqlspec.adapters.adbc.type_converter import get_adbc_type_converter from sqlspec.core import ( @@ -68,6 +74,7 @@ "normalize_script_rowcount", "prepare_parameters_with_casts", "prepare_postgres_parameters", + "prepare_postgres_uuid_bindings", "resolve_column_names", "resolve_dialect_from_config", "resolve_dialect_from_driver_path", @@ -84,6 +91,7 @@ ) COLUMN_CACHE_MAX_SIZE: int = 256 +_UUID_TYPES: Final[tuple[type[Any], ...]] = tuple(build_uuid_coercions()) DIALECT_PATTERNS: "dict[str, tuple[str, ...]]" = { "postgres": ("postgres", "postgresql"), @@ -847,6 +855,177 @@ def prepare_parameters_with_casts( return parameters +def prepare_postgres_uuid_bindings( + compiled_sql: str, prepared_parameters: Any, *, is_many: bool, dialect: str +) -> "tuple[str, object]": + """Bind top-level UUID values through PostgreSQL-family ADBC drivers. + + Args: + compiled_sql: SQL after normal statement compilation and cache rebinding. + prepared_parameters: Current execution's driver parameters. + is_many: Whether parameters contain an execute-many batch. + dialect: Resolved ADBC dialect name. + + Returns: + Rewritten SQL and current-call parameters with UUID values normalized, + or the inputs unchanged when no placeholder needs a new UUID cast. + """ + if not is_postgres_dialect(dialect): + return compiled_sql, prepared_parameters + if is_many: + uuid_ordinals = _detect_batch_uuid_ordinals(prepared_parameters) + elif isinstance(prepared_parameters, (list, tuple)): + uuid_ordinals = tuple( + index for index, value in enumerate(prepared_parameters, 1) if isinstance(value, _UUID_TYPES) + ) + else: + uuid_ordinals = () + if not uuid_ordinals: + return compiled_sql, prepared_parameters + sqlglot_dialect = "postgres" if dialect == "postgresql" else dialect + rewritten_sql, effective_ordinals = _rewrite_postgres_uuid_placeholders( + compiled_sql, tuple(sorted(uuid_ordinals)), sqlglot_dialect + ) + if not effective_ordinals: + return compiled_sql, prepared_parameters + if is_many: + converted_rows = [ + _convert_uuid_row(row, effective_ordinals, row_number) + for row_number, row in enumerate(prepared_parameters, 1) + ] + converted_parameters = tuple(converted_rows) if isinstance(prepared_parameters, tuple) else converted_rows + else: + converted_parameters = _convert_uuid_row(prepared_parameters, effective_ordinals, 1) + return rewritten_sql, converted_parameters + + +def _parameter_ordinal(parameter: exp.Parameter) -> "int | None": + value = parameter.this + if not isinstance(value, exp.Literal) or value.is_string: + return None + try: + return int(value.this) + except (TypeError, ValueError): + return None + + +def _direct_parameter_cast(parameter: exp.Parameter) -> "exp.Cast | None": + current: exp.Expression = parameter + parent = current.parent + while isinstance(parent, exp.Paren): + current = parent + parent = current.parent + if isinstance(parent, exp.Cast) and parent.this is current: + return parent + return None + + +@lru_cache(maxsize=256) +def _rewrite_postgres_uuid_placeholders( + sql: str, uuid_ordinals: "tuple[int, ...]", dialect: str +) -> "tuple[str, tuple[int, ...]]": + """Wrap requested UUID parameter placeholders in explicit UUID casts. + + Placeholders that already carry a different explicit cast stay authoritative + and are excluded from the returned effective ordinals. Results are memoized + per ``(sql, uuid_ordinals, dialect)`` so repeated executions of the same + statement shape reuse the rewrite without parsing again. + """ + try: + expressions = cast( + "list[exp.Expression]", + [expression for expression in sqlglot.parse(sql, read=dialect) if expression is not None], + ) + except ParseError as exc: + msg = f"Failed to parse PostgreSQL ADBC SQL for UUID parameter binding: {exc}" + raise SQLSpecError(msg) from exc + if not expressions: + msg = "Failed to parse PostgreSQL ADBC SQL for UUID parameter binding: SQLGlot returned no statements." + raise SQLSpecError(msg) + + requested = set(uuid_ordinals) + authoritative: set[int] = set() + for expression in expressions: + for parameter in expression.find_all(exp.Parameter): + ordinal = _parameter_ordinal(parameter) + if ordinal is None or ordinal not in requested: + continue + cast_expression = _direct_parameter_cast(parameter) + if cast_expression is not None: + target = cast_expression.args.get("to") + target_kind = target.args.get("kind") if isinstance(target, exp.DataType) else None + if isinstance(target_kind, exp.Dot): + target_kind = target_kind.expression + is_uuid_cast = isinstance(target, exp.DataType) and ( + target.this == exp.DataType.Type.UUID + or ( + target.this == exp.DataType.Type.USERDEFINED + and isinstance(target_kind, exp.Identifier) + and target_kind.name.lower() == "uuid" + ) + ) + if not is_uuid_cast: + authoritative.add(ordinal) + effective = tuple(ordinal for ordinal in uuid_ordinals if ordinal not in authoritative) + effective_set = set(effective) + + def wrap_parameter(node: exp.Expression) -> exp.Expression: + if not isinstance(node, exp.Parameter): + return node + ordinal = _parameter_ordinal(node) + if ordinal is None or ordinal not in effective_set: + return node + if _direct_parameter_cast(node) is not None: + return node + return exp.Cast(this=node.copy(), to=exp.DataType.build("UUID")) + + for expression in expressions: + expression.transform(wrap_parameter, copy=False) + return "; ".join(expression.sql(dialect=dialect) for expression in expressions), effective + + +def _detect_batch_uuid_ordinals(parameters: Any) -> "tuple[int, ...]": + if not isinstance(parameters, (list, tuple)) or not parameters: + return () + first_row = parameters[0] + if not isinstance(first_row, (list, tuple)): + return () + expected_size = len(first_row) + for row in parameters: + if not isinstance(row, (list, tuple)) or len(row) != expected_size: + return () + return tuple( + ordinal + for ordinal in range(1, expected_size + 1) + if any(isinstance(row[ordinal - 1], _UUID_TYPES) for row in parameters) + ) + + +def _convert_uuid_row(row: Any, ordinals: "tuple[int, ...]", row_number: int) -> Any: + converted = list(row) + for ordinal in ordinals: + value = converted[ordinal - 1] + if value is None: + continue + if not isinstance(value, _UUID_TYPES) and not isinstance(value, str): + msg = ( + f"ADBC PostgreSQL UUID parameter ordinal {ordinal} has incompatible " + f"{type(value).__name__} value in batch row {row_number}; expected a UUID object, " + "parseable UUID string, or None." + ) + raise SQLSpecError(msg) + try: + converted[ordinal - 1] = str(UUID(str(value))) + except (AttributeError, TypeError, ValueError) as exc: + msg = ( + f"ADBC PostgreSQL UUID parameter ordinal {ordinal} has incompatible " + f"{type(value).__name__} value in batch row {row_number}; expected a UUID object, " + "parseable UUID string, or None." + ) + raise SQLSpecError(msg) from exc + return tuple(converted) if isinstance(row, tuple) else converted + + def _create_adbc_error(error: Any, error_class: type[SQLSpecError], description: str) -> SQLSpecError: """Create an ADBC error instance without raising it.""" msg = f"ADBC {description}: {error}" diff --git a/sqlspec/adapters/adbc/driver.py b/sqlspec/adapters/adbc/driver.py index c53975122..374c7bee3 100644 --- a/sqlspec/adapters/adbc/driver.py +++ b/sqlspec/adapters/adbc/driver.py @@ -22,6 +22,7 @@ normalize_postgres_empty_parameters, normalize_script_rowcount, prepare_postgres_parameters, + prepare_postgres_uuid_bindings, resolve_column_names, resolve_dialect_name, resolve_many_rowcount, @@ -189,6 +190,16 @@ def __init__( # CORE DISPATCH METHODS # ───────────────────────────────────────────────────────────────────────────── + def _compiled_sql( + self, statement: "SQL", statement_config: "StatementConfig", flatten_single_parameters: bool = False + ) -> "tuple[str, object]": + compiled_sql, prepared_parameters = super()._compiled_sql( + statement, statement_config, flatten_single_parameters=flatten_single_parameters + ) + return prepare_postgres_uuid_bindings( + compiled_sql, prepared_parameters, is_many=statement.is_many, dialect=self._dialect_name + ) + def dispatch_execute(self, cursor: "AdbcRawCursor", statement: SQL) -> "ExecutionResult": """Execute single SQL statement. diff --git a/tests/integration/adapters/postgres/adbc/test_driver.py b/tests/integration/adapters/postgres/adbc/test_driver.py index dfe5b7808..a20abb09c 100644 --- a/tests/integration/adapters/postgres/adbc/test_driver.py +++ b/tests/integration/adapters/postgres/adbc/test_driver.py @@ -1,5 +1,10 @@ """PostgreSQL-backed ADBC driver residuals.""" +from uuid import uuid4 + +import pytest + +from sqlspec.adapters.adbc import AdbcDriver from tests.integration.adapters._shared.adbc_backends import postgresql_session, test_postgresql_specific_features from tests.integration.adapters._shared.adbc_connection import ( test_connection, @@ -22,3 +27,67 @@ "test_execute_script_edge_cases", "test_postgresql_specific_features", ) + + +@pytest.mark.xdist_group("postgres") +@pytest.mark.adbc +def test_postgresql_uuid_identity_and_same_sql_cache_reuse(postgresql_session: AdbcDriver) -> None: + """Distinct UUID objects bind through one cached INSERT statement without losing identity.""" + table_name = "adbc_uuid_identity" + values = [uuid4(), uuid4()] + insert_sql = f"INSERT INTO {table_name} (position, value) VALUES (?, ?)" + + try: + postgresql_session.execute_script(f"DROP TABLE IF EXISTS {table_name}") + postgresql_session.execute_script(f"CREATE TABLE {table_name} (position INTEGER PRIMARY KEY, value UUID)") + for position, value in enumerate(values, 1): + postgresql_session.execute(insert_sql, (position, value)) + + rows = postgresql_session.execute( + f"SELECT position, value::text AS value FROM {table_name} ORDER BY position" + ).get_data() + assert [row["value"] for row in rows] == [str(value) for value in values] + finally: + postgresql_session.execute_script(f"DROP TABLE IF EXISTS {table_name}") + + +@pytest.mark.xdist_group("postgres") +@pytest.mark.adbc +@pytest.mark.parametrize("uuid_first", [False, True], ids=["ordinary-first", "uuid-first"]) +def test_postgresql_uuid_binding_does_not_leak_through_same_sql_cache( + postgresql_session: AdbcDriver, uuid_first: bool +) -> None: + """Value-aware UUID SQL never replaces the stable cached statement in either value order.""" + statement = "SELECT pg_typeof($1)::text AS bound_type" + uuid_value = uuid4() + values = (uuid_value, "ordinary") if uuid_first else ("ordinary", uuid_value) + + bound_types = [postgresql_session.select_value(statement, value) for value in values] + uuid_type, ordinary_type = bound_types if uuid_first else reversed(bound_types) + + assert uuid_type == "uuid" + assert ordinary_type != uuid_type + + +@pytest.mark.xdist_group("postgres") +@pytest.mark.adbc +def test_postgresql_uuid_batch_inference(postgresql_session: AdbcDriver) -> None: + """Batch binding infers UUID columns across rows and accepts strings and nulls.""" + table_name = "adbc_uuid_batch" + first_value = uuid4() + last_value = uuid4() + + try: + postgresql_session.execute_script(f"DROP TABLE IF EXISTS {table_name}") + postgresql_session.execute_script(f"CREATE TABLE {table_name} (position INTEGER PRIMARY KEY, value UUID)") + postgresql_session.execute_many( + f"INSERT INTO {table_name} (position, value) VALUES (?, ?)", + [(1, str(first_value).upper()), (2, None), (3, last_value)], + ) + + rows = postgresql_session.execute( + f"SELECT position, value::text AS value FROM {table_name} ORDER BY position" + ).get_data() + assert [row["value"] for row in rows] == [str(first_value), None, str(last_value)] + finally: + postgresql_session.execute_script(f"DROP TABLE IF EXISTS {table_name}") diff --git a/tests/unit/adapters/test_adbc/test_uuid_binding.py b/tests/unit/adapters/test_adbc/test_uuid_binding.py new file mode 100644 index 000000000..b31f2fe0f --- /dev/null +++ b/tests/unit/adapters/test_adbc/test_uuid_binding.py @@ -0,0 +1,351 @@ +"""Unit tests for PostgreSQL-family ADBC UUID parameter binding.""" + +from typing import Any, cast +from uuid import UUID + +import pyarrow as pa +import pytest + +from sqlspec.adapters.adbc import core as adbc_core +from sqlspec.adapters.adbc._typing import AdbcConnection +from sqlspec.adapters.adbc.core import get_statement_config, prepare_postgres_uuid_bindings +from sqlspec.adapters.adbc.driver import AdbcDriver +from sqlspec.core import SQL, StatementConfig +from sqlspec.exceptions import SQLSpecError +from sqlspec.typing import UUID_UTILS_INSTALLED + +UUID_VALUE = UUID("550e8400-e29b-41d4-a716-446655440000") +OTHER_UUID_VALUE = UUID("550e8400-e29b-41d4-a716-446655440001") + + +class _AdbcUuidCursor: + def __init__(self) -> None: + self.closed = False + self.executed: list[tuple[str, object]] = [] + + def execute(self, sql: str, parameters: object = None) -> None: + self.executed.append((sql, parameters)) + + def fetch_arrow_table(self) -> pa.Table: + return pa.table({"value": [str(UUID_VALUE)]}) + + def close(self) -> None: + self.closed = True + + +class _AdbcUuidConnection: + def __init__(self, dialect: str = "postgres") -> None: + self.dialect = dialect + self.cursor_obj = _AdbcUuidCursor() + + def adbc_get_info(self) -> dict[str, str]: + return {"vendor_name": self.dialect, "driver_name": self.dialect} + + def cursor(self) -> _AdbcUuidCursor: + return self.cursor_obj + + +def _make_driver(dialect: str = "postgres") -> tuple[AdbcDriver, StatementConfig]: + base_dialect = "postgres" if dialect in {"pgvector", "paradedb"} else dialect + config = get_statement_config(base_dialect) + if dialect in {"pgvector", "paradedb"}: + config = config.replace(dialect=dialect) + connection = _AdbcUuidConnection(dialect) + driver = AdbcDriver(cast("AdbcConnection", connection), statement_config=config, dialect=dialect) + return driver, config + + +def _compile(sql: str, parameters: object, *, dialect: str = "postgres", is_many: bool = False) -> tuple[str, object]: + driver, config = _make_driver(dialect) + statement = SQL(sql, cast("Any", parameters), statement_config=config, is_many=is_many) + return driver._compiled_sql(statement, config) # pyright: ignore[reportPrivateUsage] + + +@pytest.mark.parametrize( + ("sql", "parameters"), + [ + ("SELECT ?, ?", (UUID_VALUE, "ordinary")), + ("SELECT :identifier, :label", {"identifier": UUID_VALUE, "label": "ordinary"}), + ("SELECT @identifier, @label", {"identifier": UUID_VALUE, "label": "ordinary"}), + ("SELECT $identifier, $label", {"identifier": UUID_VALUE, "label": "ordinary"}), + ("SELECT %(identifier)s, %(label)s", {"identifier": UUID_VALUE, "label": "ordinary"}), + ("SELECT %s, %s", (UUID_VALUE, "ordinary")), + ("SELECT :1, :2", (UUID_VALUE, "ordinary")), + ("SELECT $1, $2", (UUID_VALUE, "ordinary")), + ], +) +def test_postgres_uuid_binding_normalizes_supported_authoring_styles(sql: str, parameters: object) -> None: + compiled_sql, compiled_parameters = _compile(sql, parameters) + + assert compiled_sql == "SELECT CAST($1 AS UUID), $2" + assert list(cast("Any", compiled_parameters)) == [str(UUID_VALUE), "ordinary"] + + +def test_out_of_text_order_numeric_ordinals_use_parameter_ordinals() -> None: + compiled_sql, compiled_parameters = _compile("SELECT $2, $1", ("ordinary", UUID_VALUE)) + + assert compiled_sql == "SELECT CAST($2 AS UUID), $1" + assert compiled_parameters == ["ordinary", str(UUID_VALUE)] + + +def test_out_of_text_order_named_parameters_use_normalized_ordinals() -> None: + compiled_sql, compiled_parameters = _compile( + "SELECT :label, :identifier", {"identifier": UUID_VALUE, "label": "ordinary"} + ) + + assert compiled_sql == "SELECT $1, CAST($2 AS UUID)" + assert compiled_parameters == ("ordinary", str(UUID_VALUE)) + + +@pytest.mark.parametrize("dialect", ["postgres", "postgresql", "pgvector", "paradedb"]) +def test_postgres_family_aliases_bind_uuid_parameters(dialect: str) -> None: + compiled_sql, compiled_parameters = _compile("SELECT ?", (UUID_VALUE,), dialect=dialect) + + assert compiled_sql == "SELECT CAST($1 AS UUID)" + assert list(cast("Any", compiled_parameters)) == [str(UUID_VALUE)] + + +def test_reused_uuid_placeholder_is_rewritten_at_every_occurrence() -> None: + compiled_sql, compiled_parameters = _compile("SELECT $1, $1", (UUID_VALUE,)) + + assert compiled_sql == "SELECT CAST($1 AS UUID), CAST($1 AS UUID)" + assert compiled_parameters == [str(UUID_VALUE)] + + +@pytest.mark.parametrize( + ("sql", "expected_sql"), + [ + ("SELECT CAST($1 AS UUID)", "SELECT CAST($1 AS UUID)"), + ("SELECT $1::uuid", "SELECT CAST($1 AS UUID)"), + ("SELECT $1::public.uuid", "SELECT CAST($1 AS public.uuid)"), + ('SELECT $1::"public"."uuid"', 'SELECT CAST($1 AS "public"."uuid")'), + ], +) +def test_existing_uuid_cast_is_reused(sql: str, expected_sql: str) -> None: + compiled_sql, compiled_parameters = _compile(sql, (UUID_VALUE,)) + + assert compiled_sql == expected_sql + assert compiled_parameters == [str(UUID_VALUE)] + + +@pytest.mark.parametrize( + ("sql", "expected_sql"), + [ + ("SELECT CAST($1 AS TEXT)", "SELECT CAST($1 AS TEXT)"), + ("SELECT $1::varchar", "SELECT $1::varchar"), + ("SELECT $1::public.my_uuid", "SELECT $1::public.my_uuid"), + ], +) +def test_different_explicit_cast_remains_authoritative(sql: str, expected_sql: str) -> None: + compiled_sql, compiled_parameters = _compile(sql, (UUID_VALUE,)) + + assert compiled_sql == expected_sql + assert compiled_parameters == [UUID_VALUE] + + +def test_different_explicit_cast_skips_all_reused_occurrences() -> None: + compiled_sql, compiled_parameters = _compile("SELECT $1::text, $1", (UUID_VALUE,)) + + assert compiled_sql == "SELECT $1::text, $1" + assert compiled_parameters == [UUID_VALUE] + + +def test_existing_uuid_cast_is_reused_for_other_occurrences() -> None: + compiled_sql, compiled_parameters = _compile("SELECT CAST($1 AS UUID), $1", (UUID_VALUE,)) + + assert compiled_sql == "SELECT CAST($1 AS UUID), CAST($1 AS UUID)" + assert compiled_parameters == [str(UUID_VALUE)] + + +def test_parenthesized_different_cast_remains_authoritative() -> None: + compiled_sql, compiled_parameters = _compile("SELECT ($1)::text", (UUID_VALUE,)) + + assert compiled_sql == "SELECT ($1)::text" + assert compiled_parameters == [UUID_VALUE] + + +def test_parenthesized_uuid_cast_is_reused() -> None: + compiled_sql, compiled_parameters = _compile("SELECT ($1)::uuid", (UUID_VALUE,)) + + assert compiled_sql == "SELECT CAST(($1) AS UUID)" + assert compiled_parameters == [str(UUID_VALUE)] + + +def test_non_postgres_dialect_leaves_uuid_unchanged() -> None: + compiled_sql, compiled_parameters = _compile("SELECT ?", (UUID_VALUE,), dialect="sqlite") + + assert compiled_sql == "SELECT ?" + assert compiled_parameters == [UUID_VALUE] + + +@pytest.mark.parametrize("nested", [[UUID_VALUE], {"identifier": UUID_VALUE}]) +def test_nested_uuid_values_are_not_rewritten(nested: object) -> None: + compiled_sql, compiled_parameters = _compile("SELECT ?", (nested,)) + + assert compiled_sql == "SELECT $1" + if isinstance(nested, dict): + assert compiled_parameters == ['{"identifier":"550e8400-e29b-41d4-a716-446655440000"}'] + else: + assert list(cast("Any", compiled_parameters)) == [nested] + + +@pytest.mark.skipif(not UUID_UTILS_INSTALLED, reason="uuid_utils not installed") +def test_uuid_utils_values_are_normalized() -> None: + import uuid_utils + + value = uuid_utils.UUID(str(UUID_VALUE)) + compiled_sql, compiled_parameters = _compile("SELECT ?", (value,)) + + assert compiled_sql == "SELECT CAST($1 AS UUID)" + assert compiled_parameters == [str(UUID_VALUE)] + + +def test_batch_uuid_inference_uses_every_row_and_accepts_none_and_uuid_strings() -> None: + compiled_sql, compiled_parameters = _compile( + "INSERT INTO values_table (identifier, label) VALUES (?, ?)", + [(str(UUID_VALUE).upper(), "first"), (None, "second"), (UUID_VALUE, "third")], + is_many=True, + ) + + assert compiled_sql == "INSERT INTO values_table (identifier, label) VALUES (CAST($1 AS UUID), $2)" + assert compiled_parameters == [(str(UUID_VALUE), "first"), (None, "second"), (str(UUID_VALUE), "third")] + + +@pytest.mark.parametrize( + "parameters", [[(UUID_VALUE, "first"), (42, "second")], [(UUID_VALUE, "first"), ("not-a-uuid", "second")]] +) +def test_batch_uuid_inference_rejects_incompatible_values(parameters: object) -> None: + with pytest.raises(SQLSpecError, match="UUID parameter ordinal 1"): + _compile("INSERT INTO values_table VALUES (?, ?)", parameters, is_many=True) + + +def test_batch_uuid_inference_rejects_inconsistent_row_shapes() -> None: + with pytest.raises(SQLSpecError, match="Parameter count mismatch"): + _compile("INSERT INTO values_table VALUES (?, ?)", [(UUID_VALUE, "first"), (UUID_VALUE,)], is_many=True) + + +def test_batch_without_uuid_objects_is_unchanged() -> None: + parameters = [(str(UUID_VALUE),), (None,)] + + compiled_sql, compiled_parameters = _compile("INSERT INTO values_table VALUES (?)", parameters, is_many=True) + + assert compiled_sql == "INSERT INTO values_table VALUES ($1)" + assert compiled_parameters == parameters + + +def test_batch_with_scalar_rows_is_left_unchanged() -> None: + parameters = ["first", "second"] + + sql, converted = prepare_postgres_uuid_bindings("SELECT $1", parameters, is_many=True, dialect="postgres") + + assert sql == "SELECT $1" + assert converted is parameters + + +def test_batch_with_inconsistent_row_shapes_is_left_unchanged() -> None: + parameters = [(UUID_VALUE, "first"), (UUID_VALUE,)] + + sql, converted = prepare_postgres_uuid_bindings("SELECT $1, $2", parameters, is_many=True, dialect="postgres") + + assert sql == "SELECT $1, $2" + assert converted is parameters + + +def test_structural_rewrite_cache_is_bounded_and_keyed_by_dialect_and_ordinals() -> None: + rewrite = cast("Any", adbc_core)._rewrite_postgres_uuid_placeholders + rewrite.cache_clear() + + first = rewrite("SELECT $1, $2", (1,), "postgres") + second = rewrite("SELECT $1, $2", (1,), "postgres") + third = rewrite("SELECT $1, $2", (2,), "postgres") + fourth = rewrite("SELECT $1, $2", (1,), "pgvector") + + assert first == second == ("SELECT CAST($1 AS UUID), $2", (1,)) + assert third == ("SELECT $1, CAST($2 AS UUID)", (2,)) + assert fourth == first + assert rewrite.cache_info().hits == 1 + assert rewrite.cache_info().misses == 3 + assert rewrite.cache_info().maxsize == 256 + + +def test_structural_rewrite_uses_ast_and_ignores_literal_and_comment_placeholders() -> None: + rewrite = cast("Any", adbc_core)._rewrite_postgres_uuid_placeholders + + rewritten, effective = rewrite("SELECT '$1' AS literal, $1 -- $1\n", (1,), "postgres") + + assert "'$1' AS literal" in rewritten + assert "CAST($1" in rewritten + assert rewritten.count("AS UUID)") == 1 + assert "/* $1 */" in rewritten + assert effective == (1,) + + +@pytest.mark.parametrize( + ("dialect", "sql", "operator"), + [ + ("pgvector", "SELECT $1, embedding <=> $2", "<=>"), + ("pgvector", "SELECT $1, embedding <#> $2", "<#>"), + ("pgvector", "SELECT $1, embedding <+> $2", "<+>"), + ("pgvector", "SELECT $1, embedding <~> $2", "<~>"), + ("paradedb", "SELECT $1 FROM documents WHERE title @@@ $2", "@@@"), + ("paradedb", "SELECT $1 FROM documents WHERE tags &&& $2", "&&&"), + ("paradedb", "SELECT $1 FROM documents WHERE title ||| $2", "|||"), + ], +) +def test_extension_dialect_operators_survive_uuid_ast_rewrite(dialect: str, sql: str, operator: str) -> None: + compiled_sql, compiled_parameters = _compile(sql, (UUID_VALUE, "rhs"), dialect=dialect) + + assert "CAST($1 AS UUID)" in compiled_sql + assert operator in compiled_sql + assert list(cast("Any", compiled_parameters)) == [str(UUID_VALUE), "rhs"] + + +def test_structural_rewrite_reports_sqlglot_parse_errors() -> None: + rewrite = cast("Any", adbc_core)._rewrite_postgres_uuid_placeholders + + with pytest.raises(SQLSpecError, match="Failed to parse PostgreSQL ADBC SQL for UUID parameter binding"): + rewrite("SELECT (", (1,), "postgres") + + +def test_value_dependent_binding_does_not_leak_between_same_sql_calls() -> None: + driver, config = _make_driver() + + def compile_value(value: object) -> tuple[str, object]: + statement = SQL("SELECT ?", (value,), statement_config=config) + return driver._compiled_sql(statement, config) # pyright: ignore[reportPrivateUsage] + + ordinary_before = compile_value("ordinary") + uuid_call = compile_value(UUID_VALUE) + ordinary_after = compile_value("ordinary") + + assert ordinary_before == ("SELECT $1", ["ordinary"]) + assert uuid_call == ("SELECT CAST($1 AS UUID)", [str(UUID_VALUE)]) + assert ordinary_after == ordinary_before + + +def test_value_dependent_binding_does_not_leak_in_reverse_order() -> None: + driver, config = _make_driver() + + def compile_value(value: object) -> tuple[str, object]: + statement = SQL("SELECT ?", (value,), statement_config=config) + return driver._compiled_sql(statement, config) # pyright: ignore[reportPrivateUsage] + + uuid_before = compile_value(UUID_VALUE) + ordinary_call = compile_value("ordinary") + uuid_after = compile_value(OTHER_UUID_VALUE) + + assert uuid_before == ("SELECT CAST($1 AS UUID)", [str(UUID_VALUE)]) + assert ordinary_call == ("SELECT $1", ["ordinary"]) + assert uuid_after == ("SELECT CAST($1 AS UUID)", [str(OTHER_UUID_VALUE)]) + + +def test_select_to_arrow_uses_uuid_rewrite() -> None: + connection = _AdbcUuidConnection() + config = get_statement_config("postgres") + driver = AdbcDriver(cast("AdbcConnection", connection), statement_config=config, dialect="postgres") + + result = driver.select_to_arrow("SELECT ? AS value", UUID_VALUE) + + assert result.data.to_pydict() == {"value": [str(UUID_VALUE)]} + assert connection.cursor_obj.executed == [("SELECT CAST($1 AS UUID) AS value", [str(UUID_VALUE)])] + assert connection.cursor_obj.closed is True diff --git a/tests/unit/utils/test_mypyc_inventory.py b/tests/unit/utils/test_mypyc_inventory.py index 9a7e748ac..d995c1c39 100644 --- a/tests/unit/utils/test_mypyc_inventory.py +++ b/tests/unit/utils/test_mypyc_inventory.py @@ -8,10 +8,12 @@ import sys from pathlib import Path from types import ModuleType +from uuid import UUID import sqlspec.utils.correlation as correlation_module import sqlspec.utils.schema as schema_module import sqlspec.utils.sync_tools as sync_tools_module +from sqlspec.adapters.adbc.core import prepare_postgres_uuid_bindings from sqlspec.utils.correlation import CorrelationContext if sys.version_info >= (3, 11): @@ -21,6 +23,16 @@ PROJECT_ROOT = Path(__file__).resolve().parents[3] +def test_adbc_postgres_uuid_binding_runs_across_mypyc_boundary() -> None: + """The compiled ADBC core should retain UUID rewrite behavior.""" + value = UUID("550e8400-e29b-41d4-a716-446655440000") + + sql, parameters = prepare_postgres_uuid_bindings("SELECT $1", [value], is_many=False, dialect="postgres") + + assert sql == "SELECT CAST($1 AS UUID)" + assert parameters == [str(value)] + + def _load_mypyc_boundary_map_module() -> ModuleType: module_path = PROJECT_ROOT / "tools" / "scripts" / "mypyc_boundary_map.py" spec = importlib.util.spec_from_file_location("mypyc_boundary_map_for_tests", module_path) @@ -83,6 +95,7 @@ def test_pyproject_mypyc_include_patterns_cover_smoke_critical_modules() -> None "sqlspec/core/result/_base.py", "sqlspec/core/splitter.py", "sqlspec/driver/_query_cache.py", + "sqlspec/adapters/adbc/core.py", "sqlspec/adapters/sqlite/core.py", "sqlspec/adapters/psqlpy/core.py", "sqlspec/adapters/sqlite/pool.py",