"""CockroachDB psycopg driver implementation."""
import asyncio
import contextlib
import time
from typing import TYPE_CHECKING, Any, TypeVar, cast
import psycopg
from sqlspec.adapters.cockroach_psycopg._typing import (
CockroachAsyncConnection,
CockroachPsycopgAsyncSessionContext,
CockroachPsycopgSyncSessionContext,
CockroachSyncConnection,
)
from sqlspec.adapters.cockroach_psycopg.core import (
CockroachPsycopgRetryConfig,
build_statement_config,
calculate_backoff_seconds,
driver_profile,
is_retryable_error,
)
from sqlspec.adapters.cockroach_psycopg.data_dictionary import (
CockroachPsycopgAsyncDataDictionary,
CockroachPsycopgSyncDataDictionary,
)
from sqlspec.adapters.psycopg.core import create_mapped_exception
from sqlspec.adapters.psycopg.driver import PsycopgAsyncDriver, PsycopgSyncDriver
from sqlspec.core import SQL, StatementConfig, get_cache_config, register_driver_profile
from sqlspec.driver import BaseAsyncExceptionHandler, BaseSyncExceptionHandler
from sqlspec.exceptions import SerializationConflictError, TransactionRetryError
from sqlspec.utils.logging import get_logger
from sqlspec.utils.type_guards import has_sqlstate
if TYPE_CHECKING:
from collections.abc import Awaitable, Callable
from sqlspec.adapters.cockroach_psycopg._typing import CockroachAsyncCursor, CockroachSyncCursor
from sqlspec.driver import ExecutionResult
__all__ = (
"CockroachPsycopgAsyncDriver",
"CockroachPsycopgAsyncExceptionHandler",
"CockroachPsycopgAsyncSessionContext",
"CockroachPsycopgSyncDriver",
"CockroachPsycopgSyncExceptionHandler",
"CockroachPsycopgSyncSessionContext",
)
logger = get_logger("sqlspec.adapters.cockroach_psycopg")
_T = TypeVar("_T")
class CockroachPsycopgSyncExceptionHandler(BaseSyncExceptionHandler):
"""Context manager for handling CockroachDB psycopg exceptions."""
__slots__ = ()
def _handle_exception(self, exc_type: "type[BaseException] | None", exc_val: "BaseException") -> bool:
if exc_type is None:
return False
if issubclass(exc_type, psycopg.Error):
if has_sqlstate(exc_val) and str(exc_val.sqlstate) == "40001":
self.pending_exception = SerializationConflictError(str(exc_val))
return True
self.pending_exception = create_mapped_exception(exc_val)
return True
return False
class CockroachPsycopgAsyncExceptionHandler(BaseAsyncExceptionHandler):
"""Async context manager for handling CockroachDB psycopg exceptions."""
__slots__ = ()
def _handle_exception(self, exc_type: "type[BaseException] | None", exc_val: "BaseException") -> bool:
if exc_type is None:
return False
if issubclass(exc_type, psycopg.Error):
if has_sqlstate(exc_val) and str(exc_val.sqlstate) == "40001":
self.pending_exception = SerializationConflictError(str(exc_val))
return True
self.pending_exception = create_mapped_exception(exc_val)
return True
return False
[docs]
class CockroachPsycopgSyncDriver(PsycopgSyncDriver):
"""CockroachDB sync driver using psycopg.crdb."""
__slots__ = ("_enable_retry", "_follower_staleness", "_retry_config")
dialect = "postgres"
[docs]
def __init__(
self,
connection: CockroachSyncConnection,
statement_config: "StatementConfig | None" = None,
driver_features: "dict[str, Any] | None" = None,
) -> None:
if statement_config is None:
statement_config = build_statement_config().replace(
enable_caching=get_cache_config().compiled_cache_enabled
)
driver_features = dict(driver_features) if driver_features else {}
super().__init__(connection=connection, statement_config=statement_config, driver_features=driver_features)
self._retry_config = CockroachPsycopgRetryConfig.from_features(self.driver_features)
self._enable_retry = bool(self.driver_features.get("enable_auto_retry", True))
self._follower_staleness = cast("str | None", self.driver_features.get("default_staleness"))
# Data dictionary is lazily initialized in property; use parent slot
self._data_dictionary = None
[docs]
def run_transaction_with_retry(self, operation: "Callable[[], _T]") -> _T:
"""Execute a full CockroachDB transaction callback with serialization retries."""
if not self._enable_retry or self._connection_in_transaction():
return operation()
last_error: BaseException | None = None
for attempt in range(self._retry_config.max_retries + 1):
try:
self.begin()
result = operation()
self.commit()
except Exception as exc:
last_error = exc
with contextlib.suppress(Exception):
self.rollback()
if not is_retryable_error(exc) or attempt >= self._retry_config.max_retries:
raise
else:
return result
delay = calculate_backoff_seconds(attempt, self._retry_config)
if self._retry_config.enable_logging:
logger.debug("CockroachDB retry %s/%s after %.3fs", attempt + 1, self._retry_config.max_retries, delay)
time.sleep(delay)
msg = "CockroachDB transaction retry limit exceeded"
raise TransactionRetryError(msg) from last_error
[docs]
def dispatch_execute(self, cursor: "CockroachSyncCursor", statement: SQL) -> "ExecutionResult":
return self._dispatch_execute_impl(cursor, statement)
[docs]
def dispatch_execute_many(self, cursor: "CockroachSyncCursor", statement: SQL) -> "ExecutionResult":
return self._dispatch_execute_many_impl(cursor, statement)
[docs]
def dispatch_execute_script(self, cursor: "CockroachSyncCursor", statement: SQL) -> "ExecutionResult":
return self._dispatch_execute_script_impl(cursor, statement)
[docs]
def handle_database_exceptions(self) -> "CockroachPsycopgSyncExceptionHandler": # type: ignore[override]
return CockroachPsycopgSyncExceptionHandler()
@property
def data_dictionary(self) -> "CockroachPsycopgSyncDataDictionary": # type: ignore[override]
if self._data_dictionary is None:
# Intentionally assign CockroachDB-specific data dictionary to parent slot
self._data_dictionary = CockroachPsycopgSyncDataDictionary() # type: ignore[assignment]
return cast("CockroachPsycopgSyncDataDictionary", self._data_dictionary)
def _apply_follower_reads(self, cursor: "CockroachSyncCursor") -> None:
if not self.driver_features.get("enable_follower_reads", False):
return
if not self._follower_staleness:
return
cursor.execute(cast("Any", f"SET TRANSACTION AS OF SYSTEM TIME {self._follower_staleness}"))
def _dispatch_execute_impl(self, cursor: "CockroachSyncCursor", statement: SQL) -> "ExecutionResult":
if statement.returns_rows():
self._apply_follower_reads(cursor)
return super().dispatch_execute(cursor, statement)
def _dispatch_execute_many_impl(self, cursor: "CockroachSyncCursor", statement: SQL) -> "ExecutionResult":
return PsycopgSyncDriver.dispatch_execute_many(self, cursor, statement)
def _dispatch_execute_script_impl(self, cursor: "CockroachSyncCursor", statement: SQL) -> "ExecutionResult":
return PsycopgSyncDriver.dispatch_execute_script(self, cursor, statement)
[docs]
class CockroachPsycopgAsyncDriver(PsycopgAsyncDriver):
"""CockroachDB async driver using psycopg.crdb."""
__slots__ = ("_enable_retry", "_follower_staleness", "_retry_config")
dialect = "postgres"
[docs]
def __init__(
self,
connection: CockroachAsyncConnection,
statement_config: "StatementConfig | None" = None,
driver_features: "dict[str, Any] | None" = None,
) -> None:
if statement_config is None:
statement_config = build_statement_config().replace(
enable_caching=get_cache_config().compiled_cache_enabled
)
driver_features = dict(driver_features) if driver_features else {}
super().__init__(connection=connection, statement_config=statement_config, driver_features=driver_features)
self._retry_config = CockroachPsycopgRetryConfig.from_features(self.driver_features)
self._enable_retry = bool(self.driver_features.get("enable_auto_retry", True))
self._follower_staleness = cast("str | None", self.driver_features.get("default_staleness"))
# Data dictionary is lazily initialized in property; use parent slot
self._data_dictionary = None
[docs]
async def run_transaction_with_retry(self, operation: "Callable[[], Awaitable[_T]]") -> _T:
"""Execute a full CockroachDB transaction callback with serialization retries."""
if not self._enable_retry or self._connection_in_transaction():
return await operation()
last_error: BaseException | None = None
for attempt in range(self._retry_config.max_retries + 1):
try:
await self.begin()
result = await operation()
await self.commit()
except Exception as exc:
last_error = exc
with contextlib.suppress(Exception):
await self.rollback()
if not is_retryable_error(exc) or attempt >= self._retry_config.max_retries:
raise
else:
return result
delay = calculate_backoff_seconds(attempt, self._retry_config)
if self._retry_config.enable_logging:
logger.debug("CockroachDB retry %s/%s after %.3fs", attempt + 1, self._retry_config.max_retries, delay)
await asyncio.sleep(delay)
msg = "CockroachDB transaction retry limit exceeded"
raise TransactionRetryError(msg) from last_error
[docs]
async def dispatch_execute(self, cursor: "CockroachAsyncCursor", statement: SQL) -> "ExecutionResult":
return await self._dispatch_execute_impl(cursor, statement)
[docs]
async def dispatch_execute_many(self, cursor: "CockroachAsyncCursor", statement: SQL) -> "ExecutionResult":
return await self._dispatch_execute_many_impl(cursor, statement)
[docs]
async def dispatch_execute_script(self, cursor: "CockroachAsyncCursor", statement: SQL) -> "ExecutionResult":
return await self._dispatch_execute_script_impl(cursor, statement)
[docs]
def handle_database_exceptions(self) -> "CockroachPsycopgAsyncExceptionHandler": # type: ignore[override]
return CockroachPsycopgAsyncExceptionHandler()
@property
def data_dictionary(self) -> "CockroachPsycopgAsyncDataDictionary": # type: ignore[override]
if self._data_dictionary is None:
# Intentionally assign CockroachDB-specific data dictionary to parent slot
self._data_dictionary = CockroachPsycopgAsyncDataDictionary() # type: ignore[assignment]
return cast("CockroachPsycopgAsyncDataDictionary", self._data_dictionary)
async def _apply_follower_reads(self, cursor: "CockroachAsyncCursor") -> None:
if not self.driver_features.get("enable_follower_reads", False):
return
if not self._follower_staleness:
return
await cursor.execute(cast("Any", f"SET TRANSACTION AS OF SYSTEM TIME {self._follower_staleness}"))
async def _dispatch_execute_impl(self, cursor: "CockroachAsyncCursor", statement: SQL) -> "ExecutionResult":
if statement.returns_rows():
await self._apply_follower_reads(cursor)
return await super().dispatch_execute(cursor, statement)
async def _dispatch_execute_many_impl(self, cursor: "CockroachAsyncCursor", statement: SQL) -> "ExecutionResult":
return await PsycopgAsyncDriver.dispatch_execute_many(self, cursor, statement)
async def _dispatch_execute_script_impl(self, cursor: "CockroachAsyncCursor", statement: SQL) -> "ExecutionResult":
return await PsycopgAsyncDriver.dispatch_execute_script(self, cursor, statement)
register_driver_profile("cockroach_psycopg", driver_profile)
