CockroachDB + AsyncPG#

CockroachDB adapter using asyncpg with automatic transaction retry logic.

Configuration#

class sqlspec.adapters.cockroach_asyncpg.CockroachAsyncpgConfig[source]#

Bases: AsyncDatabaseConfig[PoolConnectionProxy, Pool, CockroachAsyncpgDriver]

Configuration for CockroachDB using AsyncPG.

driver_type#: alias of CockroachAsyncpgDriver

connection_type#: alias of PoolConnectionProxy

__init__(*, connection_config=None, connection_instance=None, migration_config=None, statement_config=None, driver_features=None, bind_key=None, extension_config=None, observability_config=None, **kwargs)[source]#

async create_connection()[source]#

Open a standalone connection owned by the caller.

The connection carries the same connection settings and init hook the pool applies, consumes no pool slot, and must be closed by the caller.

Return type:: PoolConnectionProxy
Returns:: A CockroachDB asyncpg connection.

provide_session(*_args, statement_config=None, follower_reads=None, staleness=None, **_kwargs)[source]#

Provide a database session context manager.

Return type:: CockroachAsyncpgSessionContext

async provide_pool(*args, **kwargs)[source]#

Provide pool instance.

Return type:: Pool

get_signature_namespace()[source]#

Get the signature namespace for this database configuration.

Returns a dictionary of type names to objects (classes, functions, or other callables) that should be registered with Litestar's signature namespace to prevent serialization attempts on database-specific structures.

Return type:: dict[str, typing.Any]
Returns:: Dictionary mapping type names to objects.

get_event_runtime_hints()[source]#

Return default event runtime hints for this configuration.

Return type:: EventRuntimeHints

Connection Parameters#

class sqlspec.adapters.cockroach_asyncpg.CockroachAsyncpgConnectionConfig[source]#

Bases: TypedDict

AsyncPG connection parameters for CockroachDB.

Pool Parameters#

class sqlspec.adapters.cockroach_asyncpg.CockroachAsyncpgPoolConfig[source]#

Bases: CockroachAsyncpgConnectionConfig

AsyncPG pool parameters for CockroachDB.

Driver Features#

class sqlspec.adapters.cockroach_asyncpg.CockroachAsyncpgDriverFeatures[source]#

Bases: TypedDict

Driver feature flags for CockroachDB AsyncPG adapter.

enable_native_storage: Opt in to server-side storage operations. Exports create: generated files under a prefix; imports temporarily take the table offline.
native_storage_csv_options: Explicit nullas/nullif markers and skip count.: Export is headerless; native CSV import requires an explicit skip count. Markers must not occur as literal data. No NULL marker is inferred.
on_connection_create: Async callback executed when a connection is acquired from pool.: Receives the raw asyncpg connection for low-level driver configuration. Called after internal setup (JSON codecs, pgvector registration).

Retry Configuration#

class sqlspec.adapters.cockroach_asyncpg.CockroachAsyncpgRetryConfig[source]#

Bases: object

CockroachDB asyncpg transaction retry configuration.

__init__(max_retries=10, base_delay_ms=50.0, max_delay_ms=5000.0, enable_logging=True)[source]#

classmethod from_features(driver_features)[source]#

Build retry config from driver feature mappings.

Return type:: CockroachAsyncpgRetryConfig

Driver#

class sqlspec.adapters.cockroach_asyncpg.CockroachAsyncpgDriver[source]#

Bases: AsyncpgDriver

CockroachDB AsyncPG driver with retry support.

__init__(connection, statement_config=None, driver_features=None)[source]#

Initialize driver adapter with connection and configuration.

Parameters:

connection¶ (PoolConnectionProxy) -- Database connection instance
statement_config¶ (StatementConfig | None) -- Statement configuration for the driver
driver_features¶ (dict[str, typing.Any] | None) -- Driver-specific features like extensions, secrets, and connection callbacks
observability¶ -- Optional runtime handling lifecycle hooks, observers, and spans

async select_to_storage(statement, destination, /, *parameters, statement_config=None, partitioner=None, format_hint=None, telemetry=None, **kwargs)[source]#

Export native CSV/Parquet to generated files under a remote prefix.

CSV is headerless. NULL values require an explicit nullas convention; server failures propagate without replay through the inherited path.

Return type:: StorageBridgeJob

async load_from_storage(table, source, *, file_format, partitioner=None, overwrite=False)[source]#

Append via native IMPORT when eligible, taking the table offline.

CockroachDB invalidates foreign keys during IMPORT. Overwrite and active transactions use the inherited path. CSV needs explicit skip (0 for no header); nullif is never inferred. Native failures are never replayed.

Return type:: StorageBridgeJob

async begin()[source]#

Begin a transaction and apply follower-read staleness to it.

The staleness clause must be the transaction's first statement, so it is applied only when this call is what opened the transaction. A session configured for follower reads is read-only: CockroachDB rejects writes against a historical timestamp.

Return type:: None

async run_transaction_with_retry(operation)[source]#

Execute a full CockroachDB transaction callback with serialization retries.

A rollback that itself fails leaves the transaction aborted, and every statement on it would then fail with a non-retryable error that hides the real conflict. The original exception is raised instead of retrying, so the rollback failure never replaces the transaction's outcome.

Return type:: TypeVar(_T)

async dispatch_execute(cursor, statement)[source]#

Execute single SQL statement.

Handles both SELECT queries and non-SELECT operations.

Parameters:

cursor¶ (Any) -- AsyncPG connection object
statement¶ (SQL) -- SQL statement to execute

Return type:

ExecutionResult

Returns:

ExecutionResult with statement execution details

async dispatch_execute_many(cursor, statement)[source]#

Execute SQL with multiple parameter sets using AsyncPG's executemany.

Parameters:

cursor¶ (Any) -- AsyncPG connection object
statement¶ (SQL) -- SQL statement with multiple parameter sets

Return type:

ExecutionResult

Returns:

ExecutionResult with batch execution details

async dispatch_execute_script(cursor, statement)[source]#

Execute SQL script with statement splitting and parameter handling.

Parameters:

cursor¶ (Any) -- AsyncPG connection object
statement¶ (SQL) -- SQL statement containing multiple statements

Return type:

ExecutionResult

Returns:

ExecutionResult with script execution details

handle_database_exceptions()[source]#

Handle database exceptions with PostgreSQL error codes.

Return type:: CockroachAsyncpgExceptionHandler

property data_dictionary: CockroachAsyncpgDataDictionary#

Get the data dictionary for this driver.

Returns:: Data dictionary instance for metadata queries

Data Dictionary#

class sqlspec.adapters.cockroach_asyncpg.data_dictionary.CockroachAsyncpgDataDictionary[source]#

Bases: AsyncDataDictionaryBase

CockroachDB async data dictionary (AsyncPG).

dialect: ClassVar[str] = 'cockroachdb'#: Dialect identifier. Must be defined by subclasses as a class attribute.

async get_metadata_capabilities(driver, domains=None)[source]#

Get CockroachDB data-dictionary capability profile.

Return type:: MetadataCapabilityProfile

async get_system_metadata_capabilities(driver, domains=None)[source]#

Get CockroachDB opt-in system metadata capability disclosures.

Return type:: tuple[SystemMetadataCapability, ...]

async get_schemas(driver)[source]#

Get schema metadata.

Return type:: MetadataResult

async get_objects(driver, schema=None)[source]#

Get database object metadata.

Return type:: MetadataResult

async get_table_details(driver, table, schema=None)[source]#

Get rich table metadata.

Return type:: MetadataResult

async get_constraints(driver, table=None, schema=None)[source]#

Get constraint metadata.

Return type:: MetadataResult

async get_views(driver, schema=None)[source]#

Get view metadata.

Return type:: MetadataResult

async get_routines(driver, schema=None)[source]#

Get routine metadata.

Return type:: MetadataResult

async get_privileges(driver, object_name=None, schema=None)[source]#

Get privilege metadata.

Return type:: MetadataResult

async get_dependencies(driver, object_name=None, schema=None)[source]#

Get stable dependency metadata.

Return type:: MetadataResult

async get_ddl(driver, object_name, schema=None, *, object_type='table', include_dependencies=True, prefer_native=True, redact=True)[source]#

Get lossy CockroachDB DDL status without parameterized identifiers.

Return type:: DDLResult

async get_system_metadata(driver, request=None, **kwargs)[source]#

Get opt-in CockroachDB system metadata.

Return type:: SystemMetadataResult

async get_version(driver)[source]#

Get database version information.

Parameters:: driver¶ (CockroachAsyncpgDriver) -- Async database driver instance
Return type:: VersionInfo | None
Returns:: Version information or None if detection fails

async get_feature_flag(driver, feature)[source]#

Check if database supports a specific feature.

Parameters:

driver¶ (CockroachAsyncpgDriver) -- Async database driver instance
feature¶ (str) -- Feature name to check

Return type:

bool

Returns:

True if feature is supported, False otherwise

async get_optimal_type(driver, type_category)[source]#

Get optimal database type for a category.

Parameters:

driver¶ (CockroachAsyncpgDriver) -- Async database driver instance
type_category¶ (str) -- Type category

Return type:

str

Returns:

Database-specific type name

async get_tables(driver, schema=None)[source]#

Get list of tables in schema.

Parameters:

driver¶ (CockroachAsyncpgDriver) -- Async database driver instance
schema¶ (str | None) -- Schema name (None for default)

Return type:

list[TableMetadata]

Returns:

List of table metadata dictionaries

async get_columns(driver, table=None, schema=None)[source]#

Get column information for a table or schema.

Parameters:

driver¶ (CockroachAsyncpgDriver) -- Async database driver instance
table¶ (str | None) -- Table name (None to fetch columns for all tables in schema)
schema¶ (str | None) -- Schema name (None for default)

Return type:

list[ColumnMetadata]

Returns:

List of column metadata dictionaries

async get_indexes(driver, table=None, schema=None)[source]#

Get index information for a table or schema.

Parameters:

driver¶ (CockroachAsyncpgDriver) -- Async database driver instance
table¶ (str | None) -- Table name (None to fetch indexes for all tables in schema)
schema¶ (str | None) -- Schema name (None for default)

Return type:

list[IndexMetadata]

Returns:

List of index metadata dictionaries

async get_foreign_keys(driver, table=None, schema=None)[source]#

Get foreign key metadata.

Parameters:

driver¶ (CockroachAsyncpgDriver) -- Async database driver instance
table¶ (str | None) -- Optional table name filter
schema¶ (str | None) -- Optional schema name filter

Return type:

list[ForeignKeyMetadata]

Returns:

List of foreign key metadata

Native object storage#

Set driver_features={"enable_native_storage": True} to use CockroachDB EXPORT/IMPORT for remote CSV and Parquet. The feature is disabled by default. Exports create generated files below a prefix; imports append into existing tables, take them offline and invalidate foreign keys. Asyncpg requires no additional autocommit setting. Calls made inside an active transaction use the inherited client-side path.

import os
from urllib.parse import urlsplit, urlunsplit
from uuid import uuid4

from sqlspec.adapters.cockroach_asyncpg import CockroachAsyncpgConfig

config = CockroachAsyncpgConfig(
    connection_config={"dsn": os.environ["DATABASE_URL"]},
    driver_features={
        "enable_native_storage": True,
        "native_storage_csv_options": {"skip": 0},
    },
)
parts = urlsplit(os.environ["COCKROACH_STORAGE_URI"])
destination = urlunsplit(
    parts._replace(path=parts.path.rstrip("/") + "/" + uuid4().hex)
)
try:
    async with config.provide_session() as driver:
        exported = await driver.select_to_storage(
            "SELECT :id::INT8 AS id, :label::STRING AS label",
            destination,
            id=7,
            label="O'Reilly",
            format_hint="csv",
        )
        prefix = urlsplit(exported.telemetry["destination"])
        for filename in exported.telemetry["extra"]["files"]:
            source = urlunsplit(
                prefix._replace(path=prefix.path.rstrip("/") + "/" + filename)
            )
            await driver.load_from_storage("archive_items", source, file_format="csv")
finally:
    await config.close_pool()

The example assumes an existing archive_items(id INT8, label STRING) table and a server-readable storage URI. CSV export is headerless: skip=0 imports all rows. Use skip=1 for an external file with one header. Without explicit skip, CSV import uses the inherited header-aware reader. NULL markers (nullas for export, nullif for import) are optional explicit strings; choose values absent from literal data. Unconfigured NULL export fails. Parquet avoids marker collisions and is the default export format.

Remote aliases resolve through the storage registry. Disabled, local, unsupported-format/protocol and overwrite=True calls use the inherited path. Native failures propagate without an inherited retry. File metadata, privileges, cleanup, supported query forms, server restrictions and benchmark measurement details match CockroachDB + Psycopg.

Local tests cover CockroachDB CCL v26.2.5 with RustFS. CockroachDB Cloud tiers, hosted privileges, cloud credentials, GCS and Azure remain unverified.

Extension Settings#

Use the configuration types below in their corresponding extension_config namespace: "litestar", "events", or "adk" as supported by this adapter.

class sqlspec.adapters.cockroach_asyncpg.litestar.CockroachAsyncpgLitestarConfig[source]#

Bases: LitestarConfig

CockroachAsyncpg-specific Litestar settings.

Use inside extension_config["litestar"] with this adapter's session store.

enable_hash_sharded_indexes: NotRequired[bool]#: Enable hash-sharded session indexes.

hash_shard_bucket_count: NotRequired[int]#: Number of hash index buckets.

ttl_expiration_expression: NotRequired[Literal[False, 'expires_at']]#: Enable row-level TTL using expires_at, or disable it with False.

class sqlspec.adapters.cockroach_asyncpg.adk.CockroachAsyncpgADKConfig[source]#

Bases: ADKConfig

CockroachDB asyncpg ADK extension settings.

Use these keys inside extension_config["adk"] with the CockroachDB asyncpg ADK stores.

table_locality: NotRequired[str]#: Default raw CockroachDB LOCALITY clause for ADK tables.

session_table_locality: NotRequired[str]#: Raw CockroachDB LOCALITY clause for the ADK session table.

events_table_locality: NotRequired[str]#: Raw CockroachDB LOCALITY clause for the ADK events table.

app_state_table_locality: NotRequired[str]#: Raw CockroachDB LOCALITY clause for the ADK app state table.

user_state_table_locality: NotRequired[str]#: Raw CockroachDB LOCALITY clause for the ADK user state table.

metadata_table_locality: NotRequired[str]#: Raw CockroachDB LOCALITY clause for the ADK metadata table.

memory_table_locality: NotRequired[str]#: Raw CockroachDB LOCALITY clause for the ADK memory table.

enable_hash_sharded_indexes: NotRequired[bool]#: Create CockroachDB hash-sharded secondary indexes on hot timestamp access paths.

hash_shard_bucket_count: NotRequired[int]#: Optional bucket count for CockroachDB hash-sharded secondary indexes.

enable_storing_indexes: NotRequired[bool]#: Add CockroachDB STORING columns to common ADK replay/session indexes.

enable_memory_trigram_index: NotRequired[bool]#: Create a CockroachDB trigram GIN index for simple ILIKE memory search.

Native startup settings#

connection_config accepts application_name, default_transaction_use_follower_reads (boolean), and results_buffer_size (non-negative integer bytes). These map to asyncpg server_settings; explicit entries in server_settings take precedence.