CockroachDB + Psycopg#
CockroachDB adapter using psycopg with automatic transaction retry logic. Provides both sync and async support.
Sync Configuration#
- class sqlspec.adapters.cockroach_psycopg.CockroachPsycopgSyncConfig[source]#
Bases:
SyncDatabaseConfig[CrdbConnection,ConnectionPool,CockroachPsycopgSyncDriver]Configuration for CockroachDB synchronous connections using psycopg.
- driver_type#
alias of
CockroachPsycopgSyncDriver
- __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]#
- create_connection()[source]#
Open a standalone connection owned by the caller.
- Return type:
CrdbConnection- Returns:
A connection that is not bound to the pool and can be closed directly.
- provide_session(*_args, statement_config=None, follower_reads=None, staleness=None, **_kwargs)[source]#
Provide a database session context manager.
- Return type:
CockroachPsycopgSyncSessionContext
- 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.
Async Configuration#
- class sqlspec.adapters.cockroach_psycopg.CockroachPsycopgAsyncConfig[source]#
Bases:
AsyncDatabaseConfig[AsyncCrdbConnection,AsyncConnectionPool,CockroachPsycopgAsyncDriver]Configuration for CockroachDB async connections using psycopg.
- driver_type#
alias of
CockroachPsycopgAsyncDriver
- __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.
- Return type:
AsyncCrdbConnection- Returns:
A connection that is not bound to the pool and can be closed directly.
- provide_session(*_args, statement_config=None, follower_reads=None, staleness=None, **_kwargs)[source]#
Provide a database session context manager.
- Return type:
CockroachPsycopgAsyncSessionContext
- async provide_pool(*args, **kwargs)[source]#
Provide pool instance.
- Return type:
AsyncConnectionPool
- 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.
Connection Parameters#
Pool Parameters#
- class sqlspec.adapters.cockroach_psycopg.CockroachPsycopgPoolConfig[source]#
Bases:
CockroachPsycopgConnectionConfigCockroachDB pool parameters.
Driver Features#
- class sqlspec.adapters.cockroach_psycopg.CockroachPsycopgDriverFeatures[source]#
Bases:
TypedDictCockroachDB driver feature configuration.
- 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: Callback executed when a connection is acquired from pool.
For sync: Callable[[CockroachSyncConnection], None] For async: Callable[[CockroachAsyncConnection], Awaitable[None]] Called after internal setup.
Sync Driver#
- class sqlspec.adapters.cockroach_psycopg.CockroachPsycopgSyncDriver[source]#
Bases:
PsycopgSyncDriverCockroachDB sync driver using psycopg.crdb.
- 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:
- 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:
- 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:
- 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)
- dispatch_execute(cursor, statement)[source]#
Execute single SQL statement.
- Parameters:
- Return type:
- Returns:
ExecutionResult with statement execution details
- dispatch_execute_many(cursor, statement)[source]#
Execute SQL with multiple parameter sets.
- Parameters:
- Return type:
- Returns:
ExecutionResult with batch execution details
- dispatch_execute_script(cursor, statement)[source]#
Execute SQL script with multiple statements.
- Parameters:
- Return type:
- Returns:
ExecutionResult with script execution details
- handle_database_exceptions()[source]#
Handle database-specific exceptions and wrap them appropriately.
- Return type:
CockroachPsycopgSyncExceptionHandler
- property data_dictionary: CockroachPsycopgSyncDataDictionary#
Get the data dictionary for this driver.
- Returns:
Data dictionary instance for metadata queries
Async Driver#
- class sqlspec.adapters.cockroach_psycopg.CockroachPsycopgAsyncDriver[source]#
Bases:
PsycopgAsyncDriverCockroachDB async driver using psycopg.crdb.
- 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:
- 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:
- 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:
- 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 (async).
- Parameters:
- Return type:
- Returns:
ExecutionResult with statement execution details
- async dispatch_execute_many(cursor, statement)[source]#
Execute SQL with multiple parameter sets (async).
- Parameters:
- Return type:
- Returns:
ExecutionResult with batch execution details
- async dispatch_execute_script(cursor, statement)[source]#
Execute SQL script with multiple statements (async).
- Parameters:
- Return type:
- Returns:
ExecutionResult with script execution details
- handle_database_exceptions()[source]#
Handle database-specific exceptions and wrap them appropriately.
- Return type:
CockroachPsycopgAsyncExceptionHandler
- property data_dictionary: CockroachPsycopgAsyncDataDictionary#
Get the data dictionary for this driver.
- Returns:
Data dictionary instance for metadata queries
Retry Configuration#
Sync Data Dictionary#
- class sqlspec.adapters.cockroach_psycopg.data_dictionary.CockroachPsycopgSyncDataDictionary[source]#
Bases:
SyncDataDictionaryBaseCockroachDB sync data dictionary.
- dialect: ClassVar[str] = 'cockroachdb'#
Dialect identifier. Must be defined by subclasses as a class attribute.
- get_metadata_capabilities(driver, domains=None)[source]#
Get CockroachDB data-dictionary capability profile.
- Return type:
- get_system_metadata_capabilities(driver, domains=None)[source]#
Get CockroachDB opt-in system metadata capability disclosures.
- get_dependencies(driver, object_name=None, schema=None)[source]#
Get stable dependency metadata.
- Return type:
- 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
- get_system_metadata(driver, request=None, **kwargs)[source]#
Get opt-in CockroachDB system metadata.
- Return type:
SystemMetadataResult
- get_feature_flag(driver, feature)[source]#
Check if CockroachDB supports a specific feature.
- Return type:
- get_optimal_type(driver, type_category)[source]#
Get optimal CockroachDB type for a category.
- Return type:
- get_columns(driver, table=None, schema=None)[source]#
Get column information for a table or schema.
- Return type:
- get_indexes(driver, table=None, schema=None)[source]#
Get index metadata for a table or schema.
- Return type:
Async Data Dictionary#
- class sqlspec.adapters.cockroach_psycopg.data_dictionary.CockroachPsycopgAsyncDataDictionary[source]#
Bases:
AsyncDataDictionaryBaseCockroachDB async data dictionary.
- 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:
- async get_system_metadata_capabilities(driver, domains=None)[source]#
Get CockroachDB opt-in system metadata capability disclosures.
- async get_constraints(driver, table=None, schema=None)[source]#
Get constraint metadata.
- Return type:
- async get_privileges(driver, object_name=None, schema=None)[source]#
Get privilege metadata.
- Return type:
- async get_dependencies(driver, object_name=None, schema=None)[source]#
Get stable dependency metadata.
- Return type:
- 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_feature_flag(driver, feature)[source]#
Check if CockroachDB supports a specific feature.
- Return type:
- async get_optimal_type(driver, type_category)[source]#
Get optimal CockroachDB type for a category.
- Return type:
- async get_columns(driver, table=None, schema=None)[source]#
Get column information for a table or schema.
- Return type:
- async get_indexes(driver, table=None, schema=None)[source]#
Get index metadata for a table or schema.
- Return type:
Native object storage#
Set driver_features={"enable_native_storage": True} to let CockroachDB
write or read CSV and Parquet directly. This changes the storage contract:
select_to_storage writes generated files below a destination prefix;
load_from_storage appends into an existing table and takes it offline.
Imports invalidate foreign keys, which must be validated afterward.
The feature is disabled by default.
Install sqlspec[cockroachdb,obstore] for both drivers and URI resolution.
The inherited CSV/Parquet paths additionally require pyarrow.
Both psycopg configurations require connection_config={"autocommit": True}
when native storage is enabled. SQLSpec never enables autocommit automatically.
The driver also checks the actual connection before each native call. Active
transactions, disabled features, unsupported formats/protocols, and
overwrite=True imports use the inherited client-side path. Native execution
or result errors propagate once; SQLSpec does not retry them through that path.
Native export supports a single SELECT, including query builders, filters and
bound parameters that compile to positional driver parameters. Scripts,
executemany statements and other statement types use the inherited path.
Parquet is the default export format. Eligible resolved protocols are s3,
gs, gcs and azure; the URI itself must use a scheme and credentials
CockroachDB understands. Local paths and local aliases use client-side storage.
Aliases resolve through the storage registry before database execution.
The destination must be reachable by every CockroachDB node. Credentials stored only in a Python storage backend are not transferred to the database. Supply a server-readable URI or configure the server's external-storage identity.
Export and import filenames#
This example assumes an existing archive_items(id INT8, label STRING) table
and a server-readable destination in COCKROACH_STORAGE_URI. Use a fresh
prefix for each export. extra["files"] contains the generated relative
filenames; preserve URI query parameters when constructing import sources.
import os
from urllib.parse import urlsplit, urlunsplit
from uuid import uuid4
from sqlspec.adapters.cockroach_psycopg import CockroachPsycopgSyncConfig
config = CockroachPsycopgSyncConfig(
connection_config={
"conninfo": os.environ["DATABASE_URL"],
"autocommit": True,
},
driver_features={"enable_native_storage": True},
)
parts = urlsplit(os.environ["COCKROACH_STORAGE_URI"])
destination = urlunsplit(
parts._replace(path=parts.path.rstrip("/") + "/" + uuid4().hex)
)
try:
with config.provide_session() as driver:
exported = driver.select_to_storage(
"SELECT :id::INT8 AS id, :label::STRING AS label",
destination,
id=7,
label="O'Reilly",
format_hint="parquet",
)
prefix = urlsplit(exported.telemetry["destination"])
for filename in exported.telemetry["extra"]["files"]:
source = urlunsplit(
prefix._replace(path=prefix.path.rstrip("/") + "/" + filename)
)
driver.load_from_storage("archive_items", source, file_format="parquet")
finally:
config.close_pool()
Export telemetry sums server-reported rows and file bytes. Import telemetry
reports rows and extra["job_id"]; its logical job byte count is not reported
as an artifact size. Partition metadata and supplied export telemetry retain
the standard storage job merge behavior. Generated files, including partial
files after failure, remain the caller's responsibility to clean up.
CSV conventions#
Native CSV exports have no header. For the non-NULL example above, select CSV
and configure native_storage_csv_options={"skip": 0}; import the returned
files with file_format="csv". skip=1 instead acknowledges one header
in an externally produced CSV file. Without an explicit skip, CSV import
uses the inherited header-aware reader. A boolean is not a valid skip count.
nullas sets an export NULL marker; nullif sets an import NULL marker.
Both accept strings, including an explicit empty string. No marker is inferred:
CSV export with NULL values and no nullas fails at the server. When using
markers, choose matching values that cannot occur as literal data. In particular,
using an empty marker can turn an empty string into NULL. Use Parquet to preserve
arbitrary NULL, empty-string and literal-marker values without this convention.
Only skip, nullas and nullif are accepted; option values are bound.
Privileges and validation boundary#
Export requires SELECT on the source table; import requires INSERT and DROP on
the target. Custom S3 endpoints and implicit cloud credentials require the admin
role or EXTERNALIOIMPLICITACCESS. Explicit credentials for standard cloud
storage have different privileges. Storage-provider permissions still apply.
Import cannot run during a rolling upgrade. Consult CockroachDB's
EXPORT reference and
IMPORT INTO reference
for server-version restrictions and foreign-key revalidation.
The import reference documents Parquet support as subject to change.
Local validation used CockroachDB CCL v26.2.5 and RustFS. It covers all three SQLSpec Cockroach drivers, CSV/Parquet roundtrips, bound values, transaction refusal, NULL conventions and custom-endpoint privileges. CockroachDB Cloud tiers, hosted privileges, cloud identity chains, GCS and Azure are unverified. Local results do not establish support on every hosted tier.
Performance measurement#
tools/scripts/bench_cockroach_storage.py compares native and inherited
Parquet export/import with explicit local connection and storage configuration.
Run --help for environment variables, row counts, warmups and iterations.
It records per-sample timings, medians and ranges. A second connection samples
table availability; its polling interval and query timeout limit the resolution.
Zero failed probes does not prove the table was never offline. Each run removes
only its generated tables and random storage prefix. Native job startup can
outweigh avoided client transfers; benchmark the intended workload.
A local Python 3.14 run against the services above used an INT8 key and a 64-character string, one warmup and three measured samples per case. Times below are median milliseconds (minimum–maximum). Native operations were slower at all three sizes in this run; these small local samples do not predict cloud or large distributed workloads.
Rows |
Client export |
Native export |
Client import |
Native import |
|---|---|---|---|---|
100 |
5.34 (4.89–5.36) |
12.78 (12.11–13.06) |
8.61 (7.83–9.08) |
107.03 (103.43–108.72) |
1,000 |
6.59 (6.24–6.73) |
13.55 (13.25–13.73) |
18.92 (18.75–21.98) |
112.92 (110.56–119.71) |
10,000 |
16.62 (16.50–17.99) |
23.18 (20.85–28.36) |
76.41 (75.57–92.64) |
129.86 (122.95–139.74) |
Native imports produced two to four failed availability probes per sample, with observed spans between 21.3 and 63.2 ms. Client imports produced none. Polling was every 20 ms with a 250 ms query timeout. These spans are sampled observations of one run, not exact table-offline durations, and repeated runs move both the probe counts and the spans.
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_psycopg.litestar.CockroachPsycopgLitestarConfig[source]#
Bases:
LitestarConfigCockroachPsycopg-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_psycopg.adk.CockroachPsycopgADKConfig[source]#
Bases:
ADKConfigCockroachDB psycopg ADK extension settings.
Use these keys inside
extension_config["adk"]with the CockroachDB psycopg 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 default_transaction_use_follower_reads as a
boolean, plus non-negative integer results_buffer_size (bytes),
statement_timeout and idle_in_transaction_session_timeout (milliseconds).
These settings append to the existing libpq options string.