"""Spanner configuration."""
from typing import TYPE_CHECKING, Any, ClassVar, Literal, TypedDict, cast
from typing_extensions import NotRequired
from sqlspec.adapters.spanner._typing import SpannerConnection
from sqlspec.adapters.spanner.core import apply_driver_features, default_statement_config
from sqlspec.adapters.spanner.driver import SpannerSessionContext, SpannerSyncDriver
from sqlspec.config import SyncDatabaseConfig
from sqlspec.driver._sync import SyncPoolConnectionContext, SyncPoolSessionFactory
from sqlspec.exceptions import ImproperConfigurationError
from sqlspec.extensions.events import EventRuntimeHints
from sqlspec.utils.config_tools import normalize_connection_config
from sqlspec.utils.type_guards import supports_close
if TYPE_CHECKING:
from collections.abc import Callable
from logging import Logger
from types import TracebackType
from google.api_core.client_info import ClientInfo
from google.api_core.client_options import ClientOptions
from google.api_core.retry import Retry
from google.auth.credentials import Credentials
from google.cloud.spanner_admin_database_v1.types import DatabaseDialect, EncryptionConfig
from google.cloud.spanner_v1 import Client, DirectedReadOptions, ExecuteSqlRequest, RequestOptions
from google.cloud.spanner_v1.database import Database
from google.cloud.spanner_v1.pool import AbstractSessionPool
from google.cloud.spanner_v1.transaction import DefaultTransactionOptions
from sqlspec.config import ExtensionConfigs
from sqlspec.core import StatementConfig
from sqlspec.observability import ObservabilityConfig
__all__ = ("SpannerConnectionParams", "SpannerDriverFeatures", "SpannerPoolParams", "SpannerSyncConfig")
_DEFAULT_SESSION_TRANSACTION: bool = True
"""Default ``transaction`` flag for ``provide_session`` / ``provide_connection``.
``True`` yields a write-capable :class:`Transaction` context matching every
other sqlspec adapter. Read-only :class:`Snapshot` contexts are available via
:meth:`SpannerSyncConfig.provide_read_session`. Pulled into a module-level
constant so an eventual ``SpannerAsyncConfig`` can import the same default."""
_CLIENT_CONFIG_FIELDS = frozenset({
"project",
"credentials",
"client_info",
"client_options",
"query_options",
"route_to_leader_enabled",
"directed_read_options",
"observability_options",
"default_transaction_options",
"experimental_host",
"disable_builtin_metrics",
"client_context",
"use_plain_text",
"ca_certificate",
"client_certificate",
"client_key",
"instance_type",
})
_INSTANCE_CONFIG_FIELDS = frozenset({"configuration_name", "display_name", "node_count", "processing_units"})
_DATABASE_CONFIG_FIELDS = frozenset({
"ddl_statements",
"logger",
"encryption_config",
"database_dialect",
"database_role",
"enable_drop_protection",
"enable_interceptors_in_tests",
"proto_descriptors",
})
class SpannerConnectionParams(TypedDict):
"""Spanner connection parameters."""
project: "NotRequired[str]"
credentials: "NotRequired[Credentials]"
client_info: "NotRequired[ClientInfo]"
client_options: "NotRequired[ClientOptions | dict[str, Any]]"
query_options: "NotRequired[ExecuteSqlRequest.QueryOptions]"
route_to_leader_enabled: "NotRequired[bool]"
directed_read_options: "NotRequired[DirectedReadOptions]"
observability_options: "NotRequired[Any]"
default_transaction_options: "NotRequired[DefaultTransactionOptions]"
experimental_host: "NotRequired[str]"
disable_builtin_metrics: "NotRequired[bool]"
client_context: "NotRequired[dict[str, str]]"
use_plain_text: "NotRequired[bool]"
ca_certificate: "NotRequired[str]"
client_certificate: "NotRequired[str]"
client_key: "NotRequired[str]"
instance_type: "NotRequired[str]"
instance_id: "NotRequired[str]"
configuration_name: "NotRequired[str]"
display_name: "NotRequired[str]"
node_count: "NotRequired[int]"
processing_units: "NotRequired[int]"
instance_labels: "NotRequired[dict[str, str]]"
database_id: "NotRequired[str]"
ddl_statements: "NotRequired[tuple[str, ...] | list[str]]"
logger: "NotRequired[Logger]"
encryption_config: "NotRequired[EncryptionConfig | dict[str, Any]]"
database_dialect: "NotRequired[DatabaseDialect]"
database_role: "NotRequired[str]"
enable_drop_protection: "NotRequired[bool]"
enable_interceptors_in_tests: "NotRequired[bool]"
proto_descriptors: "NotRequired[bytes]"
extra: "NotRequired[dict[str, Any]]"
class SpannerPoolParams(SpannerConnectionParams):
"""Session pool configuration."""
pool_type: "NotRequired[type[AbstractSessionPool]]"
size: "NotRequired[int]"
target_size: "NotRequired[int]"
max_sessions: "NotRequired[int]"
default_timeout: "NotRequired[int | float]"
session_labels: "NotRequired[dict[str, str]]"
labels: "NotRequired[dict[str, str]]"
ping_interval: "NotRequired[int]"
max_age_minutes: "NotRequired[int]"
class SpannerDriverFeatures(TypedDict):
"""Driver feature flags for Spanner.
Attributes:
enable_uuid_conversion: Enable automatic UUID string conversion.
json_serializer: Custom JSON serializer for parameter conversion.
json_deserializer: Custom JSON deserializer for result conversion.
retry: Per-request retry policy passed to execute_sql(), execute_update(), and batch_update().
timeout: Per-request timeout in seconds passed to execute_sql(), execute_update(), and batch_update().
request_options: Default RequestOptions forwarded to execute_sql(), execute_update(),
and batch_update(). Per-call overrides are available through normal
driver execution methods.
directed_read_options: Default DirectedReadOptions forwarded to execute_sql().
session_labels: Deprecated compatibility alias for pool session labels.
Prefer ``connection_config["session_labels"]``.
enable_events: Enable database event channel support.
Defaults to True when extension_config["events"] is configured.
events_backend: Backend type for event handling.
Spanner only supports "poll_queue" (no native pub/sub).
enable_batch_write_api: Route load_from_arrow through the Spanner Batch Write API
(Database.mutation_groups().batch_write()) for high-throughput, independently
committed mutation groups instead of a single in-transaction insert_or_update.
Defaults to False.
"""
enable_uuid_conversion: "NotRequired[bool]"
json_serializer: "NotRequired[Callable[[Any], str]]"
json_deserializer: "NotRequired[Callable[[str], Any]]"
retry: "NotRequired[Retry | None]"
timeout: "NotRequired[float | None]"
request_options: "NotRequired[RequestOptions | dict[str, Any] | None]"
directed_read_options: "NotRequired[DirectedReadOptions | None]"
session_labels: "NotRequired[dict[str, str]]"
enable_events: "NotRequired[bool]"
events_backend: "NotRequired[Literal['poll_queue']]"
enable_batch_write_api: "NotRequired[bool]"
class SpannerConnectionContext(SyncPoolConnectionContext):
"""Context manager for Spanner connections."""
__slots__ = ("_connection", "_session", "_transaction")
def __init__(self, config: "SpannerSyncConfig", transaction: bool = False) -> None:
super().__init__(config)
self._transaction = transaction
self._connection: SpannerConnection | None = None
self._session: Any = None
def __enter__(self) -> SpannerConnection:
database = self._config.get_database()
if self._transaction:
self._session = cast("Any", database).session()
self._session.create()
try:
txn = self._session.transaction()
txn.__enter__()
self._connection = cast("SpannerConnection", txn)
except Exception:
self._session.delete()
raise
else:
return self._connection
else:
self._session = cast("Any", database).snapshot(multi_use=True)
self._connection = cast("SpannerConnection", self._session.__enter__())
return self._connection
def __exit__(
self, exc_type: "type[BaseException] | None", exc_val: "BaseException | None", exc_tb: "TracebackType | None"
) -> bool | None:
if self._transaction and self._connection:
txn = cast("Any", self._connection)
try:
if exc_type is None:
try:
txn_id = txn._transaction_id
except AttributeError:
txn_id = None
mutations = cast("list[Any] | None", getattr(txn, "_mutations", None))
try:
committed = txn.committed
except AttributeError:
committed = None
if committed is None and (txn_id is not None or bool(mutations)):
txn.commit()
else:
try:
rollback_txn_id = txn._transaction_id
except AttributeError:
rollback_txn_id = None
if rollback_txn_id is not None:
txn.rollback()
finally:
if self._session:
self._session.delete()
elif self._session:
self._session.__exit__(exc_type, exc_val, exc_tb)
self._connection = None
self._session = None
return None
class _SpannerSessionConnectionHandler(SyncPoolSessionFactory):
__slots__ = ("_connection_ctx",)
def __init__(self, config: "SpannerSyncConfig", connection_ctx: "SpannerConnectionContext") -> None:
super().__init__(config)
self._connection_ctx = connection_ctx
def acquire_connection(self) -> "SpannerConnection":
return self._connection_ctx.__enter__()
def release_connection(self, _conn: "SpannerConnection", **kwargs: Any) -> None:
self._connection_ctx.__exit__(kwargs.get("exc_type"), kwargs.get("exc_val"), kwargs.get("exc_tb"))
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
_connection_context_class: "ClassVar[type[SpannerConnectionContext]]" = SpannerConnectionContext
_session_factory_class: "ClassVar[type[_SpannerSessionConnectionHandler]]" = _SpannerSessionConnectionHandler
_session_context_class: "ClassVar[type[SpannerSessionContext]]" = SpannerSessionContext
_default_statement_config = default_statement_config
def __init__(
self,
*,
connection_config: "SpannerPoolParams | dict[str, Any] | None" = None,
connection_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,
**kwargs: Any,
) -> None:
self.connection_config = normalize_connection_config(connection_config)
if "min_sessions" in self.connection_config:
msg = "Spanner session pools do not support 'min_sessions'; use 'size' or 'target_size'."
raise ImproperConfigurationError(msg)
raw_driver_features: dict[str, Any] = dict(driver_features) if driver_features else {}
legacy_session_labels = raw_driver_features.pop("session_labels", None)
if (
legacy_session_labels is not None
and "session_labels" not in self.connection_config
and "labels" not in self.connection_config
):
self.connection_config["session_labels"] = legacy_session_labels
from google.cloud.spanner_v1.pool import FixedSizePool
self.connection_config.setdefault("size", self.connection_config.pop("max_sessions", 10))
self.connection_config.setdefault("pool_type", FixedSizePool)
statement_config = statement_config or default_statement_config
statement_config, driver_features = apply_driver_features(statement_config, raw_driver_features)
super().__init__(
connection_config=self.connection_config,
connection_instance=connection_instance,
migration_config=migration_config,
statement_config=statement_config,
driver_features=driver_features,
bind_key=bind_key,
extension_config=extension_config,
observability_config=observability_config,
**kwargs,
)
self._client: Client | None = None
self._database: Database | None = None