sqlspec 0.32.0__py3-none-any.whl

sqlspec/adapters/spanner/config.py

@@ -0,0 +1,241 @@
+"""Spanner configuration."""
+
+from collections.abc import Callable, Generator
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Any, ClassVar, TypedDict, cast
+
+from google.cloud.spanner_v1 import Client
+from google.cloud.spanner_v1.pool import AbstractSessionPool, FixedSizePool
+from typing_extensions import NotRequired
+
+from sqlspec.adapters.spanner._types import SpannerConnection
+from sqlspec.adapters.spanner.driver import SpannerSyncDriver, spanner_statement_config
+from sqlspec.config import SyncDatabaseConfig
+from sqlspec.exceptions import ImproperConfigurationError
+from sqlspec.utils.serializers import from_json, to_json
+
+if TYPE_CHECKING:
+    from google.auth.credentials import Credentials
+    from google.cloud.spanner_v1.database import Database
+
+    from sqlspec.config import ExtensionConfigs
+    from sqlspec.core import StatementConfig
+    from sqlspec.observability import ObservabilityConfig
+
+__all__ = ("SpannerConnectionParams", "SpannerDriverFeatures", "SpannerPoolParams", "SpannerSyncConfig")
+
+
+class SpannerConnectionParams(TypedDict):
+    """Spanner connection parameters."""
+
+    project: "NotRequired[str]"
+    instance_id: "NotRequired[str]"
+    database_id: "NotRequired[str]"
+    credentials: "NotRequired[Credentials]"
+    client_options: "NotRequired[dict[str, Any]]"
+    extra: "NotRequired[dict[str, Any]]"
+
+
+class SpannerPoolParams(SpannerConnectionParams):
+    """Session pool configuration."""
+
+    pool_type: "NotRequired[type[AbstractSessionPool]]"
+    min_sessions: "NotRequired[int]"
+    max_sessions: "NotRequired[int]"
+    labels: "NotRequired[dict[str, str]]"
+    ping_interval: "NotRequired[int]"
+
+
+class SpannerDriverFeatures(TypedDict):
+    """Driver feature flags for Spanner."""
+
+    enable_uuid_conversion: "NotRequired[bool]"
+    json_serializer: "NotRequired[Callable[[Any], str]]"
+    json_deserializer: "NotRequired[Callable[[str], Any]]"
+    session_labels: "NotRequired[dict[str, str]]"
+
+
+class SpannerSyncConfig(SyncDatabaseConfig["SpannerConnection", "AbstractSessionPool", SpannerSyncDriver]):
+    """Spanner configuration and session management."""
+
+    driver_type: ClassVar[type["SpannerSyncDriver"]] = SpannerSyncDriver
+    connection_type: ClassVar[type["SpannerConnection"]] = cast("type[SpannerConnection]", SpannerConnection)
+    supports_transactional_ddl: ClassVar[bool] = False
+    supports_native_arrow_export: ClassVar[bool] = True
+    supports_native_arrow_import: ClassVar[bool] = True
+    supports_native_parquet_export: ClassVar[bool] = False
+    supports_native_parquet_import: ClassVar[bool] = False
+    requires_staging_for_load: ClassVar[bool] = False
+
+    def __init__(
+        self,
+        *,
+        pool_config: "SpannerPoolParams | dict[str, Any] | None" = None,
+        pool_instance: "AbstractSessionPool | None" = None,
+        migration_config: "dict[str, Any] | None" = None,
+        statement_config: "StatementConfig | None" = None,
+        driver_features: "SpannerDriverFeatures | dict[str, Any] | None" = None,
+        bind_key: "str | None" = None,
+        extension_config: "ExtensionConfigs | None" = None,
+        observability_config: "ObservabilityConfig | None" = None,
+    ) -> None:
+        self.pool_config = dict(pool_config) if pool_config else {}
+
+        self.pool_config.setdefault("min_sessions", 1)
+        self.pool_config.setdefault("max_sessions", 10)
+        self.pool_config.setdefault("pool_type", FixedSizePool)
+
+        features: dict[str, Any] = dict(driver_features) if driver_features else {}
+        features.setdefault("enable_uuid_conversion", True)
+        features.setdefault("json_serializer", to_json)
+        features.setdefault("json_deserializer", from_json)
+
+        base_statement_config = statement_config or spanner_statement_config
+
+        super().__init__(
+            pool_config=self.pool_config,
+            pool_instance=pool_instance,
+            migration_config=migration_config,
+            statement_config=base_statement_config,
+            driver_features=features,
+            bind_key=bind_key,
+            extension_config=extension_config,
+            observability_config=observability_config,
+        )
+
+        self._client: Client | None = None
+        self._database: Database | None = None
+
+    def _get_client(self) -> Client:
+        if self._client is None:
+            self._client = Client(
+                project=self.pool_config.get("project"),
+                credentials=self.pool_config.get("credentials"),
+                client_options=self.pool_config.get("client_options"),
+            )
+        return self._client
+
+    def get_database(self) -> "Database":
+        instance_id = self.pool_config.get("instance_id")
+        database_id = self.pool_config.get("database_id")
+        if not instance_id or not database_id:
+            msg = "instance_id and database_id are required."
+            raise ImproperConfigurationError(msg)
+
+        if self.pool_instance is None:
+            self.pool_instance = self.provide_pool()
+
+        if self._database is None:
+            client = self._get_client()
+            self._database = client.instance(instance_id).database(database_id, pool=self.pool_instance)  # type: ignore[no-untyped-call]
+        return self._database
+
+    def create_connection(self) -> SpannerConnection:
+        instance_id = self.pool_config.get("instance_id")
+        database_id = self.pool_config.get("database_id")
+        if not instance_id or not database_id:
+            msg = "instance_id and database_id are required."
+            raise ImproperConfigurationError(msg)
+
+        if self.pool_instance is None:
+            self.pool_instance = self.provide_pool()
+
+        client = self._get_client()
+        database = client.instance(instance_id).database(database_id, pool=self.pool_instance)  # type: ignore[no-untyped-call]
+        return cast("SpannerConnection", database.snapshot())
+
+    def _create_pool(self) -> AbstractSessionPool:
+        instance_id = self.pool_config.get("instance_id")
+        database_id = self.pool_config.get("database_id")
+        if not instance_id or not database_id:
+            msg = "instance_id and database_id are required."
+            raise ImproperConfigurationError(msg)
+
+        pool_type = cast("type[AbstractSessionPool]", self.pool_config.get("pool_type", FixedSizePool))
+
+        pool_kwargs: dict[str, Any] = {}
+        if pool_type is FixedSizePool:
+            if "size" in self.pool_config:
+                pool_kwargs["size"] = self.pool_config["size"]
+            elif "max_sessions" in self.pool_config:
+                pool_kwargs["size"] = self.pool_config["max_sessions"]
+            if "labels" in self.pool_config:
+                pool_kwargs["labels"] = self.pool_config["labels"]
+        else:
+            valid_pool_keys = {"size", "labels", "ping_interval"}
+            pool_kwargs = {k: v for k, v in self.pool_config.items() if k in valid_pool_keys and v is not None}
+            if "size" not in pool_kwargs and "max_sessions" in self.pool_config:
+                pool_kwargs["size"] = self.pool_config["max_sessions"]
+
+        pool_factory = cast("Callable[..., AbstractSessionPool]", pool_type)
+        return pool_factory(**pool_kwargs)
+
+    def _close_pool(self) -> None:
+        if self.pool_instance and hasattr(self.pool_instance, "close"):
+            cast("Any", self.pool_instance).close()
+
+    @contextmanager
+    def provide_connection(
+        self, *args: Any, transaction: "bool" = False, **kwargs: Any
+    ) -> Generator[SpannerConnection, None, None]:
+        """Yield a Snapshot (default) or Transaction context from the configured pool.
+
+        Args:
+            *args: Additional positional arguments (unused, for interface compatibility).
+            transaction: If True, yields a Transaction context that supports
+                execute_update() for DML statements. If False (default), yields
+                a read-only Snapshot context for SELECT queries.
+            **kwargs: Additional keyword arguments (unused, for interface compatibility).
+
+        Note: For complex transactional logic with retries, use database.run_in_transaction()
+        directly. The Transaction context here auto-commits on successful exit.
+        """
+        database = self.get_database()
+        if transaction:
+            session = cast("Any", database).session()
+            session.create()
+            try:
+                txn = session.transaction()
+                txn.__enter__()
+                try:
+                    yield cast("SpannerConnection", txn)
+                    if hasattr(txn, "_transaction_id") and txn._transaction_id is not None:
+                        txn.commit()
+                except Exception:
+                    if hasattr(txn, "_transaction_id") and txn._transaction_id is not None:
+                        txn.rollback()
+                    raise
+            finally:
+                session.delete()
+        else:
+            with cast("Any", database).snapshot(multi_use=True) as snapshot:
+                yield cast("SpannerConnection", snapshot)
+
+    @contextmanager
+    def provide_session(
+        self, *args: Any, statement_config: "StatementConfig | None" = None, transaction: "bool" = False, **kwargs: Any
+    ) -> Generator[SpannerSyncDriver, None, None]:
+        with self.provide_connection(*args, transaction=transaction, **kwargs) as connection:
+            driver = self.driver_type(
+                connection=connection,
+                statement_config=statement_config or self.statement_config,
+                driver_features=self.driver_features,
+            )
+            yield self._prepare_driver(driver)
+
+    @contextmanager
+    def provide_write_session(
+        self, *args: Any, statement_config: "StatementConfig | None" = None, **kwargs: Any
+    ) -> Generator[SpannerSyncDriver, None, None]:
+        with self.provide_session(*args, statement_config=statement_config, transaction=True, **kwargs) as driver:
+            yield driver
+
+    def get_signature_namespace(self) -> dict[str, Any]:
+        namespace = super().get_signature_namespace()
+        namespace.update({
+            "SpannerSyncConfig": SpannerSyncConfig,
+            "SpannerConnectionParams": SpannerConnectionParams,
+            "SpannerDriverFeatures": SpannerDriverFeatures,
+            "SpannerSyncDriver": SpannerSyncDriver,
+        })
+        return namespace

sqlspec/adapters/spanner/data_dictionary.py

@@ -0,0 +1,95 @@
+"""Spanner metadata queries using INFORMATION_SCHEMA."""
+
+from typing import TYPE_CHECKING, Any, cast
+
+from sqlspec.driver import SyncDataDictionaryBase, SyncDriverAdapterBase
+
+if TYPE_CHECKING:
+    from sqlspec.driver import VersionInfo
+
+
+__all__ = ("SpannerDataDictionary",)
+
+
+class SpannerDataDictionary(SyncDataDictionaryBase):
+    """Fetch table, column, and index metadata from Spanner."""
+
+    def get_version(self, driver: "SyncDriverAdapterBase") -> "VersionInfo | None":
+        _ = driver
+        return None
+
+    def get_feature_flag(self, driver: "SyncDriverAdapterBase", feature: str) -> bool:
+        _ = driver
+        feature_checks = {
+            "supports_json": True,
+            "supports_generators": False,
+            "supports_index_clustering": True,
+            "supports_interleaved_tables": True,
+        }
+        return feature_checks.get(feature, False)
+
+    def get_optimal_type(self, driver: "SyncDriverAdapterBase", type_category: str) -> str:
+        _ = driver
+        type_map = {
+            "json": "JSON",
+            "uuid": "BYTES(16)",
+            "boolean": "BOOL",
+            "timestamp": "TIMESTAMP",
+            "text": "STRING(MAX)",
+            "blob": "BYTES(MAX)",
+            "numeric": "NUMERIC",
+            "bignumeric": "NUMERIC",
+            "array": "ARRAY",
+        }
+        return type_map.get(type_category, "STRING(MAX)")
+
+    def get_tables(self, driver: "SyncDriverAdapterBase", schema: "str | None" = None) -> "list[str]":
+        sql = 'SELECT TABLE_NAME AS "table_name" FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_SCHEMA = @schema'
+        params: dict[str, Any]
+        if schema is None:
+            sql = "SELECT TABLE_NAME AS \"table_name\" FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_SCHEMA = ''"
+            params = {}
+        else:
+            params = {"schema": schema}
+
+        results = driver.select(sql, params)
+        return [cast("str", row["table_name"]) for row in results]
+
+    def get_columns(
+        self, driver: "SyncDriverAdapterBase", table: str, schema: "str | None" = None
+    ) -> "list[dict[str, Any]]":
+        sql = """
+            SELECT COLUMN_NAME AS "column_name", SPANNER_TYPE AS "spanner_type", IS_NULLABLE AS "is_nullable"
+            FROM INFORMATION_SCHEMA.COLUMNS
+            WHERE TABLE_NAME = @table
+        """
+        params: dict[str, Any] = {"table": table}
+        if schema is not None:
+            sql = f"{sql} AND TABLE_SCHEMA = @schema"
+            params["schema"] = schema
+        else:
+            sql = f"{sql} AND TABLE_SCHEMA = ''"
+
+        results = driver.select(sql, params)
+        return [
+            {"name": row["column_name"], "type": row["spanner_type"], "nullable": row["is_nullable"] == "YES"}
+            for row in results
+        ]
+
+    def get_indexes(
+        self, driver: "SyncDriverAdapterBase", table: str, schema: "str | None" = None
+    ) -> "list[dict[str, Any]]":
+        sql = """
+            SELECT INDEX_NAME AS "index_name", INDEX_TYPE AS "index_type", IS_UNIQUE AS "is_unique"
+            FROM INFORMATION_SCHEMA.INDEXES
+            WHERE TABLE_NAME = @table
+        """
+        params: dict[str, Any] = {"table": table}
+        if schema is not None:
+            sql = f"{sql} AND TABLE_SCHEMA = @schema"
+            params["schema"] = schema
+        else:
+            sql = f"{sql} AND TABLE_SCHEMA = ''"
+
+        results = driver.select(sql, params)
+        return [{"name": row["index_name"], "type": row["index_type"], "unique": row["is_unique"]} for row in results]

sqlspec/adapters/spanner/dialect/__init__.py

@@ -0,0 +1,6 @@
+"""Spanner dialect submodule."""
+
+from sqlspec.adapters.spanner.dialect._spangres import Spangres
+from sqlspec.adapters.spanner.dialect._spanner import Spanner
+
+__all__ = ("Spangres", "Spanner")

sqlspec/adapters/spanner/dialect/_spangres.py

@@ -0,0 +1,52 @@
+r"""Google Cloud Spanner PostgreSQL-interface dialect ("Spangres")."""
+
+from typing import Any, cast
+
+from sqlglot import exp
+from sqlglot.dialects.postgres import Postgres
+from sqlglot.tokens import TokenType
+
+__all__ = ("Spangres",)
+
+_ROW_DELETION_NAME = "ROW_DELETION_POLICY"
+_TTL_MIN_COMPONENTS = 2
+
+
+class Spangres(Postgres):
+    """Spanner PostgreSQL-compatible dialect."""
+
+    class Parser(Postgres.Parser):
+        """Parse Spanner row deletion policies."""
+
+        def _parse_property(self) -> exp.Expression:
+            if self._match_text_seq("ROW", "DELETION", "POLICY"):  # type: ignore[no-untyped-call]
+                self._match(TokenType.L_PAREN)  # type: ignore[no-untyped-call]
+                self._match_text_seq("OLDER_THAN")  # type: ignore[no-untyped-call]
+                self._match(TokenType.L_PAREN)  # type: ignore[no-untyped-call]
+                column = cast("exp.Expression", self._parse_id_var())
+                self._match(TokenType.COMMA)  # type: ignore[no-untyped-call]
+                self._match_text_seq("INTERVAL")  # type: ignore[no-untyped-call]
+                interval = cast("exp.Expression", self._parse_expression())
+                self._match(TokenType.R_PAREN)  # type: ignore[no-untyped-call]
+                self._match(TokenType.R_PAREN)  # type: ignore[no-untyped-call]
+
+                return exp.Property(
+                    this=exp.Literal.string(_ROW_DELETION_NAME), value=exp.Tuple(expressions=[column, interval])
+                )
+
+            return cast("exp.Expression", super()._parse_property())
+
+    class Generator(Postgres.Generator):
+        """Generate Spanner row deletion policies."""
+
+        def property_sql(self, expression: exp.Property) -> str:
+            if getattr(expression.this, "name", "").upper() == _ROW_DELETION_NAME:
+                values = cast("Any", expression.args.get("value"))
+                if values and getattr(values, "expressions", None) and len(values.expressions) >= _TTL_MIN_COMPONENTS:
+                    column = self.sql(values.expressions[0])
+                    interval_sql = self.sql(values.expressions[1])
+                    if not interval_sql.upper().startswith("INTERVAL"):
+                        interval_sql = f"INTERVAL {interval_sql}"
+                    return f"ROW DELETION POLICY (OLDER_THAN({column}, {interval_sql}))"
+
+            return super().property_sql(expression)

sqlspec/adapters/spanner/dialect/_spanner.py

@@ -0,0 +1,123 @@
+"""Google Cloud Spanner SQL dialect (GoogleSQL variant).
+
+Extends the BigQuery dialect with Spanner-only DDL features:
+`INTERLEAVE IN PARENT` for interleaved tables and `ROW DELETION POLICY`
+for row-level time-to-live policies (GoogleSQL).
+"""
+
+from typing import Any, cast
+
+from sqlglot import exp
+from sqlglot.dialects.bigquery import BigQuery
+from sqlglot.tokens import TokenType
+
+__all__ = ("Spanner",)
+
+
+_SPANNER_KEYWORDS: dict[str, TokenType] = {}
+interleave_token = getattr(TokenType, "INTERLEAVE", None)
+if interleave_token is not None:
+    _SPANNER_KEYWORDS["INTERLEAVE"] = interleave_token
+ttl_token = getattr(TokenType, "TTL", None)
+if ttl_token is not None:
+    _SPANNER_KEYWORDS["TTL"] = ttl_token
+
+_TTL_MIN_COMPONENTS = 2
+_ROW_DELETION_NAME = "ROW_DELETION_POLICY"
+
+
+class Spanner(BigQuery):
+    """Google Cloud Spanner SQL dialect."""
+
+    class Tokenizer(BigQuery.Tokenizer):
+        """Tokenizer adds Spanner-only keywords when supported by sqlglot."""
+
+        KEYWORDS = {**BigQuery.Tokenizer.KEYWORDS, **_SPANNER_KEYWORDS}
+
+    class Parser(BigQuery.Parser):
+        """Parse Spanner extensions such as INTERLEAVE and row deletion policies."""
+
+        def _parse_table_parts(
+            self, schema: "bool" = False, is_db_reference: "bool" = False, wildcard: "bool" = False
+        ) -> exp.Table:
+            """Parse Spanner table options including interleaving metadata."""
+            table = super()._parse_table_parts(schema=schema, is_db_reference=is_db_reference, wildcard=wildcard)
+
+            if self._match_text_seq("INTERLEAVE", "IN", "PARENT"):  # type: ignore[no-untyped-call]
+                parent = cast("exp.Expression", self._parse_table(schema=True, is_db_reference=True))
+                on_delete: str | None = None
+
+                if self._match_text_seq("ON", "DELETE"):  # type: ignore[no-untyped-call]
+                    if self._match_text_seq("CASCADE"):  # type: ignore[no-untyped-call]
+                        on_delete = "CASCADE"
+                    elif self._match_text_seq("NO", "ACTION"):  # type: ignore[no-untyped-call]
+                        on_delete = "NO ACTION"
+
+                table.set("interleave_parent", parent)
+                if on_delete:
+                    table.set("interleave_on_delete", on_delete)
+
+            return table
+
+        def _parse_property(self) -> exp.Expression:
+            """Parse Spanner row deletion policy or PostgreSQL-style TTL."""
+            if self._match_text_seq("ROW", "DELETION", "POLICY"):  # type: ignore[no-untyped-call]
+                self._match(TokenType.L_PAREN)  # type: ignore[no-untyped-call]
+                self._match_text_seq("OLDER_THAN")  # type: ignore[no-untyped-call]
+                self._match(TokenType.L_PAREN)  # type: ignore[no-untyped-call]
+                column = cast("exp.Expression", self._parse_id_var())
+                self._match(TokenType.COMMA)  # type: ignore[no-untyped-call]
+                self._match_text_seq("INTERVAL")  # type: ignore[no-untyped-call]
+                interval = cast("exp.Expression", self._parse_expression())
+                self._match(TokenType.R_PAREN)  # type: ignore[no-untyped-call]
+                self._match(TokenType.R_PAREN)  # type: ignore[no-untyped-call]
+
+                return exp.Property(
+                    this=exp.Literal.string(_ROW_DELETION_NAME), value=exp.Tuple(expressions=[column, interval])
+                )
+
+            if self._match_text_seq("TTL"):  # type: ignore[no-untyped-call]  # PostgreSQL-dialect style, keep for compatibility
+                self._match_text_seq("INTERVAL")  # type: ignore[no-untyped-call]
+                interval = cast("exp.Expression", self._parse_expression())
+                self._match_text_seq("ON")  # type: ignore[no-untyped-call]
+                column = cast("exp.Expression", self._parse_id_var())
+
+                return exp.Property(this=exp.Literal.string("TTL"), value=exp.Tuple(expressions=[interval, column]))
+
+            return cast("exp.Expression", super()._parse_property())
+
+    class Generator(BigQuery.Generator):
+        """Generate Spanner-specific DDL syntax."""
+
+        def table_sql(self, expression: exp.Table, sep: str = " ") -> str:
+            """Render INTERLEAVE clause when present on a table expression."""
+            sql = super().table_sql(expression, sep=sep)
+
+            parent = expression.args.get("interleave_parent")
+            if parent:
+                sql = f"{sql}\nINTERLEAVE IN PARENT {self.sql(parent)}"
+                on_delete = expression.args.get("interleave_on_delete")
+                if on_delete:
+                    sql = f"{sql} ON DELETE {on_delete}"
+
+            return sql
+
+        def property_sql(self, expression: exp.Property) -> str:
+            """Render row deletion policy or TTL."""
+            if getattr(expression.this, "name", "").upper() == _ROW_DELETION_NAME:
+                values = cast("Any", expression.args.get("value"))
+                if values and getattr(values, "expressions", None) and len(values.expressions) >= _TTL_MIN_COMPONENTS:
+                    column = self.sql(values.expressions[0])
+                    interval_sql = self.sql(values.expressions[1])
+                    if not interval_sql.upper().startswith("INTERVAL"):
+                        interval_sql = f"INTERVAL {interval_sql}"
+                    return f"ROW DELETION POLICY (OLDER_THAN({column}, {interval_sql}))"
+
+            if getattr(expression.this, "name", "").upper() == "TTL":
+                values = cast("Any", expression.args.get("value"))
+                if values and getattr(values, "expressions", None) and len(values.expressions) >= _TTL_MIN_COMPONENTS:
+                    interval = self.sql(values.expressions[0])
+                    column = self.sql(values.expressions[1])
+                    return f"TTL INTERVAL {interval} ON {column}"
+
+            return super().property_sql(expression)