sqlspec 0.55.0__py3-none-any.whl → 0.56.0__py3-none-any.whl

sqlspec/adapters/adbc/adk/store.py

@@ -72,6 +72,10 @@ class AdbcADKStore(BaseSyncADKStore["AdbcConfig"]):
 
     def create_tables(self) -> None:
         """Create tables if they don't exist."""
+        if not self.create_schema_enabled:
+            self.reconcile_schema()
+            return
+
         self._create_tables()
 
     def create_session(
@@ -496,19 +500,6 @@ class AdbcADKStore(BaseSyncADKStore["AdbcConfig"]):
         )
         """
 
-    def _metadata_seed_sql(self) -> str:
-        if self._dialect in {DIALECT_POSTGRESQL, DIALECT_SQLITE, DIALECT_DUCKDB}:
-            return f"""
-            INSERT INTO {self._metadata_table} (key, value)
-            VALUES ('schema_version', '1')
-            ON CONFLICT(key) DO NOTHING
-            """
-        return f"""
-        INSERT INTO {self._metadata_table} (key, value)
-        SELECT 'schema_version', '1'
-        WHERE NOT EXISTS (SELECT 1 FROM {self._metadata_table} WHERE key = 'schema_version')
-        """
-
     def _drop_app_states_table_sql(self) -> str:
         return f"DROP TABLE IF EXISTS {self._app_state_table}"
 
@@ -565,9 +556,6 @@ class AdbcADKStore(BaseSyncADKStore["AdbcConfig"]):
 
                 cursor.execute(self._metadata_table_ddl())
                 conn.commit()
-
-                cursor.execute(self._metadata_seed_sql())
-                conn.commit()
             finally:
                 cursor.close()
 
@@ -1192,6 +1180,10 @@ class AdbcADKMemoryStore(BaseSyncADKMemoryStore["AdbcConfig"]):
 
     def create_tables(self) -> None:
         """Create tables if they don't exist."""
+        if not self.create_schema_enabled:
+            self.reconcile_schema()
+            return
+
         self._create_tables()
 
     def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object | None" = None) -> int:

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 @@ __all__ = (
     "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 @@ __all__ = (
 )
 
 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"),
@@ -665,9 +673,6 @@ def build_profile() -> "DriverParameterProfile":
     )
 
 
-driver_profile: Final[DriverParameterProfile] = build_profile()
-
-
 def get_statement_config(detected_dialect: str) -> StatementConfig:
     """Create statement configuration for the specified dialect."""
     default_style, supported_styles = DIALECT_PARAMETER_STYLES.get(
@@ -850,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}"
@@ -1005,3 +1181,6 @@ def _lift_flightsql_db_kwargs(config: "dict[str, Any]") -> None:
     config.pop("gizmosql_backend", None)
     if db_kwargs_dict:
         config["db_kwargs"] = db_kwargs_dict
+
+
+driver_profile: Final[DriverParameterProfile] = build_profile()

sqlspec/adapters/adbc/driver.py

@@ -22,6 +22,7 @@ from sqlspec.adapters.adbc.core import (
     normalize_postgres_empty_parameters,
     normalize_script_rowcount,
     prepare_postgres_parameters,
+    prepare_postgres_uuid_bindings,
     resolve_column_names,
     resolve_dialect_name,
     resolve_many_rowcount,
@@ -141,14 +142,6 @@ class AdbcSelectStreamSource:
                 cursor_manager.__exit__(None, None, None)
 
 
-def _fetch_record_batch(cursor: Any) -> Any:
-    fetch_record_batch = getattr(cursor, "fetch_record_batch", None)
-    if fetch_record_batch is None:
-        msg = "ADBC cursor does not expose fetch_record_batch() for native row streaming."
-        raise SQLSpecError(msg)
-    return fetch_record_batch()
-
-
 @final
 class AdbcDriver(SyncDriverAdapterBase):
     """ADBC driver for Arrow Database Connectivity.
@@ -197,6 +190,16 @@ class AdbcDriver(SyncDriverAdapterBase):
     # 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.
 
@@ -755,4 +758,12 @@ class AdbcDriver(SyncDriverAdapterBase):
         return fallback
 
 
+def _fetch_record_batch(cursor: Any) -> Any:
+    fetch_record_batch = getattr(cursor, "fetch_record_batch", None)
+    if fetch_record_batch is None:
+        msg = "ADBC cursor does not expose fetch_record_batch() for native row streaming."
+        raise SQLSpecError(msg)
+    return fetch_record_batch()
+
+
 register_driver_profile("adbc", driver_profile)

sqlspec/adapters/adbc/litestar/store.py

@@ -59,6 +59,78 @@ class ADBCStore(BaseSQLSpecStore["AdbcConfig"]):
         super().__init__(config)
         self._dialect: str | None = None
 
+    async def create_table(self) -> None:
+        """Create the session table if it doesn't exist."""
+        if not self.create_schema_enabled:
+            await self.reconcile_schema()
+            return
+        await async_(self._create_table)()
+        await self.reconcile_schema(assume_existing=True)
+
+    async def get(self, key: str, renew_for: "int | timedelta | None" = None) -> "bytes | None":
+        """Get a session value by key.
+
+        Args:
+            key: Session ID to retrieve.
+            renew_for: If given, renew the expiry time for this duration.
+
+        Returns:
+            Session data as bytes if found and not expired, None otherwise.
+        """
+        return await async_(self._get)(key, renew_for)
+
+    async def set(self, key: str, value: "str | bytes", expires_in: "int | timedelta | None" = None) -> None:
+        """Store a session value.
+
+        Args:
+            key: Session ID.
+            value: Session data.
+            expires_in: Time until expiration.
+        """
+        await async_(self._set)(key, value, expires_in)
+
+    async def delete(self, key: str) -> None:
+        """Delete a session by key.
+
+        Args:
+            key: Session ID to delete.
+        """
+        await async_(self._delete)(key)
+
+    async def delete_all(self) -> None:
+        """Delete all sessions from the store."""
+        await async_(self._delete_all)()
+
+    async def exists(self, key: str) -> bool:
+        """Check if a session key exists and is not expired.
+
+        Args:
+            key: Session ID to check.
+
+        Returns:
+            True if the session exists and is not expired.
+        """
+        return await async_(self._exists)(key)
+
+    async def expires_in(self, key: str) -> "int | None":
+        """Get the time in seconds until the session expires.
+
+        Args:
+            key: Session ID to check.
+
+        Returns:
+            Seconds until expiration, or None if no expiry or key doesn't exist.
+        """
+        return await async_(self._expires_in)(key)
+
+    async def delete_expired(self) -> int:
+        """Delete all expired sessions.
+
+        Returns:
+            Number of sessions deleted.
+        """
+        return await async_(self._delete_expired)()
+
     def _get_dialect(self) -> str:
         """Get the database dialect, caching it after first access.
 
@@ -220,10 +292,6 @@ class ADBCStore(BaseSQLSpecStore["AdbcConfig"]):
 
         return [f"DROP INDEX IF EXISTS idx_{self._table_name}_expires_at", f"DROP TABLE IF EXISTS {self._table_name}"]
 
-    async def create_table(self) -> None:
-        """Create the session table if it doesn't exist."""
-        await async_(self._create_table)()
-
     def _get(self, key: str, renew_for: "int | timedelta | None" = None) -> "bytes | None":
         """Synchronous implementation of get using ADBC driver."""
         p1 = self._get_param_placeholder(1)
@@ -258,18 +326,6 @@ class ADBCStore(BaseSQLSpecStore["AdbcConfig"]):
 
             return bytes(data)
 
-    async def get(self, key: str, renew_for: "int | timedelta | None" = None) -> "bytes | None":
-        """Get a session value by key.
-
-        Args:
-            key: Session ID to retrieve.
-            renew_for: If given, renew the expiry time for this duration.
-
-        Returns:
-            Session data as bytes if found and not expired, None otherwise.
-        """
-        return await async_(self._get)(key, renew_for)
-
     def _set(self, key: str, value: "str | bytes", expires_in: "int | timedelta | None" = None) -> None:
         """Synchronous implementation of set using ADBC driver with dialect-specific UPSERT."""
         data = self._value_to_bytes(value)
@@ -364,16 +420,6 @@ class ADBCStore(BaseSQLSpecStore["AdbcConfig"]):
             driver.execute(sql, key, data, expires_at)
             driver.commit()
 
-    async def set(self, key: str, value: "str | bytes", expires_in: "int | timedelta | None" = None) -> None:
-        """Store a session value.
-
-        Args:
-            key: Session ID.
-            value: Session data.
-            expires_in: Time until expiration.
-        """
-        await async_(self._set)(key, value, expires_in)
-
     def _delete(self, key: str) -> None:
         """Synchronous implementation of delete using ADBC driver."""
         p1 = self._get_param_placeholder(1)
@@ -383,14 +429,6 @@ class ADBCStore(BaseSQLSpecStore["AdbcConfig"]):
             driver.execute(sql, key)
             driver.commit()
 
-    async def delete(self, key: str) -> None:
-        """Delete a session by key.
-
-        Args:
-            key: Session ID to delete.
-        """
-        await async_(self._delete)(key)
-
     def _delete_all(self) -> None:
         """Synchronous implementation of delete_all using ADBC driver."""
 
@@ -400,10 +438,6 @@ class ADBCStore(BaseSQLSpecStore["AdbcConfig"]):
             driver.execute(sql)
             driver.commit()
 
-    async def delete_all(self) -> None:
-        """Delete all sessions from the store."""
-        await async_(self._delete_all)()
-
     def _exists(self, key: str) -> bool:
         """Synchronous implementation of exists using ADBC driver."""
 
@@ -419,17 +453,6 @@ class ADBCStore(BaseSQLSpecStore["AdbcConfig"]):
         with self._config.provide_session() as driver:
             return bool(driver.select_one_or_none(sql, key) is not None)
 
-    async def exists(self, key: str) -> bool:
-        """Check if a session key exists and is not expired.
-
-        Args:
-            key: Session ID to check.
-
-        Returns:
-            True if the session exists and is not expired.
-        """
-        return await async_(self._exists)(key)
-
     def _expires_in(self, key: str) -> "int | None":
         """Synchronous implementation of expires_in using ADBC driver."""
         p1 = self._get_param_placeholder(1)
@@ -460,17 +483,6 @@ class ADBCStore(BaseSQLSpecStore["AdbcConfig"]):
             delta = expires_at - now
             return int(delta.total_seconds())
 
-    async def expires_in(self, key: str) -> "int | None":
-        """Get the time in seconds until the session expires.
-
-        Args:
-            key: Session ID to check.
-
-        Returns:
-            Seconds until expiration, or None if no expiry or key doesn't exist.
-        """
-        return await async_(self._expires_in)(key)
-
     def _delete_expired(self) -> int:
         """Synchronous implementation of delete_expired using ADBC driver."""
         current_ts = self._get_current_timestamp_expr()
@@ -500,11 +512,3 @@ class ADBCStore(BaseSQLSpecStore["AdbcConfig"]):
             if count > 0:
                 self._log_delete_expired(count)
             return count
-
-    async def delete_expired(self) -> int:
-        """Delete all expired sessions.
-
-        Returns:
-            Number of sessions deleted.
-        """
-        return await async_(self._delete_expired)()

sqlspec/adapters/aiomysql/adk/store.py

@@ -64,13 +64,16 @@ class AiomysqlADKStore(BaseAsyncADKStore["AiomysqlConfig"]):
 
     async def create_tables(self) -> None:
         """Create all ADK session tables if they don't exist."""
+        if not self.create_schema_enabled:
+            await self.reconcile_schema()
+            return
+
         async with self._config.provide_session() as driver:
             await driver.execute_script(await self._sessions_table_ddl())
             await driver.execute_script(await self._events_table_ddl())
             await driver.execute_script(await self._app_states_table_ddl())
             await driver.execute_script(await self._user_states_table_ddl())
             await driver.execute_script(await self._metadata_table_ddl())
-            await driver.execute_script(await self._metadata_seed_sql())
 
     async def create_session(
         self, session_id: str, app_name: str, user_id: str, state: "dict[str, Any]", owner_id: "Any | None" = None
@@ -465,10 +468,6 @@ class AiomysqlADKStore(BaseAsyncADKStore["AiomysqlConfig"]):
         """Get MySQL CREATE TABLE SQL for ADK metadata."""
         return _mysql_metadata_ddl(self._metadata_table)
 
-    async def _metadata_seed_sql(self) -> str:
-        """Get MySQL metadata seed SQL."""
-        return f"INSERT IGNORE INTO {self._metadata_table} (`key`, value) VALUES ('schema_version', '1')"
-
     def _drop_app_states_table_sql(self) -> str:
         """Get MySQL DROP TABLE SQL for app-scoped state."""
         return f"DROP TABLE IF EXISTS {self._app_state_table}"
@@ -503,6 +502,10 @@ class AiomysqlADKMemoryStore(BaseAsyncADKMemoryStore["AiomysqlConfig"]):
 
     async def create_tables(self) -> None:
         """Create the memory table and indexes if they don't exist."""
+        if not self.create_schema_enabled:
+            await self.reconcile_schema()
+            return
+
         if not self._enabled:
             return