sqlspec 0.36.0__cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl

sqlspec/adapters/cockroach_asyncpg/adk/store.py

@@ -0,0 +1,410 @@
+"""CockroachDB ADK store for Google Agent Development Kit session/event storage (asyncpg)."""
+
+from typing import TYPE_CHECKING, Any, cast
+
+from sqlspec.extensions.adk import BaseAsyncADKStore, EventRecord, SessionRecord
+from sqlspec.extensions.adk.memory.store import BaseAsyncADKMemoryStore
+from sqlspec.utils.logging import get_logger
+
+if TYPE_CHECKING:
+    from sqlspec.adapters.cockroach_asyncpg.config import CockroachAsyncpgConfig
+    from sqlspec.extensions.adk import MemoryRecord
+
+
+__all__ = ("CockroachAsyncpgADKMemoryStore", "CockroachAsyncpgADKStore")
+
+logger = get_logger("sqlspec.adapters.cockroach_asyncpg.adk.store")
+
+
+class CockroachAsyncpgADKStore(BaseAsyncADKStore["CockroachAsyncpgConfig"]):
+    """CockroachDB ADK store using asyncpg driver."""
+
+    __slots__ = ()
+
+    def __init__(self, config: "CockroachAsyncpgConfig") -> None:
+        super().__init__(config)
+
+    async def _get_create_sessions_table_sql(self) -> str:
+        owner_id_line = ""
+        if self._owner_id_column_ddl:
+            owner_id_line = f",\n            {self._owner_id_column_ddl}"
+
+        return f"""
+        CREATE TABLE IF NOT EXISTS {self._session_table} (
+            id VARCHAR(128) PRIMARY KEY,
+            app_name VARCHAR(128) NOT NULL,
+            user_id VARCHAR(128) NOT NULL{owner_id_line},
+            state JSONB NOT NULL DEFAULT '{{}}'::jsonb,
+            create_time TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP,
+            update_time TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP
+        );
+
+        CREATE INDEX IF NOT EXISTS idx_{self._session_table}_app_user
+            ON {self._session_table}(app_name, user_id);
+
+        CREATE INDEX IF NOT EXISTS idx_{self._session_table}_update_time
+            ON {self._session_table}(update_time DESC);
+        """
+
+    async def _get_create_events_table_sql(self) -> str:
+        return f"""
+        CREATE TABLE IF NOT EXISTS {self._events_table} (
+            id VARCHAR(128) PRIMARY KEY,
+            session_id VARCHAR(128) NOT NULL,
+            app_name VARCHAR(128) NOT NULL,
+            user_id VARCHAR(128) NOT NULL,
+            invocation_id VARCHAR(256) NOT NULL,
+            author VARCHAR(256) NOT NULL,
+            actions BYTEA NOT NULL,
+            long_running_tool_ids_json JSONB,
+            branch VARCHAR(256),
+            timestamp TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP,
+            content JSONB,
+            grounding_metadata JSONB,
+            custom_metadata JSONB,
+            partial BOOLEAN,
+            turn_complete BOOLEAN,
+            interrupted BOOLEAN,
+            error_code VARCHAR(256),
+            error_message VARCHAR(1024),
+            FOREIGN KEY (session_id) REFERENCES {self._session_table}(id) ON DELETE CASCADE
+        );
+
+        CREATE INDEX IF NOT EXISTS idx_{self._events_table}_session
+            ON {self._events_table}(session_id, timestamp ASC);
+        """
+
+    def _get_drop_tables_sql(self) -> "list[str]":
+        return [f"DROP TABLE IF EXISTS {self._events_table}", f"DROP TABLE IF EXISTS {self._session_table}"]
+
+    async def create_tables(self) -> None:
+        async with self._config.provide_session() as driver:
+            await driver.execute_script(await self._get_create_sessions_table_sql())
+            await driver.execute_script(await self._get_create_events_table_sql())
+
+    async def create_session(
+        self, session_id: str, app_name: str, user_id: str, state: "dict[str, Any]", owner_id: "Any | None" = None
+    ) -> SessionRecord:
+        params: tuple[Any, ...]
+        if self._owner_id_column_name:
+            sql = f"""
+            INSERT INTO {self._session_table} (id, app_name, user_id, {self._owner_id_column_name}, state, create_time, update_time)
+            VALUES ($1, $2, $3, $4, $5, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)
+            """
+            params = (session_id, app_name, user_id, owner_id, state)
+        else:
+            sql = f"""
+            INSERT INTO {self._session_table} (id, app_name, user_id, state, create_time, update_time)
+            VALUES ($1, $2, $3, $4, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)
+            """
+            params = (session_id, app_name, user_id, state)
+
+        async with self._config.provide_connection() as conn:
+            await conn.execute(sql, *params)
+
+        result = await self.get_session(session_id)
+        if result is None:
+            msg = "Session creation failed"
+            raise RuntimeError(msg)
+        return result
+
+    async def get_session(self, session_id: str) -> "SessionRecord | None":
+        sql = f"""
+        SELECT id, app_name, user_id, state, create_time, update_time
+        FROM {self._session_table}
+        WHERE id = $1
+        """
+
+        async with self._config.provide_connection() as conn:
+            row = await conn.fetchrow(sql, session_id)
+            if row is None:
+                return None
+
+            return SessionRecord(
+                id=row["id"],
+                app_name=row["app_name"],
+                user_id=row["user_id"],
+                state=row["state"],
+                create_time=row["create_time"],
+                update_time=row["update_time"],
+            )
+
+    async def update_session_state(self, session_id: str, state: "dict[str, Any]") -> None:
+        sql = f"""
+        UPDATE {self._session_table}
+        SET state = $1, update_time = CURRENT_TIMESTAMP
+        WHERE id = $2
+        """
+
+        async with self._config.provide_connection() as conn:
+            await conn.execute(sql, state, session_id)
+
+    async def delete_session(self, session_id: str) -> None:
+        sql = f"DELETE FROM {self._session_table} WHERE id = $1"
+
+        async with self._config.provide_connection() as conn:
+            await conn.execute(sql, session_id)
+
+    async def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[SessionRecord]":
+        if user_id is None:
+            sql = f"""
+            SELECT id, app_name, user_id, state, create_time, update_time
+            FROM {self._session_table}
+            WHERE app_name = $1
+            ORDER BY update_time DESC
+            """
+            params: tuple[Any, ...] = (app_name,)
+        else:
+            sql = f"""
+            SELECT id, app_name, user_id, state, create_time, update_time
+            FROM {self._session_table}
+            WHERE app_name = $1 AND user_id = $2
+            ORDER BY update_time DESC
+            """
+            params = (app_name, user_id)
+
+        async with self._config.provide_connection() as conn:
+            rows = await conn.fetch(sql, *params)
+
+        return [
+            SessionRecord(
+                id=row["id"],
+                app_name=row["app_name"],
+                user_id=row["user_id"],
+                state=row["state"],
+                create_time=row["create_time"],
+                update_time=row["update_time"],
+            )
+            for row in rows
+        ]
+
+    async def append_event(self, event_record: EventRecord) -> None:
+        sql = f"""
+        INSERT INTO {self._events_table} (
+            id, session_id, app_name, user_id, invocation_id, author, actions,
+            long_running_tool_ids_json, branch, timestamp, content,
+            grounding_metadata, custom_metadata, partial, turn_complete,
+            interrupted, error_code, error_message
+        ) VALUES (
+            $1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18
+        )
+        """
+
+        async with self._config.provide_connection() as conn:
+            await conn.execute(
+                sql,
+                event_record["id"],
+                event_record["session_id"],
+                event_record["app_name"],
+                event_record["user_id"],
+                event_record["invocation_id"],
+                event_record["author"],
+                event_record["actions"],
+                event_record.get("long_running_tool_ids_json"),
+                event_record.get("branch"),
+                event_record["timestamp"],
+                event_record.get("content"),
+                event_record.get("grounding_metadata"),
+                event_record.get("custom_metadata"),
+                event_record.get("partial"),
+                event_record.get("turn_complete"),
+                event_record.get("interrupted"),
+                event_record.get("error_code"),
+                event_record.get("error_message"),
+            )
+
+    async def list_events(self, session_id: str) -> "list[EventRecord]":
+        sql = f"""
+        SELECT id, session_id, app_name, user_id, invocation_id, author, actions,
+               long_running_tool_ids_json, branch, timestamp, content,
+               grounding_metadata, custom_metadata, partial, turn_complete,
+               interrupted, error_code, error_message
+        FROM {self._events_table}
+        WHERE session_id = $1
+        ORDER BY timestamp ASC
+        """
+
+        async with self._config.provide_connection() as conn:
+            rows = await conn.fetch(sql, session_id)
+
+        return [
+            EventRecord(
+                id=row["id"],
+                session_id=row["session_id"],
+                app_name=row["app_name"],
+                user_id=row["user_id"],
+                invocation_id=row["invocation_id"],
+                author=row["author"],
+                actions=bytes(row["actions"]),
+                long_running_tool_ids_json=row["long_running_tool_ids_json"],
+                branch=row["branch"],
+                timestamp=row["timestamp"],
+                content=row["content"],
+                grounding_metadata=row["grounding_metadata"],
+                custom_metadata=row["custom_metadata"],
+                partial=row["partial"],
+                turn_complete=row["turn_complete"],
+                interrupted=row["interrupted"],
+                error_code=row["error_code"],
+                error_message=row["error_message"],
+            )
+            for row in rows
+        ]
+
+
+class CockroachAsyncpgADKMemoryStore(BaseAsyncADKMemoryStore["CockroachAsyncpgConfig"]):
+    """CockroachDB ADK memory store using asyncpg driver."""
+
+    __slots__ = ()
+
+    def __init__(self, config: "CockroachAsyncpgConfig") -> None:
+        super().__init__(config)
+
+    async def _get_create_memory_table_sql(self) -> str:
+        owner_id_line = ""
+        if self._owner_id_column_ddl:
+            owner_id_line = f",\n            {self._owner_id_column_ddl}"
+
+        return f"""
+        CREATE TABLE IF NOT EXISTS {self._memory_table} (
+            id VARCHAR(128) PRIMARY KEY,
+            session_id VARCHAR(128) NOT NULL,
+            app_name VARCHAR(128) NOT NULL,
+            user_id VARCHAR(128) NOT NULL,
+            event_id VARCHAR(128) NOT NULL UNIQUE,
+            author VARCHAR(256){owner_id_line},
+            timestamp TIMESTAMPTZ NOT NULL,
+            content_json JSONB NOT NULL,
+            content_text TEXT NOT NULL,
+            metadata_json JSONB,
+            inserted_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP
+        );
+
+        CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_app_user_time
+            ON {self._memory_table}(app_name, user_id, timestamp DESC);
+
+        CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_session
+            ON {self._memory_table}(session_id);
+        """
+
+    def _get_drop_memory_table_sql(self) -> "list[str]":
+        return [f"DROP TABLE IF EXISTS {self._memory_table}"]
+
+    async def create_tables(self) -> None:
+        if not self._enabled:
+            return
+
+        async with self._config.provide_session() as driver:
+            await driver.execute_script(await self._get_create_memory_table_sql())
+
+    async def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object | None" = None) -> int:
+        if not self._enabled:
+            msg = "Memory store is disabled"
+            raise RuntimeError(msg)
+
+        if not entries:
+            return 0
+
+        inserted_count = 0
+        async with self._config.provide_connection() as conn:
+            for entry in entries:
+                if self._owner_id_column_name:
+                    sql = f"""
+                    INSERT INTO {self._memory_table}
+                    (id, session_id, app_name, user_id, event_id, author,
+                     {self._owner_id_column_name}, timestamp, content_json,
+                     content_text, metadata_json, inserted_at)
+                    VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)
+                    ON CONFLICT (event_id) DO NOTHING
+                    """
+                    result = await conn.execute(
+                        sql,
+                        entry["id"],
+                        entry["session_id"],
+                        entry["app_name"],
+                        entry["user_id"],
+                        entry["event_id"],
+                        entry["author"],
+                        owner_id,
+                        entry["timestamp"],
+                        entry["content_json"],
+                        entry["content_text"],
+                        entry["metadata_json"],
+                        entry["inserted_at"],
+                    )
+                else:
+                    sql = f"""
+                    INSERT INTO {self._memory_table}
+                    (id, session_id, app_name, user_id, event_id, author,
+                     timestamp, content_json, content_text, metadata_json, inserted_at)
+                    VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)
+                    ON CONFLICT (event_id) DO NOTHING
+                    """
+                    result = await conn.execute(
+                        sql,
+                        entry["id"],
+                        entry["session_id"],
+                        entry["app_name"],
+                        entry["user_id"],
+                        entry["event_id"],
+                        entry["author"],
+                        entry["timestamp"],
+                        entry["content_json"],
+                        entry["content_text"],
+                        entry["metadata_json"],
+                        entry["inserted_at"],
+                    )
+
+                if result and result.split()[-1].isdigit() and int(result.split()[-1]) > 0:
+                    inserted_count += 1
+
+        return inserted_count
+
+    async def search_entries(
+        self, query: str, app_name: str, user_id: str, limit: "int | None" = None
+    ) -> "list[MemoryRecord]":
+        if not self._enabled:
+            msg = "Memory store is disabled"
+            raise RuntimeError(msg)
+
+        if not query:
+            return []
+
+        effective_limit = limit if limit is not None else self._max_results
+        if self._use_fts:
+            logger.debug("CockroachDB full-text search not supported; using simple search")
+
+        sql = f"""
+        SELECT * FROM {self._memory_table}
+        WHERE app_name = $1 AND user_id = $2 AND content_text ILIKE $3
+        ORDER BY timestamp DESC
+        LIMIT $4
+        """
+
+        async with self._config.provide_connection() as conn:
+            rows = await conn.fetch(sql, app_name, user_id, f"%{query}%", effective_limit)
+
+        return [cast("MemoryRecord", dict(row)) for row in rows]
+
+    async def delete_entries_by_session(self, session_id: str) -> int:
+        if not self._enabled:
+            msg = "Memory store is disabled"
+            raise RuntimeError(msg)
+
+        sql = f"DELETE FROM {self._memory_table} WHERE session_id = $1"
+        async with self._config.provide_connection() as conn:
+            result = await conn.execute(sql, session_id)
+            return int(result.split()[-1]) if result else 0
+
+    async def delete_entries_older_than(self, days: int) -> int:
+        if not self._enabled:
+            msg = "Memory store is disabled"
+            raise RuntimeError(msg)
+
+        sql = f"""
+        DELETE FROM {self._memory_table}
+        WHERE inserted_at < (CURRENT_TIMESTAMP - INTERVAL $1 DAY)
+        """
+        async with self._config.provide_connection() as conn:
+            result = await conn.execute(sql, days)
+            return int(result.split()[-1]) if result else 0

sqlspec/adapters/cockroach_asyncpg/config.py

@@ -0,0 +1,238 @@
+"""CockroachDB AsyncPG configuration."""
+
+from typing import TYPE_CHECKING, Any, ClassVar, TypedDict, cast
+
+from asyncpg import create_pool as asyncpg_create_pool
+from typing_extensions import NotRequired
+
+from sqlspec.adapters.asyncpg.core import apply_driver_features, build_connection_config, default_statement_config
+from sqlspec.adapters.cockroach_asyncpg._typing import (
+    CockroachAsyncpgConnection,
+    CockroachAsyncpgPool,
+    CockroachAsyncpgSessionContext,
+)
+from sqlspec.adapters.cockroach_asyncpg.core import CockroachAsyncpgRetryConfig
+from sqlspec.adapters.cockroach_asyncpg.driver import CockroachAsyncpgDriver, CockroachAsyncpgExceptionHandler
+from sqlspec.config import AsyncDatabaseConfig, ExtensionConfigs
+from sqlspec.extensions.events import EventRuntimeHints
+from sqlspec.utils.config_tools import normalize_connection_config, reject_pool_aliases
+
+if TYPE_CHECKING:
+    from collections.abc import Awaitable, Callable
+
+    from sqlspec.core import StatementConfig
+    from sqlspec.observability import ObservabilityConfig
+
+__all__ = (
+    "CockroachAsyncpgConfig",
+    "CockroachAsyncpgConnectionConfig",
+    "CockroachAsyncpgDriverFeatures",
+    "CockroachAsyncpgPoolConfig",
+)
+
+
+class CockroachAsyncpgConnectionConfig(TypedDict):
+    """AsyncPG connection parameters for CockroachDB."""
+
+    dsn: NotRequired[str]
+    host: NotRequired[str]
+    port: NotRequired[int]
+    user: NotRequired[str]
+    password: NotRequired[str]
+    database: NotRequired[str]
+    ssl: NotRequired[Any]
+    passfile: NotRequired[str]
+    direct_tls: NotRequired[bool]
+    connect_timeout: NotRequired[float]
+    command_timeout: NotRequired[float]
+    statement_cache_size: NotRequired[int]
+    max_cached_statement_lifetime: NotRequired[int]
+    max_cacheable_statement_size: NotRequired[int]
+    server_settings: NotRequired["dict[str, str]"]
+
+
+class CockroachAsyncpgPoolConfig(CockroachAsyncpgConnectionConfig):
+    """AsyncPG pool parameters for CockroachDB."""
+
+    min_size: NotRequired[int]
+    max_size: NotRequired[int]
+    max_queries: NotRequired[int]
+    max_inactive_connection_lifetime: NotRequired[float]
+    setup: NotRequired["Callable[[CockroachAsyncpgConnection], Awaitable[None]]"]
+    init: NotRequired["Callable[[CockroachAsyncpgConnection], Awaitable[None]]"]
+    extra: NotRequired["dict[str, Any]"]
+
+
+class CockroachAsyncpgDriverFeatures(TypedDict):
+    """Driver feature flags for CockroachDB AsyncPG adapter."""
+
+    enable_auto_retry: NotRequired[bool]
+    max_retries: NotRequired[int]
+    retry_delay_base_ms: NotRequired[float]
+    retry_delay_max_ms: NotRequired[float]
+    enable_retry_logging: NotRequired[bool]
+    enable_follower_reads: NotRequired[bool]
+    default_staleness: NotRequired[str]
+    json_serializer: NotRequired["Callable[[Any], str]"]
+    json_deserializer: NotRequired["Callable[[str], Any]"]
+    enable_json_codecs: NotRequired[bool]
+    enable_pgvector: NotRequired[bool]
+    enable_events: NotRequired[bool]
+    events_backend: NotRequired[str]
+
+
+class _CockroachAsyncpgSessionFactory:
+    __slots__ = ("_config", "_ctx")
+
+    def __init__(self, config: "CockroachAsyncpgConfig") -> None:
+        self._config = config
+        self._ctx: Any | None = None
+
+    async def acquire_connection(self) -> "CockroachAsyncpgConnection":
+        pool = self._config.connection_instance
+        if pool is None:
+            pool = await self._config.create_pool()
+            self._config.connection_instance = pool
+        ctx = pool.acquire()
+        self._ctx = ctx
+        return cast("CockroachAsyncpgConnection", await ctx.__aenter__())
+
+    async def release_connection(self, _conn: "CockroachAsyncpgConnection") -> None:
+        if self._ctx is not None:
+            await self._ctx.__aexit__(None, None, None)
+            self._ctx = None
+
+
+class CockroachAsyncpgConnectionContext:
+    """Async context manager for CockroachDB AsyncPG connections."""
+
+    __slots__ = ("_config", "_connection")
+
+    def __init__(self, config: "CockroachAsyncpgConfig") -> None:
+        self._config = config
+        self._connection: CockroachAsyncpgConnection | None = None
+
+    async def __aenter__(self) -> "CockroachAsyncpgConnection":
+        pool = self._config.connection_instance
+        if pool is None:
+            pool = await self._config.create_pool()
+            self._config.connection_instance = pool
+        self._connection = await pool.acquire()
+        return self._connection
+
+    async def __aexit__(
+        self, exc_type: "type[BaseException] | None", exc_val: "BaseException | None", exc_tb: Any
+    ) -> bool | None:
+        if self._connection is not None:
+            if self._config.connection_instance:
+                await self._config.connection_instance.release(self._connection)  # type: ignore[arg-type]
+            self._connection = None
+        return None
+
+
+class CockroachAsyncpgConfig(
+    AsyncDatabaseConfig[CockroachAsyncpgConnection, CockroachAsyncpgPool, CockroachAsyncpgDriver]
+):
+    """Configuration for CockroachDB using AsyncPG."""
+
+    driver_type: "ClassVar[type[CockroachAsyncpgDriver]]" = CockroachAsyncpgDriver
+    connection_type: "ClassVar[type[CockroachAsyncpgConnection]]" = CockroachAsyncpgConnection  # type: ignore[assignment]
+    supports_transactional_ddl: "ClassVar[bool]" = True
+    supports_native_arrow_export: "ClassVar[bool]" = True
+    supports_native_arrow_import: "ClassVar[bool]" = True
+    supports_native_parquet_export: "ClassVar[bool]" = True
+    supports_native_parquet_import: "ClassVar[bool]" = True
+
+    def __init__(
+        self,
+        *,
+        connection_config: "CockroachAsyncpgPoolConfig | dict[str, Any] | None" = None,
+        connection_instance: "CockroachAsyncpgPool | None" = None,
+        migration_config: "dict[str, Any] | None" = None,
+        statement_config: "StatementConfig | None" = None,
+        driver_features: "CockroachAsyncpgDriverFeatures | dict[str, Any] | None" = None,
+        bind_key: "str | None" = None,
+        extension_config: "ExtensionConfigs | None" = None,
+        observability_config: "ObservabilityConfig | None" = None,
+        **kwargs: Any,
+    ) -> None:
+        reject_pool_aliases(kwargs)
+
+        connection_config = normalize_connection_config(connection_config)
+        statement_config = statement_config or default_statement_config
+        statement_config, driver_features = apply_driver_features(statement_config, driver_features)
+
+        driver_features.setdefault("enable_auto_retry", True)
+        _ = CockroachAsyncpgRetryConfig.from_features(driver_features)
+
+        super().__init__(
+            connection_config=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,
+        )
+
+    async def _create_pool(self) -> "CockroachAsyncpgPool":
+        config = build_connection_config(self.connection_config)
+        return await asyncpg_create_pool(**config)
+
+    async def _close_pool(self) -> None:
+        if not self.connection_instance:
+            return
+        await self.connection_instance.close()
+        self.connection_instance = None
+
+    async def create_connection(self) -> "CockroachAsyncpgConnection":
+        if self.connection_instance is None:
+            self.connection_instance = await self.create_pool()
+        return cast("CockroachAsyncpgConnection", await self.connection_instance.acquire())
+
+    def provide_connection(self, *args: Any, **kwargs: Any) -> "CockroachAsyncpgConnectionContext":
+        return CockroachAsyncpgConnectionContext(self)
+
+    def provide_session(
+        self,
+        *_args: Any,
+        statement_config: "StatementConfig | None" = None,
+        follower_reads: bool | None = None,
+        staleness: str | None = None,
+        **_kwargs: Any,
+    ) -> "CockroachAsyncpgSessionContext":
+        factory = _CockroachAsyncpgSessionFactory(self)
+        driver_features = dict(self.driver_features)
+        if follower_reads is not None:
+            driver_features["enable_follower_reads"] = follower_reads
+        if staleness is not None:
+            driver_features["default_staleness"] = staleness
+
+        return CockroachAsyncpgSessionContext(
+            acquire_connection=factory.acquire_connection,
+            release_connection=factory.release_connection,
+            statement_config=statement_config or self.statement_config or default_statement_config,
+            driver_features=driver_features,
+            prepare_driver=self._prepare_driver,
+        )
+
+    async def provide_pool(self, *args: Any, **kwargs: Any) -> "CockroachAsyncpgPool":
+        if not self.connection_instance:
+            self.connection_instance = await self.create_pool()
+        return self.connection_instance
+
+    def get_signature_namespace(self) -> "dict[str, Any]":
+        namespace = super().get_signature_namespace()
+        namespace.update({
+            "CockroachAsyncpgConnectionConfig": CockroachAsyncpgConnectionConfig,
+            "CockroachAsyncpgPoolConfig": CockroachAsyncpgPoolConfig,
+            "CockroachAsyncpgDriver": CockroachAsyncpgDriver,
+            "CockroachAsyncpgExceptionHandler": CockroachAsyncpgExceptionHandler,
+            "CockroachAsyncpgSessionContext": CockroachAsyncpgSessionContext,
+        })
+        return namespace
+
+    def get_event_runtime_hints(self) -> "EventRuntimeHints":
+        return EventRuntimeHints(poll_interval=0.5, select_for_update=True, skip_locked=True)