diff --git a/docs/changelog.md b/docs/changelog.md index d077efe62..dcb6bec17 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -11,6 +11,8 @@ important operational fixes. **Added:** +* Added a pool for the native PostgreSQL ADBC driver. Set `AdbcConfig(pool_config={})` to reuse connections. Each lease has a wait timeout; each return clears session state. Pooling is off by default. See the [ADBC guide](databases/adbc.md#postgresql-connection-reuse) for supported settings. + * Added Sybase Adaptive Server Enterprise (`sqlspec.adapters.sybase` with `SybaseSyncConfig` and `SybaseSyncDriver`) and SAP SQL Anywhere (`sqlspec.adapters.sqlanywhere` with `SqlAnywhereSyncConfig` and `SqlAnywhereSyncDriver`) database adapters, custom SQLGlot dialects (`Sybase` and `SQLAnywhere` in `sqlspec.dialects.sybase`), data dictionaries (`SybaseSyncDataDictionary` and `SqlAnywhereSyncDataDictionary` with `SYBASE_CONFIG` and `SQLANYWHERE_CONFIG`), type converters (`SybaseOutputConverter` and `SqlAnywhereOutputConverter`), and extension stores for Litestar sessions (`SybaseStore`, `SqlAnywhereStore`), Event queues (`SybaseEventQueueStore`, `SqlAnywhereEventQueueStore`), and Google ADK sessions and memory (`SybaseADKStore`, `SybaseADKMemoryStore`, `SqlAnywhereADKStore`, `SqlAnywhereADKMemoryStore`). ([#870](https://github.com/litestar-org/sqlspec/pull/870)) * Added `arrow_odbc` dialect auto-detection and `information_schema` / system catalog introspection for Sybase ASE (`Adaptive Server Enterprise` / `ASE`) and SAP SQL Anywhere (`SQL Anywhere` / `Sybase IQ`), routing detected connections to the `sybase` and `sqlanywhere` data dictionaries. @@ -25,24 +27,85 @@ important operational fixes. **Changed:** -* Replaced correlated per-column scalar subqueries in data dictionary `columns.sql` and `indexes.sql` catalog queries across `postgres`, `cockroach`, `mssql`, `mysql`, `oracle`, `db2`, and `duckdb` with pre-aggregated CTEs and joins, and reduced `ColumnMetadata` to its 18 core catalog fields. + +* Renamed public Litestar stores: `SQLiteStore` to `SqliteStore`, `ADBCStore` to `AdbcStore`, and `DuckdbStore` to `DuckDBStore`. Update imports and store construction accordingly. DuckDB ADK exports likewise use `DuckDBADKStore`, `DuckDBADKMemoryStore`, `DuckDBADKConfig`, and `DuckDBADKFTSOptions` instead of their `Duckdb` spellings. + +* Removed `SybaseSyncEventQueueStore` and `SqlanywhereSyncEventQueueStore`; use `SybaseEventQueueStore` and `SqlanywhereEventQueueStore` respectively. + +* Standardized database failures as SQLSpec exceptions across statement execution, cached queries, bulk loads, and transaction boundaries. Code catching native driver errors should catch the corresponding `sqlspec.exceptions` classes instead. Constraint violations retain specific classes such as `UniqueViolationError` and `ForeignKeyViolationError`; syntax errors raise `SQLParsingError`. MySQL errors 1061 and 1091 now raise `OperationalError` instead of being suppressed. Original driver exceptions remain available as causes when the driver supplies an exception object. + +* Changed default SQLite `:memory:` pools to share one database across connections within that pool. Shared cache does not enable dirty reads by default; applications that deliberately need dirty reads must explicitly configure `read_uncommitted`. + +* Successful Psycopg `on_connection_create` hooks now commit their setup transaction before a connection enters the pool. Session settings such as `search_path` and timezone persist, and failed setup is rolled back. Hooks must not rely on their writes remaining uncommitted. + +* ADK table-name validation now rejects schema or dataset qualifiers where the selected adapter cannot generate valid qualified DDL. Supported schema-qualified stores retain that capability. + +* ADK schema changes target newly initialized stores. Existing ADK tables must be reset and recreated for this release; resetting removes their stored sessions and memories. Upgrade migrations for these schema changes are deferred. + +* Removed legacy `pgbouncer` and `pgbouncer_mode` flags from PostgreSQL driver features across `asyncpg`, `psycopg`, and `psqlpy` in favor of canonical `pooler_mode="transaction"`. + +* Removed the `pgbouncer` key from `AsyncpgConnectionConfig`; set `driver_features={"pooler_mode": "transaction"}` instead. + +* Removed the empty `SybaseADKConfig`, `SqlanywhereADKConfig`, `Db2LitestarConfig`, and `Db2EventsConfig` TypedDicts; use the base `ADKConfig`, `LitestarConfig`, and `EventsConfig`. + +* Renamed `use_pool` to `enable_pooling` (default `True`) across `mssql_python` and `pymssql` driver features. + +* Renamed `arrow_odbc` `enable_driver_pooling` to `enable_pooling` and BigQuery `use_query_and_wait` to `enable_query_and_wait`. + +* Removed the ADBC `strict_type_coercion` and `arrow_extension_types` driver feature aliases; use `enable_strict_type_coercion` and `enable_arrow_extension_types`. + +* Declared driver feature keys that adapters already read: `events_backend` on `ArrowOdbcDriverFeatures` and `default_query_job_config` on `BigQueryDriverFeatures`. + +* Replaced `null_pool` in `psycopg` connection configuration and driver features with canonical `enable_pooling=False`. + +* Defaulted `enable_correlation_middleware` to `True` across all web framework extensions (`Litestar`, `FastAPI`, `Starlette`, `Sanic`, `Flask`). + +* Standardized adapter driver imports: importing `sqlspec.adapters.` eagerly imports that adapter's driver, raising the native `ModuleNotFoundError` when the driver is not installed, while non-driver optional dependencies load on first use. + +* Documented Oracle AQ dequeue commit behavior as at-most-once delivery. + +* Added SAP SQL Anywhere integration driver test suite with hermetic in-container client runtime extraction. + +* Replaced correlated per-column scalar subqueries in data dictionary `columns.sql` and `indexes.sql` catalog queries across `postgres`, `cockroach`, `mysql`, `oracle`, `db2`, and `duckdb` with pre-aggregated CTEs and joins, and reduced `ColumnMetadata` to its 18 core catalog fields. ([#866](https://github.com/litestar-org/sqlspec/pull/866)) + * Redesigned the documentation site architecture around Zensical and `mkdocstrings-python` with Diataxis-aligned navigation (`Getting Started`, `Learn`, `Adapters`, `Extensions`, `Examples`, `Reference`), interactive WebAssembly SQL playground, runnable `docs/examples/` snippets, conceptual Mermaid diagrams, and tabbed package manager and `SQLFileLoader` examples. ([#864](https://github.com/litestar-org/sqlspec/pull/864), [#867](https://github.com/litestar-org/sqlspec/pull/867), [#869](https://github.com/litestar-org/sqlspec/pull/869)) + * Migrated integration test infrastructure to clientless `pytest-databases` fixtures (`#859`). ([#859](https://github.com/litestar-org/sqlspec/pull/859)) + * Upgraded CI GitHub Actions `astral-sh/setup-uv` from `10.3.1` to `10.3.3` and pinned `actions/checkout` (`6.0.2`), `actions/setup-python` (`7.1.0`), `actions/cache` (`6.0.0`), and `actions/upload-pages-artifact` / `actions/deploy-pages` in `.github/workflows/ci.yml` and `.github/workflows/docs.yml`. ([#872](https://github.com/litestar-org/sqlspec/pull/872)) **Fixed:** + +* ADBC now closes its own connections if a hook, profile setup, driver setup, or session setup fails. It keeps the first error and leaves supplied connections open. + +* Explicit ADBC transactions now honor native `autocommit=True` on supported drivers. If rollback cleanup fails, autocommit stays off so mode restoration cannot commit pending writes. + +* Preserved supported Arrow conversion and mutation paths for Asyncmy and Spanner while reporting their native capabilities accurately; corrected ADBC lowercase readers, DuckDB stream ownership and table filtering, and SQL Server Arrow cursor lifetimes. + +* Corrected Oracle session-renewal bindings, pipeline error conversion, ADK batch-error handling, standalone connection setup, and deferred Arrow failures. Text-only Oracle memory stores remain available on servers without native vector support. + +* Preserved PostgreSQL cache-hit query semantics, affected-row counts, mapped exceptions, JSON scalar values, and ADK embeddings. Psycopg connection hooks leave pooled connections idle, COPY finalization errors are translated, and Asyncpg continue-on-error stacks recover through savepoints inside caller transactions. + +* Preserved SQL Server 2016 constraint and index metadata queries with compatible catalog aggregation. + +* Corrected SQL Anywhere implicit-transaction cleanup, Db2 ADK index-name matching and event deduplication, MySQL batch percent escaping, optional Connector/Python async-pool imports, privilege filters, and mutable DDL capture settings. + * Removed runtime `duckdb` imports from `sqlspec.adapters.adbc` (`AdbcDataDictionary` and `sqlspec/adapters/adbc/core.py`) and enforced strict adapter import isolation across all 20 database adapters in `tests/unit/test_import_time.py`. ([#865](https://github.com/litestar-org/sqlspec/pull/865)) + * Eliminated extra row copies and `UserWarning` in `DuckDBDriver.select_to_arrow` on DuckDB 2.0+ by routing parameterless `SELECT` queries through `connection.sql(query).to_arrow_table()` and `to_arrow_reader(batch_size=...)` when no open DB-API transaction is active on the connection. + * Fixed `compile_table_ddl_results` and `CreateIndex` cross-dialect DDL synthesis to emit `COMMENT ON` only on dialects that support it (`postgres`, `oracle`, `db2`, `duckdb`, `snowflake`), inline `COMMENT` clauses on `mysql` and `mariadb`, render `STORING (...)` on `spanner` and `spangres`, fold unquoted source identifiers to target case before quoting target-only reserved keywords, and preserve SQLite `NUMERIC` type affinity. + * Fixed `MigrationSquasher.fold_canonical_ddl` to fall back cleanly to raw statement concatenation when a migration contains interleaved DML or procedural blocks (`INSERT`, `UPDATE`, `DELETE`, `MERGE`, `COPY`, `DO`, `CALL`), preserve table-level constraints across `ALTER TABLE ... RENAME COLUMN` and `DROP COLUMN`, and update `sqlspec_migrations` tracker history atomically inside a single transaction when applying squashed migrations. + * Fixed `write_fixture_shards` to unlink stale shard and legacy single-file artifacts (`.json`, `
.json.gz`, and higher-numbered shards) before writing new shards so subsequent `load_fixture` calls do not read orphaned records. ### v0.69.0 - Oracle PKCS#12 wallet conversion, TNS alias resolution, proximity RRF, and server-side DBMS_CLOUD ingestion @@ -260,7 +323,7 @@ important operational fixes. savepoints, SQLSTATE/vendor exception mapping, and default `enable_lowercase_column_names=True` normalization for Oracle and Db2. ([#839](https://github.com/litestar-org/sqlspec/pull/839)) -* `ADBCStore` (Litestar sessions), `AdbcEventQueueStore` (event channels), +* `AdbcStore` (Litestar sessions), `AdbcEventQueueStore` (event channels), `AdbcADKStore`, and `AdbcADKMemoryStore` (Google ADK sessions and memory) add Oracle (`oracle`) and IBM Db2 (`db2`) backend support with idempotent DDL guards (PL/SQL `SQLCODE` blocks for Oracle and `SYSCAT` catalog probes diff --git a/docs/contributing/creating-adapters.md b/docs/contributing/creating-adapters.md index 81bb60e85..727760a5e 100644 --- a/docs/contributing/creating-adapters.md +++ b/docs/contributing/creating-adapters.md @@ -41,6 +41,18 @@ sqlspec/adapters// - **`driver.py`**: Subclasses `SyncDriverAdapterBase` or `AsyncDriverAdapterBase`, implementing `dispatch_execute`, `dispatch_execute_many`, `dispatch_execute_script`, `begin`, `commit`, `rollback`, `with_cursor`, `handle_database_exceptions`, `_connection_in_transaction`, and `data_dictionary`. - **`data_dictionary.py`**: Subclasses `SyncDataDictionaryBase` or `AsyncDataDictionaryBase`, binding the dialect's SQL query pack from `sqlspec/data_dictionary/dialects//` to introspect tables, columns, indexes, foreign keys, and server versions. +## Import Contract + +Every driver is optional: users install only the extras for the adapters they use, and any one adapter (or all of them) works without the other drivers installed. Each adapter follows the same import rules: + +- **Own driver, imported eagerly.** `import sqlspec.adapters.` imports that adapter's driver once, at the top of `_typing.py`, as aliased re-exports (`from asyncpg import Connection as AsyncpgRawConnection`). A missing driver raises the driver's own `ModuleNotFoundError`. Do not wrap the adapter's own driver in `import_optional`, `import_optional_attr`, `importlib.import_module`, `try`/`except ImportError` fallbacks, `None` placeholders, stub classes, or a module `__getattr__`. +- **No other drivers.** An adapter never imports another adapter's driver or package. The CockroachDB adapters wrapping `asyncpg` and `psycopg` are the only allowed cross-adapter imports. +- **One source of driver names.** `config.py`, `core.py`, `driver.py`, `pool.py`, and the extension stores import driver names from `_typing`, not from the driver package. +- **One `TYPE_CHECKING` block.** `_typing.py` uses a single `if TYPE_CHECKING: ... else: ...` block, only for subscripted generics (`Pool[Record]`), `Protocol` types for compiled driver classes, sqlspec imports that would be circular, and secondary optional types. +- **Secondary dependencies through the shared shim.** When code cannot avoid referencing an optional dependency at import time (a runtime-evaluated annotation, an `isinstance` target, or an availability check), import the type or flag from `sqlspec.typing` (for example `ArrowTable`, `NumpyArray`, `PYARROW_INSTALLED`, `PGVECTOR_INSTALLED`). `sqlspec/_typing.py` resolves these lazily and supplies protocol fallbacks when the package is absent. If the shim lacks a type or flag you need, add it there instead of writing a local fallback. When the module itself is needed (`pyarrow`, `pandas`, `polars`, `numpy`, `pgvector`, `rich`, cloud connectors, storage backends, observability packages), load it inside the function that uses it (`ensure_*` when required, `import_optional` when there is a fallback), never at module import. + +Register a new adapter's driver module(s) in `ADAPTER_DRIVERS` in `tests/unit/test_import_time.py`. The import-contract tests fail for unregistered adapters and check that the adapter raises the driver's own error when the driver is missing, imports with every other driver blocked, and loads no secondary optional module at import. If the driver package itself imports an optional module, list it in `DRIVER_OWNED_OPTIONAL_IMPORTS` with the reason. For development, add the driver to the `extras` dependency group in `pyproject.toml`; dependency groups are not published to PyPI. + ## The `apply_driver_features` Contract Each adapter's `core.py` normalizes user-supplied `driver_features` and applies custom JSON serializers or type converters to the adapter's `StatementConfig`: diff --git a/docs/databases/adbc.md b/docs/databases/adbc.md index 0ecf99aca..e391214e4 100644 --- a/docs/databases/adbc.md +++ b/docs/databases/adbc.md @@ -118,6 +118,66 @@ with spec.provide_session(config) as session: ) ``` +## PostgreSQL Connection Reuse + +ADBC opens and closes a connection for each session by default. For short +PostgreSQL transactions, enable connection reuse with `pool_config={}` or supply +explicit bounds: + +```python +import os + +from sqlspec.adapters.adbc import AdbcConfig + +config = AdbcConfig( + connection_config={"uri": os.environ["DATABASE_URI"]}, + pool_config={"max_size": 4, "timeout": 30.0}, +) +try: + with config.provide_session() as session: + value = session.select_value("SELECT :value", value=42) + session.commit() +finally: + config.close_pool() +``` + +`pool_config=None` keeps a fresh connection per session. An empty mapping gives +four leases and a 30-second wait limit. Each lease gives one caller sole use of a +connection. The pool holds one extra seed, for a total of `max_size + 1` native +connections. Nested sessions need separate leases and can time out when the pool +is full. Standalone connections do not count toward this limit. + +Both native packages, `adbc-driver-manager` and `adbc-driver-postgresql`, must be +version 1.12.0 or newer. This pool supports only the native PostgreSQL driver +with `autocommit=False`. It rejects supplied connections and +`on_connection_create` hooks. Other backends keep a fresh connection per +session. The `supports_connection_pooling` flag means the adapter offers a +pool; the flag does not say whether a given config has enabled it. + +On clean exit, the pool rolls back unfinished work, runs `DISCARD ALL` outside +a transaction, and turns native autocommit off. This clears committed session +settings, temporary tables, prepared statements, listeners, and advisory locks. +A failed session or reset drops that connection from reuse. Direct changes to +native connection options fall outside this reset contract. Measure complete +transactions, including reset, on your own workload before choosing a pool. + +Lazy row streams and Arrow readers belong to their session. On exit, the +session closes their native readers and cursors, then drops the connection +from reuse. This also applies to streams that were never read or fully read. +Consume streams inside the session. Materialized tables and eagerly drained +batch lists stay valid afterward and allow reuse. + +`close_pool()` blocks new checkouts and closes idle connections. Active leases +finish first, then close; their work is not interrupted. If native cleanup +fails, call `close_pool()` again to retry. The pool keeps ownership until every +handle is closed. A later explicit session can then start a new generation. +Create configs after forking: a pool from the parent process refuses use in a +child. Pools do not refresh credentials on their own. + +`create_connection()` and `provide_connection()` remain separate from the pool. +An unpooled `connection_instance` stays open when a session ends; the caller +owns it. + ## Key Driver Features and Pool Tuning | Setting | Location | Description | diff --git a/docs/databases/arrow-odbc.md b/docs/databases/arrow-odbc.md index 071b82bf7..6d1612b4f 100644 --- a/docs/databases/arrow-odbc.md +++ b/docs/databases/arrow-odbc.md @@ -50,7 +50,7 @@ config = spec.add_config( }, statement_config=StatementConfig(dialect="tsql"), driver_features={ - "enable_driver_pooling": True, + "enable_pooling": True, "chunk_size": 65536, "fetch_concurrently": True, }, @@ -60,13 +60,13 @@ config = spec.add_config( ## Key Driver Features and Pool Tuning -`ArrowOdbcConfig` subclasses `NoPoolSyncConfig` and delegates connection reuse to ODBC Driver Manager pooling when `enable_driver_pooling=True`: +`ArrowOdbcConfig` subclasses `NoPoolSyncConfig` and delegates connection reuse to ODBC Driver Manager pooling when `enable_pooling=True`: | Setting | Location | Description | | --- | --- | --- | | `connection_string` | `connection_config` | Standard semicolon-delimited ODBC connection string (`Driver={...};Server=...;`). | | `user` / `password` | `connection_config` | Optional credentials passed alongside the ODBC connection string. | -| `enable_driver_pooling` | `driver_features` | Enables process-wide ODBC Driver Manager connection pooling for low-latency session checkouts. | +| `enable_pooling` | `driver_features` | Enables process-wide ODBC Driver Manager connection pooling for low-latency session checkouts. | | `chunk_size` / `max_bytes_per_batch` | `driver_features` | Maximum rows and memory bytes allocated per Arrow `RecordBatch` buffer during streaming reads. | | `fetch_concurrently` | `driver_features` | Allocates double-buffered background ODBC fetch threads in Rust while Python processes the current Arrow batch. | | `max_text_size` / `max_binary_size` | `driver_features` | Upper byte bounds for unbounded `VARCHAR(MAX)` / `NVARCHAR(MAX)` / `VARBINARY(MAX)` columns. | diff --git a/docs/databases/bigquery.md b/docs/databases/bigquery.md index 05aa60e0e..57c1990b1 100644 --- a/docs/databases/bigquery.md +++ b/docs/databases/bigquery.md @@ -48,7 +48,7 @@ config = spec.add_config( "maximum_bytes_billed": 10 * 1024**3, }, driver_features={ - "use_query_and_wait": True, + "enable_query_and_wait": True, "enable_storage_write_api": True, "storage_write_stream_type": "PENDING", }, @@ -64,7 +64,7 @@ config = spec.add_config( | `credentials` / `credentials_path` | `connection_config` | Explicit `google.auth.credentials.Credentials` instance or service-account JSON path (uses ADC when omitted). | | `maximum_bytes_billed` / `priority` / `reservation` | `connection_config` | Cost guardrails (`INTERACTIVE` vs `BATCH` priority and BigQuery Editions reservation assignment). | | `labels` | `connection_config` | Billing and observability key-value labels attached to submitted jobs. | -| `use_query_and_wait` | `driver_features` | Uses low-latency `client.query_and_wait()` for interactive queries and falls back to `QueryJob` when destination tables, transactions, or `EXPORT DATA` require job handles. | +| `enable_query_and_wait` | `driver_features` | Uses low-latency `client.query_and_wait()` for interactive queries and falls back to `QueryJob` when destination tables, transactions, or `EXPORT DATA` require job handles. | | `enable_storage_write_api` / `storage_write_stream_type` | `driver_features` | Streams `load_from_arrow` batches via BigQuery Storage Write API (`"PENDING"` atomic commit default or `"COMMITTED"` immediate visibility). | | `enable_native_storage` | `driver_features` | Accelerates GCS `gs://` Parquet/CSV/JSON exports and loads using server-side BigQuery jobs. | | `request_timeout` / `job_result_timeout` | `driver_features` | Per-RPC and overall job completion timeout bounds in seconds. | diff --git a/docs/databases/index.md b/docs/databases/index.md index 2bac5e867..a2a7fe526 100644 --- a/docs/databases/index.md +++ b/docs/databases/index.md @@ -26,8 +26,8 @@ SQLSpec supports synchronous and asynchronous **database adapters** spanning emb | [`duckdb`](duckdb.md) | `sqlspec[duckdb]` | Sync | `?` (`qmark`) | `DuckDBConnectionPool` | Native Zero-Copy | Table Queue | Sessions + Memory (`FTS`, `ARRAY` Cosine) | | [`bigquery`](bigquery.md) | `sqlspec[bigquery]` | Sync | `@name` (`named`) | Stateless HTTP/gRPC Client | Storage Read / Write API | Table Queue | Sessions + Memory (`SEARCH` / `LIKE`) | | [`spanner`](spanner.md) | `sqlspec[spanner]` | Sync + Async | `@name` (`named`) | `PingingPool` / `TransactionPingingPool` | Row / Batched | Lock-Free Snapshot + `run_in_transaction` | Sessions + Memory (`ARRAY`, FTS `SCORE`, RRF) + Artifacts | -| [`adbc`](adbc.md) | `sqlspec[adbc]` | Sync | `?` / `$1` / `:1` | Driver-Managed (`NoPoolSyncConfig`) | Native C-Stream | Table Queue | Sessions + Memory (PostgreSQL `pgvector` RRF, DuckDB, Oracle, Db2) | -| [`arrow_odbc`](arrow-odbc.md) | `sqlspec[arrow-odbc]` | Sync | `?` (`qmark`) | ODBC Driver Pool (`enable_driver_pooling`) | Native Rust ODBC Batch | Table Queue | Sessions + Memory (`LIKE` / `POSSTR`) | +| [`adbc`](adbc.md) | `sqlspec[adbc]` | Sync | `?` / `$1` / `:1` | Optional PostgreSQL clone pool | Native C-Stream | Table Queue | Sessions + Memory (PostgreSQL `pgvector` RRF, DuckDB, Oracle, Db2) | +| [`arrow_odbc`](arrow-odbc.md) | `sqlspec[arrow-odbc]` | Sync | `?` (`qmark`) | ODBC Driver Pool (`enable_pooling`) | Native Rust ODBC Batch | Table Queue | Sessions + Memory (`LIKE` / `POSSTR`) | ## Database Setup Guides diff --git a/docs/databases/mssql.md b/docs/databases/mssql.md index 54c6f054b..40f1907ed 100644 --- a/docs/databases/mssql.md +++ b/docs/databases/mssql.md @@ -96,7 +96,7 @@ SQLSpec provides two synchronous Microsoft SQL Server and Azure SQL adapters—` | `encrypt` / `trust_server_certificate` / `server_certificate` | `mssql_python` | TLS 1.2/1.3 encryption policy (`strict`, `mandatory`, `optional`) and certificate pinning. | | `application_intent` / `multi_subnet_failover` | `mssql_python` | Always On availability group read-only routing (`ReadOnly` / `ReadWrite`) and parallel subnet failover. | | `pool_min_size` / `pool_max_size` | Both | Connection pool bounds on `MssqlPythonConnectionPool` and `PymssqlConnectionPool`. | -| `use_pool` | `mssql_python` (`driver_features`) | Toggles SQLSpec's managed `MssqlPythonConnectionPool` (`True` by default). | +| `enable_pooling` | `mssql_python`, `pymssql` (`driver_features`) | Toggles SQLSpec's managed connection pooling (`True` by default). | | `encryption` / `tds_version` / `login_timeout` | `pymssql` | FreeTDS transport encryption and TDS protocol negotiation settings. | See [MSSQL Python Reference](/reference/adapters/mssql-python/) and [PyMSSQL Reference](/reference/adapters/pymssql/) for complete TypedDict field documentation. diff --git a/docs/databases/postgres/psycopg.md b/docs/databases/postgres/psycopg.md index a23bde63e..0a50b6757 100644 --- a/docs/databases/postgres/psycopg.md +++ b/docs/databases/postgres/psycopg.md @@ -28,7 +28,7 @@ poetry add "sqlspec[psycopg]" ``` -Choose `psycopg` when your service mixes synchronous CLI/worker processes with `asyncio` HTTP servers, relies on standard `libpq` environment variables and `.pgpass` / GSSAPI / Kerberos authentication, or runs on scale-to-zero serverless platforms where `null_pool=True` avoids idle connections. +Choose `psycopg` when your service mixes synchronous CLI/worker processes with `asyncio` HTTP servers, relies on standard `libpq` environment variables and `.pgpass` / GSSAPI / Kerberos authentication, or runs on scale-to-zero serverless platforms where `enable_pooling=False` avoids idle connections. ## Minimal Connection and Pool Setup @@ -76,7 +76,7 @@ Choose `psycopg` when your service mixes synchronous CLI/worker processes with ` | `min_size` / `max_size` | `connection_config` | Pool bounds passed to `psycopg_pool.ConnectionPool` or `AsyncConnectionPool`. | | `timeout` / `max_waiting` | `connection_config` | Pool checkout timeout (seconds) and maximum queued callers before raising pool timeout errors. | | `max_lifetime` / `max_idle` | `connection_config` | Connection recycle and idle eviction intervals in seconds. | -| `null_pool` | `driver_features` | Instantiates `NullConnectionPool` / `AsyncNullConnectionPool` (zero idle connections while still enforcing `max_size` concurrency). | +| `enable_pooling` | `driver_features` | Instantiates `NullConnectionPool` / `AsyncNullConnectionPool` when `False` (zero idle connections while still enforcing `max_size` concurrency). | | `pooler_mode` | `driver_features` | `"session"`, `"transaction"`, or `"statement"`; sets `prepare_threshold=None` automatically in transaction/statement modes. | | `direct_dsn` | `driver_features` | Direct PostgreSQL URI used by `derive_direct_config()` for migrations and `LISTEN/NOTIFY`. | | `enable_alloydb` / `enable_alloydb_columnar` | `driver_features` | Connects via AlloyDB connector and enables the AlloyDB columnar engine while honoring `prepare_threshold`. | diff --git a/docs/integrations/fastapi.md b/docs/integrations/fastapi.md index a3b75d819..92188a1ab 100644 --- a/docs/integrations/fastapi.md +++ b/docs/integrations/fastapi.md @@ -58,9 +58,11 @@ Configure `extension_config["fastapi"]` (`FastAPIConfig`, which shares `Starlett | `commit_mode` | `"manual" | "autocommit" | "autocommit_include_redirect"` | `"manual"` | Controls automatic transaction commit or rollback based on response status codes. | | `extra_commit_statuses` | `set[int]` | `set()` | Additional HTTP status codes that trigger an automatic commit. | | `extra_rollback_statuses` | `set[int]` | `set()` | Additional HTTP status codes that trigger an automatic rollback. | -| `enable_correlation_middleware` | `bool` | `False` | Extracts or generates `X-Request-ID` correlation IDs into `CorrelationContext`. | +| `enable_correlation_middleware` | `bool` | `True` | Extracts or generates `X-Request-ID` correlation IDs into `CorrelationContext`. | | `enable_sqlcommenter_middleware` | `bool` | `False` | Populates `SQLCommenterContext` with FastAPI route, method, and trace attributes. | +Set `enable_correlation_middleware=False` to stop this plugin from reading request IDs or adding them to response headers. + ## Declarative Filter Dependencies FastAPI shares the same `FilterConfig` builder as Litestar via `sqlspec.extensions.fastapi.providers.create_filter_dependencies`, generating typed `Depends` callables for `LimitOffsetFilter`, `CursorFilter` (returning `HTTP 422` on malformed cursor tokens), `OrderByFilter`, `SearchFilter`, and `BeforeAfterFilter`. diff --git a/docs/integrations/flask.md b/docs/integrations/flask.md index d1181753d..48e0d2722 100644 --- a/docs/integrations/flask.md +++ b/docs/integrations/flask.md @@ -56,9 +56,11 @@ Configure distinct `session_key`, `connection_key`, and `pool_key` values under | `commit_mode` | `"manual" | "autocommit" | "autocommit_include_redirect"` | `"manual"` | Commits on `2xx` (`autocommit`) or `2xx`/`3xx` (`autocommit_include_redirect`) in `after_request` and rolls back on errors. | | `extra_commit_statuses` | `set[int]` | `set()` | Additional HTTP status codes that trigger an automatic commit. | | `extra_rollback_statuses` | `set[int]` | `set()` | Additional HTTP status codes that trigger an automatic rollback. | -| `enable_correlation_middleware` | `bool` | `False` | Binds request correlation headers into `CorrelationContext`. | +| `enable_correlation_middleware` | `bool` | `True` | Binds request correlation headers into `CorrelationContext`. | | `enable_sqlcommenter_middleware` | `bool` | `False` | Populates `SQLCommenterContext` (`framework="flask"`, route endpoint, method) for wire-level SQL comments. | +Set `enable_correlation_middleware=False` to stop this plugin from reading request IDs or adding them to response headers. + ## Recap - Initialize `SQLSpecPlugin` with `plugin.init_app(app)` and access request-scoped sessions via `plugin.get_session()`. diff --git a/docs/integrations/litestar.md b/docs/integrations/litestar.md index e150a520e..be6acab18 100644 --- a/docs/integrations/litestar.md +++ b/docs/integrations/litestar.md @@ -120,7 +120,7 @@ For bidirectional keyset pagination (`CursorPagination`), set `cursor_key="id"` ## Server-Side Session Stores -Supported database adapters expose a Litestar `Store` implementation under `sqlspec.adapters..litestar` (`AsyncpgStore`, `PsycopgSyncStore`, `PsycopgAsyncStore`, `PsqlpyStore`, `CockroachAsyncpgStore`, `CockroachPsycopgSyncStore`, `CockroachPsycopgAsyncStore`, `AiosqliteStore`, `SQLiteStore`, `AsyncmyStore`, `AiomysqlStore`, `PyMysqlStore`, `MysqlConnectorSyncStore`, `MysqlConnectorAsyncStore`, `OracleSyncStore`, `OracleAsyncStore`, `MssqlPythonStore`, `PymssqlStore`, `SybaseStore`, `SqlanywhereStore`, `Db2SyncStore`, `Db2AsyncStore`, `DuckdbStore`, `BigQueryStore`, `SpannerSyncStore`, `SpannerAsyncStore`, `ADBCStore`, and `ArrowOdbcStore`): +Supported database adapters expose a Litestar `Store` implementation under `sqlspec.adapters..litestar` (`AsyncpgStore`, `PsycopgSyncStore`, `PsycopgAsyncStore`, `PsqlpyStore`, `CockroachAsyncpgStore`, `CockroachPsycopgSyncStore`, `CockroachPsycopgAsyncStore`, `AiosqliteStore`, `SqliteStore`, `AsyncmyStore`, `AiomysqlStore`, `PyMysqlStore`, `MysqlConnectorSyncStore`, `MysqlConnectorAsyncStore`, `OracleSyncStore`, `OracleAsyncStore`, `MssqlPythonStore`, `PymssqlStore`, `SybaseStore`, `SqlanywhereStore`, `Db2SyncStore`, `Db2AsyncStore`, `DuckDBStore`, `BigQueryStore`, `SpannerSyncStore`, `SpannerAsyncStore`, `AdbcStore`, and `ArrowOdbcStore`): ```python --8<-- "docs/examples/frameworks/litestar/session_stores.py:example" diff --git a/docs/integrations/starlette.md b/docs/integrations/starlette.md index fd2b149b2..268f10b8c 100644 --- a/docs/integrations/starlette.md +++ b/docs/integrations/starlette.md @@ -54,9 +54,11 @@ Bind multiple async or sync configurations by specifying unique `session_key`, ` | `commit_mode` | `"manual" | "autocommit" | "autocommit_include_redirect"` | `"manual"` | Controls automatic commit/rollback middleware behavior based on response status codes. | | `extra_commit_statuses` | `set[int]` | `set()` | Additional HTTP status codes that trigger an automatic commit. | | `extra_rollback_statuses` | `set[int]` | `set()` | Additional HTTP status codes that trigger an automatic rollback. | -| `enable_correlation_middleware` | `bool` | `False` | Propagates `X-Request-ID` / `X-Correlation-ID` headers into `CorrelationContext`. | +| `enable_correlation_middleware` | `bool` | `True` | Propagates `X-Request-ID` / `X-Correlation-ID` headers into `CorrelationContext`. | | `enable_sqlcommenter_middleware` | `bool` | `False` | Populates `SQLCommenterContext` (`framework="starlette"`, route path, method). | +Set `enable_correlation_middleware=False` to stop this plugin from reading request IDs or adding them to response headers. + ## Recap - Call `db_plugin.init_app(app)` on your `Starlette` instance and retrieve sessions with `db_plugin.get_session(request)`. diff --git a/docs/reference/adapters/adbc.md b/docs/reference/adapters/adbc.md index bfdd4a561..2c8b5627c 100644 --- a/docs/reference/adapters/adbc.md +++ b/docs/reference/adapters/adbc.md @@ -7,6 +7,7 @@ from sqlspec.adapters.adbc import ( AdbcDriver, AdbcDriverFeatures, AdbcExceptionHandler, + AdbcPoolConfig, ) ``` @@ -24,6 +25,12 @@ Apache Arrow Database Connectivity (ADBC) adapter supporting zero-copy Arrow ing separate_signature: false show_signature_annotations: true +::: sqlspec.adapters.adbc.AdbcPoolConfig + options: + heading_level: 3 + separate_signature: false + show_signature_annotations: true + ::: sqlspec.adapters.adbc.AdbcDriverFeatures options: heading_level: 3 @@ -49,7 +56,7 @@ Apache Arrow Database Connectivity (ADBC) adapter supporting zero-copy Arrow ing ## Extension Stores & Configs -::: sqlspec.adapters.adbc.litestar.ADBCStore +::: sqlspec.adapters.adbc.litestar.AdbcStore options: heading_level: 4 diff --git a/docs/reference/adapters/db2.md b/docs/reference/adapters/db2.md index b55cead7f..985e79881 100644 --- a/docs/reference/adapters/db2.md +++ b/docs/reference/adapters/db2.md @@ -85,12 +85,6 @@ Synchronous and asynchronous IBM Db2 LUW adapters built on `ibm_db`. Inherited e ## Extension Stores & Configs -::: sqlspec.adapters.db2.litestar.Db2LitestarConfig - options: - heading_level: 4 - separate_signature: false - show_signature_annotations: true - ::: sqlspec.adapters.db2.litestar.Db2SyncStore options: heading_level: 4 @@ -115,12 +109,6 @@ Synchronous and asynchronous IBM Db2 LUW adapters built on `ibm_db`. Inherited e options: heading_level: 4 -::: sqlspec.adapters.db2.events.Db2EventsConfig - options: - heading_level: 4 - separate_signature: false - show_signature_annotations: true - ::: sqlspec.adapters.db2.events.Db2SyncEventQueueStore options: heading_level: 4 diff --git a/docs/reference/adapters/duckdb.md b/docs/reference/adapters/duckdb.md index 510c4b47d..b5f9b9565 100644 --- a/docs/reference/adapters/duckdb.md +++ b/docs/reference/adapters/duckdb.md @@ -75,27 +75,27 @@ Embedded analytical database adapter built on `duckdb` with zero-copy Apache Arr ## Extension Stores & Configs -::: sqlspec.adapters.duckdb.litestar.DuckdbStore +::: sqlspec.adapters.duckdb.litestar.DuckDBStore options: heading_level: 4 -::: sqlspec.adapters.duckdb.adk.DuckdbADKConfig +::: sqlspec.adapters.duckdb.adk.DuckDBADKConfig options: heading_level: 4 separate_signature: false show_signature_annotations: true -::: sqlspec.adapters.duckdb.adk.DuckdbADKFTSOptions +::: sqlspec.adapters.duckdb.adk.DuckDBADKFTSOptions options: heading_level: 4 separate_signature: false show_signature_annotations: true -::: sqlspec.adapters.duckdb.adk.DuckdbADKStore +::: sqlspec.adapters.duckdb.adk.DuckDBADKStore options: heading_level: 4 -::: sqlspec.adapters.duckdb.adk.DuckdbADKMemoryStore +::: sqlspec.adapters.duckdb.adk.DuckDBADKMemoryStore options: heading_level: 4 diff --git a/docs/reference/adapters/sqlanywhere.md b/docs/reference/adapters/sqlanywhere.md index e0d6c131b..279e9b69b 100644 --- a/docs/reference/adapters/sqlanywhere.md +++ b/docs/reference/adapters/sqlanywhere.md @@ -72,12 +72,6 @@ Synchronous SAP SQL Anywhere 12+ adapter built on `sqlanydb` (`dbcapi` C library options: heading_level: 4 -::: sqlspec.adapters.sqlanywhere.adk.SqlanywhereADKConfig - options: - heading_level: 4 - separate_signature: false - show_signature_annotations: true - ::: sqlspec.adapters.sqlanywhere.adk.SqlanywhereADKStore options: heading_level: 4 @@ -86,6 +80,6 @@ Synchronous SAP SQL Anywhere 12+ adapter built on `sqlanydb` (`dbcapi` C library options: heading_level: 4 -::: sqlspec.adapters.sqlanywhere.events.SqlanywhereSyncEventQueueStore +::: sqlspec.adapters.sqlanywhere.events.SqlanywhereEventQueueStore options: heading_level: 4 diff --git a/docs/reference/adapters/sqlite.md b/docs/reference/adapters/sqlite.md index 6800ca6df..98060fd51 100644 --- a/docs/reference/adapters/sqlite.md +++ b/docs/reference/adapters/sqlite.md @@ -60,7 +60,7 @@ Synchronous SQLite adapter built on Python's standard library `sqlite3` with thr separate_signature: false show_signature_annotations: true -::: sqlspec.adapters.sqlite.litestar.SQLiteStore +::: sqlspec.adapters.sqlite.litestar.SqliteStore options: heading_level: 4 diff --git a/docs/reference/adapters/sybase.md b/docs/reference/adapters/sybase.md index 4c2089776..b92a4327a 100644 --- a/docs/reference/adapters/sybase.md +++ b/docs/reference/adapters/sybase.md @@ -72,12 +72,6 @@ Synchronous SAP Adaptive Server Enterprise (Sybase ASE) 16.0+ adapter built on ` options: heading_level: 4 -::: sqlspec.adapters.sybase.adk.SybaseADKConfig - options: - heading_level: 4 - separate_signature: false - show_signature_annotations: true - ::: sqlspec.adapters.sybase.adk.SybaseADKStore options: heading_level: 4 @@ -86,6 +80,6 @@ Synchronous SAP Adaptive Server Enterprise (Sybase ASE) 16.0+ adapter built on ` options: heading_level: 4 -::: sqlspec.adapters.sybase.events.SybaseSyncEventQueueStore +::: sqlspec.adapters.sybase.events.SybaseEventQueueStore options: heading_level: 4 diff --git a/docs/reference/data-dictionary/dialects/cockroachdb.md b/docs/reference/data-dictionary/dialects/cockroachdb.md index ed9176f38..0bee49140 100644 --- a/docs/reference/data-dictionary/dialects/cockroachdb.md +++ b/docs/reference/data-dictionary/dialects/cockroachdb.md @@ -79,15 +79,16 @@ The `cockroachdb` dialect pack defines `18` named queries across `15` domain `.s ```sql -- name: by_schema -- dialect: cockroachdb - WITH target_schema AS ( + WITH resolved_table_schema AS ( + SELECT rn.nspname::text AS schema_name + FROM pg_catalog.pg_class rc + JOIN pg_catalog.pg_namespace rn ON rn.oid = rc.relnamespace + WHERE rc.oid = pg_catalog.to_regclass(pg_catalog.quote_ident(:table_name::text)) + ), + target_schema AS ( SELECT COALESCE( :schema_name::text, - ( - SELECT rn.nspname::text - FROM pg_catalog.pg_class rc - JOIN pg_catalog.pg_namespace rn ON rn.oid = rc.relnamespace - WHERE rc.oid = pg_catalog.to_regclass(pg_catalog.quote_ident(:table_name::text)) - ), + (SELECT schema_name FROM resolved_table_schema), current_schema() ) AS schema_name ), @@ -160,15 +161,16 @@ The `cockroachdb` dialect pack defines `18` named queries across `15` domain `.s -- name: by_table -- dialect: cockroachdb - WITH target_schema AS ( + WITH resolved_table_schema AS ( + SELECT rn.nspname::text AS schema_name + FROM pg_catalog.pg_class rc + JOIN pg_catalog.pg_namespace rn ON rn.oid = rc.relnamespace + WHERE rc.oid = pg_catalog.to_regclass(pg_catalog.quote_ident(:table_name::text)) + ), + target_schema AS ( SELECT COALESCE( :schema_name::text, - ( - SELECT rn.nspname::text - FROM pg_catalog.pg_class rc - JOIN pg_catalog.pg_namespace rn ON rn.oid = rc.relnamespace - WHERE rc.oid = pg_catalog.to_regclass(pg_catalog.quote_ident(:table_name::text)) - ), + (SELECT schema_name FROM resolved_table_schema), current_schema() ) AS schema_name ), @@ -260,6 +262,36 @@ The `cockroachdb` dialect pack defines `18` named queries across `15` domain `.s ```sql -- name: by_schema -- dialect: cockroachdb + WITH con_cols AS ( + SELECT + kcu.constraint_catalog, + kcu.constraint_schema, + kcu.constraint_name, + kcu.table_schema, + kcu.table_name, + array_agg(kcu.column_name::text ORDER BY kcu.ordinal_position) AS columns + FROM information_schema.key_column_usage kcu + WHERE kcu.table_schema = :schema_name + AND (:table_name::text IS NULL OR kcu.table_name = :table_name) + GROUP BY + kcu.constraint_catalog, + kcu.constraint_schema, + kcu.constraint_name, + kcu.table_schema, + kcu.table_name + ), + ref_cols AS ( + SELECT + pk_kcu.constraint_catalog, + pk_kcu.constraint_schema, + pk_kcu.constraint_name, + array_agg(pk_kcu.column_name::text ORDER BY pk_kcu.ordinal_position) AS referenced_columns + FROM information_schema.key_column_usage pk_kcu + GROUP BY + pk_kcu.constraint_catalog, + pk_kcu.constraint_schema, + pk_kcu.constraint_name + ) SELECT tc.constraint_catalog::text AS constraint_catalog, tc.constraint_schema::text AS constraint_schema, @@ -269,27 +301,17 @@ The `cockroachdb` dialect pack defines `18` named queries across `15` domain `.s tc.constraint_type::text AS constraint_type, tc.is_deferrable::text AS is_deferrable, tc.initially_deferred::text AS initially_deferred, - ARRAY( - SELECT kcu.column_name::text - FROM information_schema.key_column_usage kcu - WHERE kcu.constraint_catalog = tc.constraint_catalog - AND kcu.constraint_schema = tc.constraint_schema - AND kcu.constraint_name = tc.constraint_name - AND kcu.table_schema = tc.table_schema - AND kcu.table_name = tc.table_name - ORDER BY kcu.ordinal_position - )::text[] AS columns, + COALESCE(cc.columns, ARRAY[]::text[]) AS columns, rc.unique_constraint_schema::text AS referenced_schema, pk_tc.table_name::text AS referenced_table, - ARRAY( - SELECT pk_kcu.column_name::text - FROM information_schema.key_column_usage pk_kcu - WHERE pk_kcu.constraint_catalog = rc.unique_constraint_catalog - AND pk_kcu.constraint_schema = rc.unique_constraint_schema - AND pk_kcu.constraint_name = rc.unique_constraint_name - ORDER BY pk_kcu.ordinal_position - )::text[] AS referenced_columns + COALESCE(rc_cols.referenced_columns, ARRAY[]::text[]) AS referenced_columns FROM information_schema.table_constraints tc + LEFT JOIN con_cols cc + ON cc.constraint_catalog = tc.constraint_catalog + AND cc.constraint_schema = tc.constraint_schema + AND cc.constraint_name = tc.constraint_name + AND cc.table_schema = tc.table_schema + AND cc.table_name = tc.table_name LEFT JOIN information_schema.referential_constraints rc ON rc.constraint_catalog = tc.constraint_catalog AND rc.constraint_schema = tc.constraint_schema @@ -298,6 +320,10 @@ The `cockroachdb` dialect pack defines `18` named queries across `15` domain `.s ON pk_tc.constraint_catalog = rc.unique_constraint_catalog AND pk_tc.constraint_schema = rc.unique_constraint_schema AND pk_tc.constraint_name = rc.unique_constraint_name + LEFT JOIN ref_cols rc_cols + ON rc_cols.constraint_catalog = rc.unique_constraint_catalog + AND rc_cols.constraint_schema = rc.unique_constraint_schema + AND rc_cols.constraint_name = rc.unique_constraint_name WHERE tc.table_schema = :schema_name AND (:table_name::text IS NULL OR tc.table_name = :table_name) ORDER BY tc.table_schema, tc.table_name, tc.constraint_name; diff --git a/docs/reference/data-dictionary/dialects/db2.md b/docs/reference/data-dictionary/dialects/db2.md index ae1d7ffdc..9a4d3e983 100644 --- a/docs/reference/data-dictionary/dialects/db2.md +++ b/docs/reference/data-dictionary/dialects/db2.md @@ -16,7 +16,7 @@ Introspects IBM Db2 `SYSCAT` views (`SYSCAT.TABLES`, `SYSCAT.COLUMNS`, `SYSCAT.I | **Default Schema** | `None` | | **Version Pattern** | `(?:v|V)?(\d+)\.(\d+)(?:\.(\d+))?` | | **DDL Strategy** | AST synthesis + `SYSCAT` catalog text definitions | -| **Catalog `.sql` Files** | `9` files (`16` named queries) | +| **Catalog `.sql` Files** | `9` files (`14` named queries) | ### Capabilities & Type Mappings @@ -48,7 +48,7 @@ Introspects IBM Db2 `SYSCAT` views (`SYSCAT.TABLES`, `SYSCAT.COLUMNS`, `SYSCAT.I ## Catalog SQL Queries -The `db2` dialect pack defines `16` named queries across `9` domain `.sql` files loaded via `DataDictionaryLoader` (`sqlspec/data_dictionary/dialects/db2/`): +The `db2` dialect pack defines `14` named queries across `9` domain `.sql` files loaded via `DataDictionaryLoader` (`sqlspec/data_dictionary/dialects/db2/`): | Domain Pack | Named Queries | Bind Parameters | | --- | --- | --- | @@ -57,8 +57,8 @@ The `db2` dialect pack defines `16` named queries across `9` domain `.sql` files | `foreign_keys.sql` | `by_schema`, `by_table` | `:schema_name`, `:table_name` | | `indexes.sql` | `by_schema`, `by_table` | `:schema_name`, `:table_name` | | `objects.sql` | `by_schema` | `:object_name`, `:schema_name` | -| `schemas.sql` | `by_schema`, `current` | `:schema_name` | -| `tables.sql` | `by_schema`, `all_by_schema`, `names_by_schema` | `:schema_name`, `:table_name` | +| `schemas.sql` | `by_schema` | `:schema_name` | +| `tables.sql` | `by_schema`, `all_by_schema` | `:schema_name`, `:table_name` | | `version.sql` | `current` | *none* | | `views.sql` | `by_schema` | `:schema_name`, `:view_name` | @@ -369,7 +369,7 @@ The `db2` dialect pack defines `16` named queries across `9` domain `.sql` files ORDER BY schema_name, object_type, object_name; ``` -??? code "`schemas.sql` (2 named queries)" +??? code "`schemas.sql` (1 named queries)" ```sql -- name: by_schema @@ -385,13 +385,9 @@ The `db2` dialect pack defines `16` named queries across `9` domain `.sql` files AND s.SCHEMANAME NOT LIKE 'SQLJ%' AND (:schema_name IS NULL OR s.SCHEMANAME = :schema_name) ORDER BY s.SCHEMANAME; - - -- name: current - -- dialect: db2 - VALUES CURRENT SCHEMA; ``` -??? code "`tables.sql` (3 named queries)" +??? code "`tables.sql` (2 named queries)" ```sql -- name: by_schema @@ -471,18 +467,6 @@ The `db2` dialect pack defines `16` named queries across `9` domain `.sql` files AND t.TABSCHEMA = COALESCE(CAST(:schema_name AS VARCHAR(128)), CURRENT SCHEMA) AND (CAST(:table_name AS VARCHAR(128)) IS NULL OR t.TABNAME = :table_name) ORDER BY t.TABSCHEMA, t.TABNAME; - - -- name: names_by_schema - -- dialect: db2 - SELECT - RTRIM(t.TABNAME) AS table_name - FROM SYSCAT.TABLES t - WHERE t.TYPE = 'T' - AND t.TABSCHEMA NOT LIKE 'SYS%' - AND t.TABSCHEMA NOT LIKE 'NULLID%' - AND t.TABSCHEMA NOT LIKE 'SQLJ%' - AND t.TABSCHEMA = COALESCE(CAST(:schema_name AS VARCHAR(128)), CURRENT SCHEMA) - ORDER BY t.TABNAME; ``` ??? code "`version.sql` (1 named queries)" diff --git a/docs/reference/data-dictionary/dialects/duckdb.md b/docs/reference/data-dictionary/dialects/duckdb.md index 44505c5d5..c35da754f 100644 --- a/docs/reference/data-dictionary/dialects/duckdb.md +++ b/docs/reference/data-dictionary/dialects/duckdb.md @@ -72,7 +72,7 @@ The `duckdb` dialect pack defines `24` named queries across `17` domain `.sql` f | `schemas.sql` | `list` | *none* | | `sequences.sql` | `by_schema` | `:schema_name` | | `system.sql` | `logs`, `memory`, `settings` | *none* | -| `tables.sql` | `by_schema` | `:schema_name` | +| `tables.sql` | `by_schema` | `:schema_name`, `:table_name` | | `types.sql` | `by_schema` | `:schema_name` | | `version.sql` | `current` | *none* | | `views.sql` | `by_schema` | `:schema_name` | @@ -860,6 +860,7 @@ The `duckdb` dialect pack defines `24` named queries across `17` domain `.sql` f LEFT JOIN dependency_levels dl ON dl.table_name = t.table_name WHERE t.table_type = 'BASE TABLE' AND t.table_schema = COALESCE(:schema_name, current_schema()) + AND (:table_name IS NULL OR t.table_name = :table_name) ORDER BY level, t.table_name; ``` diff --git a/docs/reference/data-dictionary/dialects/index.md b/docs/reference/data-dictionary/dialects/index.md index 1f462c2da..391cfb602 100644 --- a/docs/reference/data-dictionary/dialects/index.md +++ b/docs/reference/data-dictionary/dialects/index.md @@ -16,14 +16,14 @@ SQLSpec ships 13 dialect-specific SQL catalog packs under `sqlspec/data_dictiona | **DuckDB** (`duckdb`) | `None` | Native `duckdb_*().sql` catalog functions + AST synthesis | `17` | `24` | [DuckDB Catalog](duckdb.md) | | **MySQL** (`mysql`) | `None` | Native `SHOW CREATE` + `information_schema` synthesis | `18` | `27` | [MySQL Catalog](mysql.md) | | **MariaDB** (`mariadb`) | `None` | Native `SHOW CREATE` + MariaDB sequence & constraint reflection | `17` | `27` | [MariaDB Catalog](mariadb.md) | -| **Oracle Database** (`oracle`) | `None` | Native `DBMS_METADATA.GET_DDL` + AST fallback | `22` | `35` | [Oracle Database Catalog](oracle.md) | +| **Oracle Database** (`oracle`) | `None` | Native `DBMS_METADATA.GET_DDL` + AST fallback | `22` | `36` | [Oracle Database Catalog](oracle.md) | | **SQL Server** (`mssql`) | `dbo` | AST synthesis + `OBJECT_DEFINITION()` for views/routines/triggers | `19` | `31` | [SQL Server Catalog](mssql.md) | | **Sybase ASE** (`sybase`) | `dbo` | AST synthesis + `dbo.syscomments` object text reassembly | `9` | `16` | [Sybase ASE Catalog](sybase.md) | | **SAP SQL Anywhere** (`sqlanywhere`) | `DBA` | AST synthesis + `SYS.SYSTAB` / `SYS.SYSPROCEDURE` definitions | `9` | `16` | [SAP SQL Anywhere Catalog](sqlanywhere.md) | -| **IBM Db2** (`db2`) | `None` | AST synthesis + `SYSCAT` catalog text definitions | `9` | `16` | [IBM Db2 Catalog](db2.md) | +| **IBM Db2** (`db2`) | `None` | AST synthesis + `SYSCAT` catalog text definitions | `9` | `14` | [IBM Db2 Catalog](db2.md) | | **Google BigQuery** (`bigquery`) | `None` | Native `INFORMATION_SCHEMA.TABLES.ddl` + AST fallback | `14` | `30` | [Google BigQuery Catalog](bigquery.md) | | **Google Cloud Spanner** (`spanner`) | `None` | AST synthesis with `INTERLEAVE IN PARENT` & `ROW DELETION POLICY` (`GoogleSQL` & `PostgreSQL` modes) | `5` | `8` | [Google Cloud Spanner Catalog](spanner.md) | -| **Total (`13` dialects)** | - | - | **`190`** | **`297`** | - | +| **Total (`13` dialects)** | - | - | **`190`** | **`296`** | - | ## Cross-Dialect Capability Matrix diff --git a/docs/reference/data-dictionary/dialects/mysql.md b/docs/reference/data-dictionary/dialects/mysql.md index 753b6109a..c5c0c6f02 100644 --- a/docs/reference/data-dictionary/dialects/mysql.md +++ b/docs/reference/data-dictionary/dialects/mysql.md @@ -85,7 +85,7 @@ The `mysql` dialect pack defines `27` named queries across `18` domain `.sql` fi | `objects.sql` | `by_schema` | `:schema_name` | | `partitions.sql` | `by_schema` | `:schema_name` | | `plugins.sql` | `list` | *none* | -| `privileges.sql` | `by_schema` | `:schema_name` | +| `privileges.sql` | `by_schema` | `:object_name`, `:schema_name` | | `routines.sql` | `by_schema` | `:schema_name` | | `schemas.sql` | `list` | `:schema_name` | | `system.sql` | `performance_schema_tables`, `sys_schema_table_statistics` | `:schema_name` | @@ -450,6 +450,7 @@ The `mysql` dialect pack defines `27` named queries across `18` domain `.sql` fi NULL AS column_name FROM information_schema.table_privileges WHERE table_schema = COALESCE(:schema_name, DATABASE()) + AND (:object_name IS NULL OR table_name = :object_name) UNION ALL SELECT grantee, @@ -461,6 +462,7 @@ The `mysql` dialect pack defines `27` named queries across `18` domain `.sql` fi column_name FROM information_schema.column_privileges WHERE table_schema = COALESCE(:schema_name, DATABASE()) + AND (:object_name IS NULL OR table_name = :object_name) ORDER BY schema_name, table_name, column_name, grantee, privilege_type; ``` diff --git a/docs/reference/data-dictionary/dialects/oracle.md b/docs/reference/data-dictionary/dialects/oracle.md index 491e11569..fcd796888 100644 --- a/docs/reference/data-dictionary/dialects/oracle.md +++ b/docs/reference/data-dictionary/dialects/oracle.md @@ -16,7 +16,7 @@ Introspects `ALL_TABLES`, `ALL_TAB_COLS`, `ALL_INDEXES`, `ALL_IND_COLUMNS`, `ALL | **Default Schema** | `None` | | **Version Pattern** | `(\d+)` | | **DDL Strategy** | Native `DBMS_METADATA.GET_DDL` + AST fallback | -| **Catalog `.sql` Files** | `22` files (`35` named queries) | +| **Catalog `.sql` Files** | `22` files (`36` named queries) | ### Capabilities & Type Mappings @@ -53,7 +53,7 @@ Introspects `ALL_TABLES`, `ALL_TAB_COLS`, `ALL_INDEXES`, `ALL_IND_COLUMNS`, `ALL ## Catalog SQL Queries -The `oracle` dialect pack defines `35` named queries across `22` domain `.sql` files loaded via `DataDictionaryLoader` (`sqlspec/data_dictionary/dialects/oracle/`): +The `oracle` dialect pack defines `36` named queries across `22` domain `.sql` files loaded via `DataDictionaryLoader` (`sqlspec/data_dictionary/dialects/oracle/`): | Domain Pack | Named Queries | Bind Parameters | | --- | --- | --- | @@ -65,7 +65,7 @@ The `oracle` dialect pack defines `35` named queries across `22` domain `.sql` f | `dependencies.sql` | `by_schema` | `:object_name`, `:schema_name` | | `foreign_keys.sql` | `by_table`, `by_schema` | `:schema_name`, `:table_name` | | `grants.sql` | `by_schema` | `:object_name`, `:schema_name` | -| `indexes.sql` | `by_schema`, `by_table` | `:schema_name`, `:table_name` | +| `indexes.sql` | `by_schema`, `by_table`, `expressions` | `:schema_name`, `:table_name` | | `lob_storage.sql` | `by_schema` | `:schema_name`, `:table_name` | | `materialized_views.sql` | `by_schema` | `:mview_name`, `:schema_name` | | `objects.sql` | `by_schema`, `cloud_ai_available` | `:object_name`, `:schema_name` | @@ -425,7 +425,7 @@ The `oracle` dialect pack defines `35` named queries across `22` domain `.sql` f ORDER BY schema_name, object_name, column_name, grantee, privilege; ``` -??? code "`indexes.sql` (2 named queries)" +??? code "`indexes.sql` (3 named queries)" ```sql -- name: by_schema @@ -450,16 +450,11 @@ The `oracle` dialect pack defines `35` named queries across `22` domain `.sql` f i.avg_leaf_blocks_per_key, i.avg_data_blocks_per_key, i.clustering_factor, - LISTAGG(ic.column_name, ',') WITHIN GROUP (ORDER BY ic.column_position) AS columns, - LISTAGG(ie.column_expression, ',') WITHIN GROUP (ORDER BY ie.column_position) AS expressions + LISTAGG(ic.column_name, ',') WITHIN GROUP (ORDER BY ic.column_position) AS columns FROM all_indexes i LEFT JOIN all_ind_columns ic ON ic.index_owner = i.owner AND ic.index_name = i.index_name - LEFT JOIN all_ind_expressions ie - ON ie.index_owner = i.owner - AND ie.index_name = i.index_name - AND ie.column_position = ic.column_position WHERE i.owner = COALESCE(:schema_name, USER) AND (:table_name IS NULL OR i.table_name = :table_name) GROUP BY @@ -487,6 +482,7 @@ The `oracle` dialect pack defines `35` named queries across `22` domain `.sql` f -- name: by_table -- dialect: oracle SELECT + i.owner AS schema_name, i.index_name, i.table_name AS table_name, CASE WHEN i.uniqueness = 'UNIQUE' THEN 1 ELSE 0 END AS is_unique, @@ -505,7 +501,20 @@ The `oracle` dialect pack defines `35` named queries across `22` domain `.sql` f AND c.constraint_type = 'P' WHERE i.table_name = :table_name AND i.owner = COALESCE(:schema_name, USER) - GROUP BY i.index_name, i.table_name, i.uniqueness; + GROUP BY i.owner, i.index_name, i.table_name, i.uniqueness; + + -- name: expressions + -- dialect: oracle + SELECT + ie.index_owner AS schema_name, + ie.table_name, + ie.index_name, + ie.column_position, + ie.column_expression + FROM all_ind_expressions ie + WHERE ie.index_owner = COALESCE(:schema_name, USER) + AND (:table_name IS NULL OR ie.table_name = :table_name) + ORDER BY ie.index_owner, ie.table_name, ie.index_name, ie.column_position; ``` ??? code "`lob_storage.sql` (1 named queries)" @@ -876,12 +885,10 @@ The `oracle` dialect pack defines `35` named queries across `22` domain `.sql` f -- name: current -- dialect: oracle SELECT - product AS "product", - version AS "version", - status AS "status" - FROM product_component_version - WHERE product LIKE 'Oracle%' - ORDER BY TO_NUMBER(REGEXP_SUBSTR(version, '^[0-9]+')) DESC, version DESC + p.* + FROM product_component_version p + WHERE p.product LIKE 'Oracle%' + ORDER BY TO_NUMBER(REGEXP_SUBSTR(p.version, '^[0-9]+')) DESC, p.version DESC FETCH FIRST 1 ROWS ONLY; -- name: autonomous_service diff --git a/docs/reference/data-dictionary/dialects/postgres.md b/docs/reference/data-dictionary/dialects/postgres.md index e7ea7c44e..20abc1d73 100644 --- a/docs/reference/data-dictionary/dialects/postgres.md +++ b/docs/reference/data-dictionary/dialects/postgres.md @@ -257,6 +257,24 @@ The `postgres` dialect pack defines `26` named queries across `21` domain `.sql` ```sql -- name: by_schema -- dialect: postgres + WITH con_cols AS ( + SELECT + con.oid AS con_oid, + array_agg(att.attname::text ORDER BY keys.ord) AS columns + FROM pg_catalog.pg_constraint con + CROSS JOIN LATERAL unnest(con.conkey) WITH ORDINALITY AS keys(attnum, ord) + JOIN pg_catalog.pg_attribute att ON att.attrelid = con.conrelid AND att.attnum = keys.attnum + GROUP BY con.oid + ), + ref_cols AS ( + SELECT + con.oid AS con_oid, + array_agg(att.attname::text ORDER BY keys.ord) AS referenced_columns + FROM pg_catalog.pg_constraint con + CROSS JOIN LATERAL unnest(con.confkey) WITH ORDINALITY AS keys(attnum, ord) + JOIN pg_catalog.pg_attribute att ON att.attrelid = con.confrelid AND att.attnum = keys.attnum + GROUP BY con.oid + ) SELECT current_database()::text AS constraint_catalog, n.nspname::text AS constraint_schema, @@ -272,20 +290,10 @@ The `postgres` dialect pack defines `26` named queries across `21` domain `.sql` ELSE con.contype::text END AS constraint_type, pg_catalog.pg_get_constraintdef(con.oid, true)::text AS definition, - ARRAY( - SELECT att.attname::text - FROM unnest(con.conkey) WITH ORDINALITY AS keys(attnum, ord) - JOIN pg_catalog.pg_attribute att ON att.attrelid = con.conrelid AND att.attnum = keys.attnum - ORDER BY keys.ord - )::text[] AS columns, + COALESCE(cc.columns, ARRAY[]::text[]) AS columns, ref_ns.nspname::text AS referenced_schema, ref.relname::text AS referenced_table, - ARRAY( - SELECT att.attname::text - FROM unnest(con.confkey) WITH ORDINALITY AS keys(attnum, ord) - JOIN pg_catalog.pg_attribute att ON att.attrelid = con.confrelid AND att.attnum = keys.attnum - ORDER BY keys.ord - )::text[] AS referenced_columns, + COALESCE(rc.referenced_columns, ARRAY[]::text[]) AS referenced_columns, con.condeferrable AS is_deferrable, con.condeferred AS initially_deferred, dep.refobjid IS NOT NULL AS is_extension_owned @@ -295,6 +303,8 @@ The `postgres` dialect pack defines `26` named queries across `21` domain `.sql` LEFT JOIN pg_catalog.pg_class ref ON ref.oid = con.confrelid LEFT JOIN pg_catalog.pg_namespace ref_ns ON ref_ns.oid = ref.relnamespace LEFT JOIN pg_catalog.pg_depend dep ON dep.objid = con.oid AND dep.deptype = 'e' + LEFT JOIN con_cols cc ON cc.con_oid = con.oid + LEFT JOIN ref_cols rc ON rc.con_oid = con.oid WHERE n.nspname = :schema_name AND (:table_name::text IS NULL OR c.relname = :table_name) ORDER BY n.nspname, c.relname, con.conname; @@ -505,6 +515,18 @@ The `postgres` dialect pack defines `26` named queries across `21` domain `.sql` ```sql -- name: by_schema -- dialect: postgres + WITH target_schema AS ( + SELECT COALESCE( + :schema_name::text, + ( + SELECT rn.nspname::text + FROM pg_catalog.pg_class rc + JOIN pg_catalog.pg_namespace rn ON rn.oid = rc.relnamespace + WHERE rc.oid = pg_catalog.to_regclass(pg_catalog.quote_ident(:table_name::text)) + ), + current_schema() + ) AS schema_name + ) SELECT tn.nspname::text AS schema_name, tbl.relname::text AS table_name, @@ -540,11 +562,11 @@ The `postgres` dialect pack defines `26` named queries across `21` domain `.sql` FROM pg_catalog.pg_index ix JOIN pg_catalog.pg_class tbl ON tbl.oid = ix.indrelid JOIN pg_catalog.pg_namespace tn ON tn.oid = tbl.relnamespace + JOIN target_schema ts ON ts.schema_name = tn.nspname JOIN pg_catalog.pg_class idx ON idx.oid = ix.indexrelid JOIN pg_catalog.pg_am am ON am.oid = idx.relam LEFT JOIN pg_catalog.pg_depend dep ON dep.objid = idx.oid AND dep.deptype = 'e' - WHERE tn.nspname = :schema_name - AND (:table_name::text IS NULL OR tbl.relname = :table_name) + WHERE (:table_name::text IS NULL OR tbl.relname = :table_name) ORDER BY tn.nspname, tbl.relname, idx.relname; -- name: by_table @@ -786,7 +808,7 @@ The `postgres` dialect pack defines `26` named queries across `21` domain `.sql` total_exec_time, mean_exec_time, rows - FROM pg_catalog.pg_stat_statements + FROM pg_stat_statements ORDER BY total_exec_time DESC LIMIT :limit; diff --git a/docs/reference/data-dictionary/dialects/sqlanywhere.md b/docs/reference/data-dictionary/dialects/sqlanywhere.md index 22c4fc4c6..e3e8de913 100644 --- a/docs/reference/data-dictionary/dialects/sqlanywhere.md +++ b/docs/reference/data-dictionary/dialects/sqlanywhere.md @@ -35,7 +35,7 @@ Introspects SAP SQL Anywhere `SYS` catalog views (`SYS.SYSTAB`, `SYS.SYSTABCOL`, | `supports_schemas` | `True` | `json` | `LONG VARCHAR` | | `supports_sequences` | `True` | `smallint` | `SMALLINT` | | `supports_skip_locked` | `True` | `string` | `VARCHAR(255)` | -| `supports_transactional_ddl` | `True` | `text` | `LONG VARCHAR` | +| `supports_transactional_ddl` | `False` | `text` | `LONG VARCHAR` | | `supports_transactions` | `True` | `time` | `TIME` | | `supports_update_from` | `True` | `timestamp` | `TIMESTAMP` | | `supports_upsert` | `True` | `uuid` | `UNIQUEIDENTIFIER` | diff --git a/docs/reference/data-dictionary/dialects/sqlite.md b/docs/reference/data-dictionary/dialects/sqlite.md index b66c029b1..e17f0de6b 100644 --- a/docs/reference/data-dictionary/dialects/sqlite.md +++ b/docs/reference/data-dictionary/dialects/sqlite.md @@ -107,7 +107,7 @@ The `sqlite` dialect pack defines `23` named queries across `15` domain `.sql` f -- name: by_table -- dialect: sqlite SELECT - :schema_name AS schema_name, + COALESCE(:schema_name, 'main') AS schema_name, :table_name AS table_name, ti.name AS column_name, ti.type AS data_type, @@ -128,7 +128,7 @@ The `sqlite` dialect pack defines `23` named queries across `15` domain `.sql` f END AS extra, ti.hidden IN (2, 3) AS is_generated, CASE WHEN ti.hidden = 1 THEN 1 ELSE 0 END AS is_hidden - FROM pragma_table_xinfo(:table_name, :schema_name) AS ti + FROM pragma_table_xinfo(:table_name, COALESCE(:schema_name, 'main')) AS ti LEFT JOIN {schema_prefix}sqlite_schema AS sm ON sm.type = 'table' AND sm.name = :table_name ORDER BY ti.cid; ``` @@ -348,7 +348,7 @@ The `sqlite` dialect pack defines `23` named queries across `15` domain `.sql` f -- name: by_table -- dialect: sqlite SELECT - :schema_name AS schema_name, + COALESCE(:schema_name, 'main') AS schema_name, :table_name AS table_name, il.name AS index_name, il."unique" AS is_unique, @@ -362,8 +362,8 @@ The `sqlite` dialect pack defines `23` named queries across `15` domain `.sql` f ix.coll AS collation, ix.key AS is_key_column, sm.sql AS index_sql - FROM pragma_index_list(:table_name, :schema_name) AS il - JOIN pragma_index_xinfo(il.name, :schema_name) AS ix + FROM pragma_index_list(:table_name, COALESCE(:schema_name, 'main')) AS il + JOIN pragma_index_xinfo(il.name, COALESCE(:schema_name, 'main')) AS ix LEFT JOIN {schema_prefix}sqlite_schema AS sm ON sm.type = 'index' AND sm.name = il.name @@ -373,7 +373,7 @@ The `sqlite` dialect pack defines `23` named queries across `15` domain `.sql` f -- name: columns_by_index -- dialect: sqlite SELECT - :schema_name AS schema_name, + COALESCE(:schema_name, 'main') AS schema_name, :index_name AS index_name, ix.seqno, ix.cid AS column_id, @@ -381,7 +381,7 @@ The `sqlite` dialect pack defines `23` named queries across `15` domain `.sql` f ix.desc AS is_descending, ix.coll AS collation, ix.key AS is_key_column - FROM pragma_index_xinfo(:index_name, :schema_name) AS ix + FROM pragma_index_xinfo(:index_name, COALESCE(:schema_name, 'main')) AS ix ORDER BY ix.seqno; -- name: index_columns_by_index diff --git a/pyproject.toml b/pyproject.toml index 6fd2f187f..10ee865c2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -113,6 +113,7 @@ extras = [ "adlfs", "gcsfs", "dishka", + "sqlanydb>=1.0.14", ] lint = [ "mypy>=2.0.0", diff --git a/sqlspec/_typing.py b/sqlspec/_typing.py index 257601710..a92e14a71 100644 --- a/sqlspec/_typing.py +++ b/sqlspec/_typing.py @@ -38,6 +38,7 @@ from pandas import DataFrame as PandasDataFrame from polars import DataFrame as PolarsDataFrame from prometheus_client import Counter, Gauge, Histogram + from pyarrow import ArrowException from pyarrow import RecordBatch as ArrowRecordBatch from pyarrow import RecordBatchReader as ArrowRecordBatchReader from pyarrow import Schema as ArrowSchema @@ -67,6 +68,7 @@ "UNSET", "UNSET_STUB", "UUID_UTILS_INSTALLED", + "ArrowException", "ArrowRecordBatch", "ArrowRecordBatchReader", "ArrowRecordBatchReaderProtocol", @@ -472,6 +474,13 @@ def read_next_batch(self) -> Any: _ARROW_RECORD_BATCH_READER_SHIM = ArrowRecordBatchReaderProtocol +class ArrowExceptionStubError(Exception): + """Placeholder for ``pyarrow.ArrowException`` that no runtime error is an instance of.""" + + +_ARROW_EXCEPTION_SHIM = ArrowExceptionStubError + + @runtime_checkable class PandasDataFrameProtocol(Protocol): """Typed shim for pandas.DataFrame.""" @@ -711,6 +720,7 @@ def labels(self, *labelvalues: str, **labelkwargs: str) -> "_MetricInstance": _CATTRS_UNSTRUCTURE_SHIM.__name__ = "cattrs_unstructure" _CATTRS_UNSTRUCTURE_SHIM.__qualname__ = "cattrs_unstructure" _LAZY_EXPORTS: "dict[str, tuple[str, str | None, Any]]" = { + "ArrowException": ("pyarrow", "ArrowException", _ARROW_EXCEPTION_SHIM), "ArrowRecordBatch": ("pyarrow", "RecordBatch", _ARROW_RECORD_BATCH_SHIM), "ArrowRecordBatchReader": ("pyarrow", "RecordBatchReader", _ARROW_RECORD_BATCH_READER_SHIM), "ArrowSchema": ("pyarrow", "Schema", _ARROW_SCHEMA_SHIM), diff --git a/sqlspec/adapters/adbc/__init__.py b/sqlspec/adapters/adbc/__init__.py index 34c638acd..3eefd962b 100644 --- a/sqlspec/adapters/adbc/__init__.py +++ b/sqlspec/adapters/adbc/__init__.py @@ -1,7 +1,7 @@ """ADBC adapter for SQLSpec.""" from sqlspec.adapters.adbc._typing import AdbcConnection -from sqlspec.adapters.adbc.config import AdbcConfig, AdbcConnectionParams, AdbcDriverFeatures +from sqlspec.adapters.adbc.config import AdbcConfig, AdbcConnectionParams, AdbcDriverFeatures, AdbcPoolConfig from sqlspec.adapters.adbc.driver import AdbcCursor, AdbcDriver, AdbcExceptionHandler __all__ = ( @@ -12,4 +12,5 @@ "AdbcDriver", "AdbcDriverFeatures", "AdbcExceptionHandler", + "AdbcPoolConfig", ) diff --git a/sqlspec/adapters/adbc/_typing.py b/sqlspec/adapters/adbc/_typing.py index ca995dfd2..d9eba6804 100644 --- a/sqlspec/adapters/adbc/_typing.py +++ b/sqlspec/adapters/adbc/_typing.py @@ -7,6 +7,7 @@ import contextlib from typing import TYPE_CHECKING, Any, TypeAlias +from adbc_driver_manager import AdbcStatusCode from adbc_driver_manager import Error as _AdbcNativeError from adbc_driver_manager import NotSupportedError as AdbcNotSupportedError from adbc_driver_manager import OperationalError as AdbcOperationalError @@ -38,6 +39,7 @@ "AdbcOperationalError", "AdbcRawCursor", "AdbcSessionContext", + "AdbcStatusCode", ) @@ -75,6 +77,7 @@ class AdbcSessionContext: "_acquire_connection", "_connection", "_driver", + "_driver_class", "_driver_features", "_prepare_driver", "_release_connection", @@ -88,34 +91,69 @@ def __init__( statement_config: "StatementConfig", driver_features: "dict[str, Any]", prepare_driver: "Callable[[AdbcDriver], AdbcDriver]", + driver_class: "type[AdbcDriver] | None" = None, ) -> None: self._acquire_connection = acquire_connection self._release_connection = release_connection self._statement_config = statement_config self._driver_features = driver_features self._prepare_driver = prepare_driver + self._driver_class = driver_class self._connection: Any = None self._driver: AdbcDriver | None = None def __enter__(self) -> "AdbcDriver": - from sqlspec.adapters.adbc.driver import AdbcDriver + driver_class = self._driver_class + if driver_class is None: + from sqlspec.adapters.adbc.driver import AdbcDriver + + driver_class = AdbcDriver self._connection = self._acquire_connection() - resolved_dialect = resolve_dialect_name(self._statement_config.dialect) or None - self._driver = AdbcDriver( - connection=self._connection, - statement_config=self._statement_config, - driver_features=self._driver_features, - dialect=resolved_dialect, - ) - return self._prepare_driver(self._driver) + try: + resolved_dialect = resolve_dialect_name(self._statement_config.dialect) or None + self._driver = driver_class( + connection=self._connection, + statement_config=self._statement_config, + driver_features=self._driver_features, + dialect=resolved_dialect, + ) + return self._prepare_driver(self._driver) + except BaseException as exc: + connection = self._connection + driver = self._driver + self._connection = None + self._driver = None + if driver is not None: + with contextlib.suppress(BaseException): + driver.close_streams() + with contextlib.suppress(BaseException): + self._release_connection( + connection, exc_type=type(exc), exc_val=exc, exc_tb=exc.__traceback__, driver=driver + ) + raise def __exit__( self, exc_type: "type[BaseException] | None", exc_val: "BaseException | None", exc_tb: "TracebackType | None" ) -> "bool | None": if self._connection is not None: - self._release_connection(self._connection, exc_type=exc_type, exc_val=exc_val, exc_tb=exc_tb) + connection = self._connection + driver = self._driver self._connection = None + self._driver = None + first_error: BaseException | None = None + if driver is not None: + try: + driver.close_streams() + except BaseException as exc: + first_error = exc + try: + self._release_connection(connection, exc_type=exc_type, exc_val=exc_val, exc_tb=exc_tb, driver=driver) + except BaseException as exc: + if first_error is None: + first_error = exc + if exc_type is None and first_error is not None: + raise first_error return None diff --git a/sqlspec/adapters/adbc/adk/store.py b/sqlspec/adapters/adbc/adk/store.py index b948ef2f0..2f19b9702 100644 --- a/sqlspec/adapters/adbc/adk/store.py +++ b/sqlspec/adapters/adbc/adk/store.py @@ -8,11 +8,14 @@ from typing import TYPE_CHECKING, Any, Final, Literal, NotRequired, cast from sqlspec.adapters.adbc.core import ( + _DB2_INDEX_EXISTS_SQL, + _DB2_TABLE_EXISTS_SQL, driver_kind_from_driver_name, is_postgres_dialect, is_shared_object_driver, resolve_dialect_from_config, ) +from sqlspec.adapters.adbc.core import split_db2_table_name as _split_db2_table_name from sqlspec.config import ADKConfig from sqlspec.core.parameters._validator import PARAMETER_REGEX from sqlspec.exceptions import ImproperConfigurationError @@ -20,7 +23,6 @@ from sqlspec.extensions.adk.memory.store import BaseSyncADKMemoryStore from sqlspec.utils.logging import get_logger from sqlspec.utils.serializers import from_json, to_json -from sqlspec.utils.text import normalize_identifier, split_qualified_identifier if TYPE_CHECKING: from collections.abc import Sequence @@ -49,7 +51,7 @@ class AdbcADKConfig(ADKConfig): """Vector index type for PostgreSQL. Defaults to 'hnsw'.""" scann_num_leaves: NotRequired[int] - """Number of leaves for ScaNN index on PostgreSQL. Defaults to 50.""" + """Number of leaves for ScaNN index on PostgreSQL. Defaults to 100.""" scann_quantizer: NotRequired[Literal["SQ8", "FP32"]] """Quantizer for ScaNN index on PostgreSQL. Defaults to 'SQ8'.""" @@ -224,6 +226,7 @@ def set_metadata(self, key: str, value: str) -> None: """Set a metadata value.""" self._set_metadata(key, value) + @property def dialect(self) -> str: """Return the detected database dialect.""" return self._dialect @@ -2194,7 +2197,23 @@ def _insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "objec with self._config.provide_connection() as conn: cursor = conn.cursor() try: - for entry in entries: + entries_to_insert = entries + if not use_returning: + event_ids = [entry["event_id"] for entry in entries if "event_id" in entry] + if event_ids: + placeholders = ", ".join("?" for _ in event_ids) + select_sql = f"SELECT event_id FROM {self._memory_table} WHERE event_id IN ({placeholders})" + self._execute(cursor, select_sql, tuple(event_ids)) + existing_rows = cursor.fetchall() + existing_event_ids = {row[0] for row in existing_rows} + entries_to_insert = [] + for entry in entries: + event_id = entry["event_id"] + if event_id not in existing_event_ids: + entries_to_insert.append(entry) + existing_event_ids.add(event_id) + + for entry in entries_to_insert: content_json = self._serialize_json_field(entry["content_json"]) metadata_json = self._serialize_json_field(entry.get("metadata_json")) raw_emb = entry.get("embedding") @@ -2272,19 +2291,25 @@ def _insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "objec else: try: self._execute(cursor, sql, params) - inserted_count += 1 except Exception as exc: - exc_str = str(exc).lower() - if ( - "unique" in exc_str - or "constraint" in exc_str - or "duplicate" in exc_str - or "ora-00001" in exc_str - or "sql0803n" in exc_str - or "23505" in exc_str + error_message = str(exc).lower() + if not ( + getattr(exc, "sqlstate", None) == "23505" + or "unique" in error_message + or "duplicate" in error_message + or "ora-00001" in error_message + or "sql0803n" in error_message ): - continue - raise + raise + self._execute( + cursor, + f"SELECT event_id FROM {self._memory_table} WHERE event_id = ?", + (entry["event_id"],), + ) + if cursor.fetchone() is None: + raise + continue + inserted_count += 1 conn.commit() except Exception: with contextlib.suppress(Exception): @@ -2622,15 +2647,6 @@ def decode_adbc_pgvector(raw_embedding: Any) -> "list[float] | None": return [float(v) for v in raw_embedding] -_DB2_TABLE_EXISTS_SQL: Final = ( - "SELECT 1 FROM SYSCAT.TABLES WHERE TABSCHEMA = COALESCE(CAST(? AS VARCHAR(128)), CURRENT SCHEMA) AND TABNAME = ?" -) -_DB2_INDEX_EXISTS_SQL: Final = ( - "SELECT 1 FROM SYSCAT.INDEXES WHERE INDSCHEMA = COALESCE(CAST(? AS VARCHAR(128)), CURRENT SCHEMA) AND INDNAME = ?" -) -_MIN_QUALIFIED_PARTS: Final = 2 - - @lru_cache(maxsize=256) def _postgresql_store_sql(sql: str) -> str: """Cache positional SQL conversion without changing quoted or commented text.""" @@ -2669,13 +2685,6 @@ def _metadata_key_column(dialect: str) -> str: return "key" -def _split_db2_table_name(table_name: str) -> "tuple[str | None, str]": - parts = [part.strip() for part in split_qualified_identifier(table_name, quote_chars='"') if part.strip()] - if len(parts) >= _MIN_QUALIFIED_PARTS: - return normalize_identifier(parts[-2], "db2"), normalize_identifier(parts[-1], "db2") - return None, normalize_identifier(parts[0] if parts else table_name, "db2") - - def _db2_table_exists(cursor: Any, table_name: str) -> bool: schema_name, catalog_table = _split_db2_table_name(table_name) cursor.execute(_DB2_TABLE_EXISTS_SQL, (schema_name, catalog_table)) diff --git a/sqlspec/adapters/adbc/config.py b/sqlspec/adapters/adbc/config.py index 18fe38548..81ed57538 100644 --- a/sqlspec/adapters/adbc/config.py +++ b/sqlspec/adapters/adbc/config.py @@ -1,8 +1,12 @@ """ADBC database configuration.""" +import contextlib +import importlib.metadata +import math +import re from collections.abc import Callable from pathlib import Path -from typing import TYPE_CHECKING, Any, ClassVar, Literal, NotRequired, TypedDict, cast +from typing import TYPE_CHECKING, Any, ClassVar, Literal, NotRequired, Self, TypedDict, cast from sqlspec.adapters.adbc._typing import AdbcConnection, AdbcCursor, AdbcSessionContext from sqlspec.adapters.adbc.core import ( @@ -18,9 +22,11 @@ resolve_dialect_from_config, resolve_dialect_name, resolve_driver_connect_func, + resolve_driver_name_from_config, ) from sqlspec.adapters.adbc.driver import AdbcDriver, AdbcExceptionHandler -from sqlspec.config import ExtensionConfigs, NoPoolSyncConfig +from sqlspec.adapters.adbc.pool import AdbcConnectionPool +from sqlspec.config import ExtensionConfigs, SyncDatabaseConfig from sqlspec.core import StatementConfig from sqlspec.core.capabilities import TypeCoercionCapabilities from sqlspec.core.config_runtime import ( @@ -30,7 +36,7 @@ resolve_runtime_statement_config, ) from sqlspec.driver import SyncPoolConnectionContext, SyncPoolSessionFactory -from sqlspec.exceptions import ImproperConfigurationError +from sqlspec.exceptions import DatabaseConnectionError, ImproperConfigurationError from sqlspec.extensions.events import EventRuntimeHints from sqlspec.utils.config_tools import normalize_connection_config @@ -40,7 +46,16 @@ from sqlspec.observability import ObservabilityConfig -__all__ = ("AdbcConfig", "AdbcConnectionParams", "AdbcDriverFeatures") +__all__ = ("AdbcConfig", "AdbcConnectionParams", "AdbcDriverFeatures", "AdbcPoolConfig") + + +class AdbcPoolConfig(TypedDict): + """Optional PostgreSQL clone-pool bounds; an empty mapping enables defaults.""" + + max_size: NotRequired[int] + """Maximum simultaneous clone leases, excluding the original seed (default 4).""" + timeout: NotRequired[float] + """Maximum seconds to wait for a lease (default 30.0).""" class AdbcConnectionParams(TypedDict): @@ -149,12 +164,8 @@ class AdbcDriverFeatures(TypedDict): """Detect explicit SQL casts such as ``::jsonb`` or ``::vector`` during parameter preparation (defaults to ``True``).""" enable_strict_type_coercion: NotRequired[bool] """Raise an error on unsupported parameter types instead of best-effort conversion (defaults to ``False``).""" - strict_type_coercion: NotRequired[bool] - """Alias normalized to ``enable_strict_type_coercion``.""" enable_arrow_extension_types: NotRequired[bool] """Preserve PyArrow extension type metadata when materializing Arrow results (defaults to ``True``).""" - arrow_extension_types: NotRequired[bool] - """Alias normalized to ``enable_arrow_extension_types``.""" enable_pgvector: NotRequired[bool] """Probe ``pg_extension`` for ``vector`` on PostgreSQL ADBC connections (defaults to ``True`` when ``pgvector`` is installed).""" enable_paradedb: NotRequired[bool] @@ -162,7 +173,7 @@ class AdbcDriverFeatures(TypedDict): enable_pg_textsearch: NotRequired[bool] """Probe ``pg_extension`` for ``pg_textsearch`` BM25 support on PostgreSQL ADBC connections (defaults to ``True``).""" enable_lowercase_column_names: NotRequired[bool] - """Normalize unquoted uppercase column names to lowercase (defaults to ``True`` on Oracle and Db2 ADBC backends).""" + """Normalize unquoted uppercase column names to lowercase (defaults to ``True`` on Oracle, Db2, and Snowflake ADBC backends).""" enable_events: NotRequired[bool] """Enable SQLSpec table-backed event channel support.""" on_connection_create: "NotRequired[Callable[[AdbcConnection], None]]" @@ -197,45 +208,63 @@ def __exit__( class _AdbcSessionConnectionHandler(SyncPoolSessionFactory): - __slots__ = ("_connection", "_driver_features", "_profile_path") + __slots__ = ("_connection", "_driver_features", "_origin_pool", "_profile_path") def __init__(self, config: "AdbcConfig") -> None: super().__init__(config) self._connection: AdbcConnection | None = None self._driver_features: dict[str, Any] = dict(config.driver_features) self._profile_path: Path | None = None + self._origin_pool: AdbcConnectionPool | None = None @property def driver_features(self) -> "dict[str, Any]": return self._driver_features def acquire_connection(self) -> "AdbcConnection": - self._connection = self._config.create_connection() - if self._config._resolved_dialect == "duckdb" and bool( - self._driver_features.get("enable_statement_profiling", False) - ): - self._profile_path = create_duckdb_profile_path() - configure_adbc_duckdb_profiling(self._connection, self._profile_path) - self._driver_features["_duckdb_profile_path"] = str(self._profile_path) - return cast("AdbcConnection", self._connection) + self._origin_pool = self._config.provide_pool() + connection: AdbcConnection = ( + self._origin_pool.acquire() if self._origin_pool is not None else self._config.create_connection() + ) + self._connection = connection + try: + if self._config._resolved_dialect == "duckdb" and bool( + self._driver_features.get("enable_statement_profiling", False) + ): + self._profile_path = create_duckdb_profile_path() + configure_adbc_duckdb_profiling(connection, self._profile_path) + self._driver_features["_duckdb_profile_path"] = str(self._profile_path) + except BaseException as exc: + with contextlib.suppress(BaseException): + self.release_connection(connection, exc_type=type(exc), exc_val=exc, exc_tb=exc.__traceback__) + raise + return connection def release_connection(self, _conn: "AdbcConnection", **kwargs: Any) -> None: if self._connection is None: return + connection = self._connection + pool, self._origin_pool = self._origin_pool, None + profile_path = self._profile_path + self._connection = None + self._profile_path = None + raw_path = self._driver_features.pop("_duckdb_profile_path", None) try: - if self._connection is not self._config.connection_instance: - self._connection.close() + if pool is not None: + driver = kwargs.get("driver") + discard = kwargs.get("exc_type") is not None or bool(driver is not None and driver.has_streams) + pool.release(connection, discard=discard) + elif connection is not self._config.connection_instance: + connection.close() finally: - self._connection = None - if self._profile_path is not None: - cleanup_duckdb_profile_path(self._profile_path) - self._profile_path = None - raw_path = self._driver_features.pop("_duckdb_profile_path", None) - if isinstance(raw_path, str): - cleanup_duckdb_profile_path(Path(raw_path)) + try: + cleanup_duckdb_profile_path(profile_path) + finally: + if isinstance(raw_path, str): + cleanup_duckdb_profile_path(Path(raw_path)) -class AdbcConfig(NoPoolSyncConfig[AdbcConnection, AdbcDriver]): +class AdbcConfig(SyncDatabaseConfig[AdbcConnection, AdbcConnectionPool | AdbcConnection | None, AdbcDriver]): """ADBC configuration for Arrow Database Connectivity. ADBC provides an interface for connecting to multiple database systems @@ -268,6 +297,8 @@ class AdbcConfig(NoPoolSyncConfig[AdbcConnection, AdbcDriver]): "_paradedb_available", "_pg_textsearch_available", "_pgvector_available", + "_pool_close_started", + "_pool_config", "_resolved_dialect", "_user_connection_hook", ) @@ -277,6 +308,7 @@ def __init__( *, connection_config: "AdbcConnectionParams | dict[str, Any] | None" = None, connection_instance: "AdbcConnection | None" = None, + pool_config: "AdbcPoolConfig | None" = None, migration_config: "dict[str, Any] | None" = None, statement_config: StatementConfig | None = None, driver_features: "AdbcDriverFeatures | dict[str, Any] | None" = None, @@ -290,6 +322,7 @@ def __init__( Args: connection_config: Connection configuration parameters connection_instance: Pre-created connection instance to use instead of creating new one + pool_config: Optional PostgreSQL clone pooling; None preserves unpooled ownership. migration_config: Migration configuration statement_config: Default SQL statement configuration driver_features: Driver feature configuration (AdbcDriverFeatures) @@ -311,12 +344,15 @@ def __init__( explicit_lowercase = bool(driver_features and "enable_lowercase_column_names" in driver_features) statement_config, driver_features = apply_driver_features(statement_config, driver_features) features_dict = dict(driver_features) if driver_features else {} - if not explicit_lowercase and self._resolved_dialect in {"oracle", "db2"}: + if not explicit_lowercase and self._resolved_dialect in {"oracle", "db2", "snowflake"}: features_dict["enable_lowercase_column_names"] = True self._user_connection_hook = cast( "Callable[[AdbcConnection], None] | None", features_dict.pop("on_connection_create", None) ) self._default_session_config = get_statement_config(self._resolved_dialect) + self._pool_config = cast("AdbcPoolConfig", dict(pool_config)) if pool_config is not None else None + self._pool_close_started = False + self._validate_pool_settings(connection_instance) super().__init__( connection_config=self.connection_config, @@ -330,6 +366,135 @@ def __init__( **kwargs, ) + def create_pool(self) -> "AdbcConnectionPool | None": + """Publish one lazy generation, rejecting a drain until every native owner closes.""" + if self._pool_config is None: + return None + while True: + existing = self.connection_instance + if isinstance(existing, AdbcConnectionPool): + existing.check_pid() + with self._pool_lock: + if self.connection_instance is not existing: + continue + if existing is not None and not isinstance(existing, AdbcConnectionPool): + msg = "A pooled ADBC configuration cannot own an injected raw connection." + raise ImproperConfigurationError(msg) + if existing is not None and not existing.is_closed: + if self._pool_close_started: + msg = "ADBC pool is draining; retry shutdown before starting another generation." + raise DatabaseConnectionError(msg) + return existing + if existing is not None: + self.connection_instance = None + self._pool_close_started = False + created = None + else: + created = self._create_pool() + self.connection_instance = created + self._pool_close_started = False + runtime = self.get_observability_runtime() + if existing is not None: + runtime.emit_pool_destroy_sync(existing) + continue + runtime.emit_pool_create_sync(created) + return created + + def provide_pool(self, *args: Any, **kwargs: Any) -> "AdbcConnectionPool | None": + """Return the current enabled pool generation, or None when pooling is disabled.""" + return self.create_pool() + + def close_pool(self) -> None: + """Drain an owned generation, retaining native cleanup failures for retry.""" + self._shutdown_pool(retire=True) + + def derive_pool(self, **overrides: Any) -> Self: + """Derive independent ownership while preserving explicit pool settings.""" + pool_config = overrides.pop("pool_config", self._pool_config) + derived = super().derive_pool(**overrides) + derived._pool_config = cast("AdbcPoolConfig", dict(pool_config)) if pool_config is not None else None + derived._validate_pool_settings(None) + return derived + + def _create_pool(self) -> "AdbcConnectionPool | None": + if self._pool_config is None: + return None + return AdbcConnectionPool(self.create_connection, **self._pool_config) + + def _close_pool(self) -> None: + self._shutdown_pool(retire=False) + + def _shutdown_pool(self, *, retire: bool) -> None: + if self._pool_config is None: + return + while True: + pool = self.connection_instance + if pool is None: + return + if not isinstance(pool, AdbcConnectionPool): + msg = "Pooled ADBC shutdown requires an owned pool generation." + raise ImproperConfigurationError(msg) + pool.check_pid() + with self._pool_lock: + if self.connection_instance is not pool: + continue + first_close = not self._pool_close_started + self._pool_close_started = True + break + runtime = self.get_observability_runtime() + if first_close: + runtime.emit_pool_destroying_sync(pool) + pool.close() + if retire and pool.is_closed: + with self._pool_lock: + retired = self.connection_instance is pool + if retired: + self.connection_instance = None + self._pool_close_started = False + if retired: + runtime.emit_pool_destroy_sync(pool) + + def _validate_pool_settings(self, connection_instance: "AdbcConnection | None") -> None: + if self._pool_config is None: + return + max_size = self._pool_config.get("max_size", 4) + timeout = self._pool_config.get("timeout", 30.0) + if isinstance(max_size, bool) or not isinstance(max_size, int) or max_size <= 0: + msg = "ADBC pool max_size must be a positive integer." + raise ImproperConfigurationError(msg) + if not math.isfinite(timeout) or timeout < 0: + msg = "ADBC pool timeout must be finite and nonnegative." + raise ImproperConfigurationError(msg) + if resolve_driver_name_from_config(self.connection_config) != "adbc_driver_postgresql.dbapi.connect": + msg = "ADBC pooling supports only the native PostgreSQL driver." + raise ImproperConfigurationError(msg) + conn_kwargs = self.connection_config.get("conn_kwargs") or {} + native_autocommit = conn_kwargs.get("adbc.connection.autocommit", False) + if self.connection_config.get("autocommit", False) or native_autocommit not in (False, "false", "0", 0): + msg = "ADBC pooling requires autocommit=False." + raise ImproperConfigurationError(msg) + if connection_instance is not None or self._user_connection_hook is not None: + msg = "ADBC pooling does not support injected connections or connection initialization hooks." + raise ImproperConfigurationError(msg) + for distribution in ("adbc-driver-manager", "adbc-driver-postgresql"): + try: + installed = importlib.metadata.version(distribution) + except importlib.metadata.PackageNotFoundError as exc: + msg = f"ADBC pooling requires {distribution}>=1.12.0." + raise ImproperConfigurationError(msg) from exc + match = re.fullmatch( + r"(\d+)\.(\d+)(?:\.(\d+))?((?:(?:a|b|rc)\d+)?(?:\.post\d+)?(?:\.dev\d+)?(?:\+[a-zA-Z0-9.-]+)?)", + installed, + ) + release = tuple(int(part or 0) for part in match.groups()[:3]) if match is not None else (0, 0, 0) + suffix = match.group(4) if match is not None else "" + if release < (1, 12, 0) or (release == (1, 12, 0) and suffix.startswith(("a", "b", "rc", ".dev"))): + msg = f"ADBC pooling requires {distribution}>=1.12.0; installed version is {installed}." + raise ImproperConfigurationError(msg) + if not callable(getattr(AdbcConnection, "adbc_clone", None)): + msg = "ADBC pooling requires the public Connection.adbc_clone capability." + raise ImproperConfigurationError(msg) + @property def supports_migration_schemas(self) -> bool: # type: ignore[override] """Migration schema support is only available for PostgreSQL-backed ADBC drivers.""" @@ -349,9 +514,10 @@ def create_connection(self) -> AdbcConnection: Raises: ImproperConfigurationError: If the connection could not be established. """ - if self.connection_instance is not None: - return self.connection_instance + if self._pool_config is None and self.connection_instance is not None: + return cast("AdbcConnection", self.connection_instance) + connection: AdbcConnection | None = None try: driver_name = cast( "str | None", self.connection_config.get("driver_name") or self.connection_config.get("driver") @@ -365,7 +531,12 @@ def create_connection(self) -> AdbcConnection: if self._user_connection_hook is not None: self._user_connection_hook(cast("AdbcConnection", connection)) return cast("AdbcConnection", connection) - except Exception as e: + except BaseException as e: + if connection is not None: + with contextlib.suppress(BaseException): + connection.close() + if not isinstance(e, Exception): + raise err_driver_name = ( self.connection_config.get("driver_name") or self.connection_config.get("driver") or "Unknown" ) @@ -400,6 +571,8 @@ def provide_session( Returns: A context manager that yields an AdbcDriver instance. """ + if self._pool_config is not None: + self.create_pool() self._detect_extensions_if_needed() statement_config = resolve_runtime_statement_config( statement_config, self.statement_config, self._default_session_config @@ -412,6 +585,7 @@ def provide_session( statement_config=statement_config, driver_features=handler.driver_features, prepare_driver=self._prepare_driver, + driver_class=AdbcDriver, ) def get_signature_namespace(self) -> "dict[str, Any]": diff --git a/sqlspec/adapters/adbc/core.py b/sqlspec/adapters/adbc/core.py index 920addaf2..9e4bcc545 100644 --- a/sqlspec/adapters/adbc/core.py +++ b/sqlspec/adapters/adbc/core.py @@ -6,7 +6,7 @@ import re import tempfile from collections.abc import Sized -from functools import lru_cache, partial +from functools import partial from pathlib import Path from typing import TYPE_CHECKING, Any, Final, cast from uuid import UUID @@ -15,6 +15,7 @@ from sqlglot import exp from sqlglot.errors import ParseError +from sqlspec.adapters.adbc._typing import AdbcStatusCode from sqlspec.adapters.adbc.type_converter import ADBCOutputConverter, get_adbc_type_converter from sqlspec.core import ( DriverParameterProfile, @@ -32,6 +33,7 @@ ForeignKeyViolationError, ImproperConfigurationError, IntegrityError, + NotFoundError, NotNullViolationError, OperationalError, OperationCancelledError, @@ -47,6 +49,7 @@ from sqlspec.typing import PGVECTOR_INSTALLED, Empty, import_optional_attr from sqlspec.utils.module_loader import import_optional, import_string from sqlspec.utils.serializers import from_json, to_json +from sqlspec.utils.text import normalize_identifier, split_qualified_identifier from sqlspec.utils.type_converters import build_uuid_coercions from sqlspec.utils.type_guards import has_rowcount, has_sqlstate from sqlspec.utils.uuids import uuid4, uuid_from_string @@ -102,8 +105,9 @@ COLUMN_CACHE_MAX_SIZE: int = 256 _IMPLICIT_UPPER_COLUMN_PATTERN: Final[re.Pattern[str]] = re.compile(r"^(?!\d)(?:[A-Z0-9_]+)$") _UUID_TYPES: Final[tuple[type[Any], ...]] = tuple(build_uuid_coercions()) +NoneType: Final = type(None) _STRICT_COERCION_ALLOWED_TYPES: Final[tuple[type[Any], ...]] = ( - type(None), + NoneType, bool, int, float, @@ -834,6 +838,39 @@ def prepare_postgres_parameters( ) +_ADBC_STATUS_CODE_MAP: Final[dict[str, tuple[type[SQLSpecError], str]]] = { + "TIMEOUT": (QueryTimeoutError, "query timeout"), + "CANCELLED": (OperationCancelledError, "operation cancelled"), + "UNAUTHENTICATED": (PermissionDeniedError, "unauthenticated"), + "UNAUTHORIZED": (PermissionDeniedError, "unauthorized"), + "NOT_FOUND": (NotFoundError, "not found"), + "ALREADY_EXISTS": (IntegrityError, "already exists"), + "INVALID_ARGUMENT": (DataError, "invalid argument"), + "INVALID_DATA": (DataError, "invalid data"), + "INTEGRITY": (IntegrityError, "integrity violation"), + "IO": (DatabaseConnectionError, "I/O error"), + "INVALID_STATE": (OperationalError, "invalid state"), + "INTERNAL": (OperationalError, "internal error"), + "NOT_IMPLEMENTED": (SQLSpecError, "not implemented"), +} + + +def _map_adbc_status_exception(error: Any) -> "SQLSpecError | None": + status_code = getattr(error, "status_code", None) + if status_code is None: + return None + status_name: str | None = getattr(status_code, "name", None) + if status_name is None and isinstance(status_code, int) and AdbcStatusCode is not None: + with contextlib.suppress(ValueError): + status_name = AdbcStatusCode(status_code).name + if status_name is None: + status_name = str(status_code).upper() + if status_name in _ADBC_STATUS_CODE_MAP: + exc_cls, desc = _ADBC_STATUS_CODE_MAP[status_name] + return _create_adbc_error(error, exc_cls, desc) + return None + + def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpecError: """Map ADBC errors to SQLSpec exceptions. @@ -842,9 +879,10 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpec avoids issues with exception control flow in different Python versions. Mapping priority: - 1. SQLSTATE codes (most reliable for ADBC drivers) - 2. Error message patterns - 3. Default SQLSpecError fallback + 1. Cancellation and timeout status codes + 2. SQLSTATE codes and vendor-specific diagnostics + 3. Error message patterns + 4. Generic ADBC status codes or SQLSpecError fallback Args: error: The ADBC exception to map @@ -853,13 +891,10 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpec Returns: A SQLSpec exception that wraps the original error """ - del logger - status_code = getattr(error, "status_code", None) - status_name = getattr(status_code, "name", str(status_code)).upper() - if status_name == "TIMEOUT": - return _create_adbc_error(error, QueryTimeoutError, "query timeout") - if status_name == "CANCELLED": - return _create_adbc_error(error, OperationCancelledError, "operation cancelled") + _ = logger + status_mapped = _map_adbc_status_exception(error) + if isinstance(status_mapped, (QueryTimeoutError, OperationCancelledError)): + return status_mapped raw_message = str(error) error_msg = raw_message.lower() @@ -900,6 +935,9 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpec if "connection" in error_msg or "connect" in error_msg: return _create_adbc_error(error, DatabaseConnectionError, "connection error") + if status_mapped is not None: + return status_mapped + return _create_adbc_error(error, SQLSpecError, "database error") @@ -961,10 +999,11 @@ def get_statement_config(detected_dialect: str) -> StatementConfig: "type_coercion_map": _BASE_TYPE_COERCION_MAP, } + statement_overrides: dict[str, Any] = {"dialect": sqlglot_dialect} if detected_dialect == "duckdb": parameter_overrides["preserve_parameter_format"] = False parameter_overrides["supported_execution_parameter_styles"] = {ParameterStyle.QMARK, ParameterStyle.NUMERIC} - parameter_overrides["ast_transformer"] = _transform_duckdb_copy_parameters + statement_overrides["statement_transformers"] = (_transform_duckdb_copy_parameters,) elif detected_dialect == "oracle": parameter_overrides["preserve_parameter_format"] = False parameter_overrides["supported_execution_parameter_styles"] = {ParameterStyle.POSITIONAL_COLON} @@ -976,7 +1015,7 @@ def get_statement_config(detected_dialect: str) -> StatementConfig: parameter_overrides["ast_transformer"] = build_null_pruning_transform(dialect=sqlglot_dialect) return build_statement_config_from_profile( - profile, parameter_overrides=parameter_overrides, statement_overrides={"dialect": sqlglot_dialect} + profile, parameter_overrides=parameter_overrides, statement_overrides=statement_overrides ) @@ -993,19 +1032,16 @@ def apply_driver_features( Updated statement configuration and normalized driver features. """ processed_features: dict[str, Any] = dict(driver_features) if driver_features else {} - processed_features = _normalize_adbc_driver_features(processed_features) resolved_dialect = resolve_dialect_name(statement_config.dialect).lower() json_serializer = cast("Callable[[Any], str] | None", processed_features.setdefault("json_serializer", to_json)) processed_features.setdefault("enable_cast_detection", True) - processed_features.setdefault("strict_type_coercion", False) - processed_features.setdefault("enable_strict_type_coercion", processed_features["strict_type_coercion"]) - processed_features.setdefault("arrow_extension_types", True) - processed_features.setdefault("enable_arrow_extension_types", processed_features["arrow_extension_types"]) + processed_features.setdefault("enable_strict_type_coercion", False) + processed_features.setdefault("enable_arrow_extension_types", True) processed_features.setdefault("enable_pgvector", PGVECTOR_INSTALLED) processed_features.setdefault("enable_paradedb", True) processed_features.setdefault("enable_pg_textsearch", True) - processed_features.setdefault("enable_lowercase_column_names", resolved_dialect in {"oracle", "db2"}) + processed_features.setdefault("enable_lowercase_column_names", resolved_dialect in {"oracle", "db2", "snowflake"}) processed_features.setdefault("enable_statement_profiling", False) if json_serializer is not None: @@ -1363,6 +1399,8 @@ def _convert_array_for_postgres_adbc(value: Any) -> Any: int: _identity, float: _identity, bytes: _identity, + bytearray: bytes, + memoryview: bytes, tuple: _convert_array_for_postgres_adbc, list: _convert_array_for_postgres_adbc, dict: _identity, @@ -1416,7 +1454,29 @@ def _is_uuid_array_data_type(target: Any) -> bool: return len(element_types) == 1 and _is_uuid_data_type(element_types[0]) -@lru_cache(maxsize=256) +class _RewriteCacheState: + __slots__ = ("cache", "hits", "maxsize", "misses") + + def __init__(self, maxsize: int = 256) -> None: + self.maxsize = maxsize + self.hits = 0 + self.misses = 0 + self.cache: dict[ + tuple[str, tuple[int, ...], tuple[int, ...], str], tuple[str, tuple[int, ...], tuple[int, ...]] + ] = {} + + def clear(self) -> None: + self.hits = 0 + self.misses = 0 + self.cache.clear() + + +_REWRITE_CACHE_STATE: Final[_RewriteCacheState] = _RewriteCacheState(256) +_REWRITE_UUID_CACHE: Final[ + dict[tuple[str, tuple[int, ...], tuple[int, ...], str], tuple[str, tuple[int, ...], tuple[int, ...]]] +] = _REWRITE_CACHE_STATE.cache + + def _rewrite_postgres_uuid_placeholders( sql: str, scalar_ordinals: "tuple[int, ...]", array_ordinals: "tuple[int, ...]", dialect: str ) -> "tuple[str, tuple[int, ...], tuple[int, ...]]": @@ -1429,6 +1489,13 @@ def _rewrite_postgres_uuid_placeholders( repeated executions of the same statement shape reuse the rewrite without parsing again. """ + cache_key = (sql, scalar_ordinals, array_ordinals, dialect) + cached = _REWRITE_CACHE_STATE.cache.get(cache_key) + if cached is not None: + _REWRITE_CACHE_STATE.hits += 1 + return cached + + _REWRITE_CACHE_STATE.misses += 1 try: expressions = cast( "list[exp.Expression]", @@ -1463,21 +1530,22 @@ def _rewrite_postgres_uuid_placeholders( effective_array_set = set(effective_arrays) effective_set = set(effective_scalars) | effective_array_set - def wrap_parameter(node: exp.Expression) -> exp.Expression: - if not isinstance(node, exp.Parameter): - return node - ordinal = _parameter_ordinal(node) - if ordinal is None or ordinal not in effective_set: - return node - if _direct_parameter_cast(node) is not None: - return node - target_type = "UUID[]" if ordinal in effective_array_set else "UUID" - return exp.Cast(this=node.copy(), to=exp.DataType.build(target_type, dialect=dialect)) - for expression in expressions: - expression.transform(wrap_parameter, copy=False) + for node in list(expression.find_all(exp.Parameter)): + ordinal = _parameter_ordinal(node) + if ordinal is None or ordinal not in effective_set: + continue + if _direct_parameter_cast(node) is not None: + continue + target_type = "UUID[]" if ordinal in effective_array_set else "UUID" + node.replace(exp.Cast(this=node.copy(), to=exp.DataType.build(target_type, dialect=dialect))) + rewritten = "; ".join(expression.sql(dialect=dialect, copy=False) for expression in expressions) - return rewritten, effective_scalars, effective_arrays + result = (rewritten, effective_scalars, effective_arrays) + if len(_REWRITE_CACHE_STATE.cache) >= _REWRITE_CACHE_STATE.maxsize: + _REWRITE_CACHE_STATE.cache.pop(next(iter(_REWRITE_CACHE_STATE.cache))) + _REWRITE_CACHE_STATE.cache[cache_key] = result + return result def _sequence_holds_uuid(values: "list[Any] | tuple[Any, ...]") -> bool: @@ -1721,20 +1789,6 @@ def _map_sqlstate_exception(error: Any, sqlstate: str) -> SQLSpecError: return _create_adbc_error(error, SQLSpecError, "database error") -def _normalize_adbc_driver_features(processed_features: "dict[str, Any]") -> "dict[str, Any]": - if "strict_type_coercion" in processed_features and "enable_strict_type_coercion" not in processed_features: - processed_features["enable_strict_type_coercion"] = processed_features["strict_type_coercion"] - if "enable_strict_type_coercion" in processed_features and "strict_type_coercion" not in processed_features: - processed_features["strict_type_coercion"] = processed_features["enable_strict_type_coercion"] - - if "arrow_extension_types" in processed_features and "enable_arrow_extension_types" not in processed_features: - processed_features["enable_arrow_extension_types"] = processed_features["arrow_extension_types"] - if "enable_arrow_extension_types" in processed_features and "arrow_extension_types" not in processed_features: - processed_features["arrow_extension_types"] = processed_features["enable_arrow_extension_types"] - - return processed_features - - def _validate_strict_parameter_value(value: Any) -> None: if isinstance(value, (list, tuple)): for item in value: @@ -1899,9 +1953,6 @@ def _prepare_parameter_sequence_with_casts( return tuple(result) if isinstance(parameters, tuple) else result -_prepare_batch_with_casts = prepare_batch_with_casts - - def _resolve_gizmosql_driver_connect() -> str: """Return the preferred GizmoSQL connect path based on installed packages.""" if import_optional("adbc_driver_gizmosql") is not None: @@ -2129,3 +2180,19 @@ def _apply_driver_manager_config(config: "dict[str, Any]", driver_kind: "str | N driver_profile: Final[DriverParameterProfile] = build_profile() + +_DB2_TABLE_EXISTS_SQL: Final = ( + "SELECT 1 FROM SYSCAT.TABLES WHERE TABSCHEMA = COALESCE(CAST(? AS VARCHAR(128)), CURRENT SCHEMA) AND TABNAME = ?" +) +_DB2_INDEX_EXISTS_SQL: Final = ( + "SELECT 1 FROM SYSCAT.INDEXES WHERE INDSCHEMA = COALESCE(CAST(? AS VARCHAR(128)), CURRENT SCHEMA) AND INDNAME = ?" +) +_MIN_QUALIFIED_PARTS: Final = 2 + + +def split_db2_table_name(table_name: str) -> "tuple[str | None, str]": + """Split qualified DB2 table name into schema and table components.""" + parts = [part.strip() for part in split_qualified_identifier(table_name, quote_chars='"') if part.strip()] + if len(parts) >= _MIN_QUALIFIED_PARTS: + return normalize_identifier(parts[-2], "db2"), normalize_identifier(parts[-1], "db2") + return None, normalize_identifier(parts[0] if parts else table_name, "db2") diff --git a/sqlspec/adapters/adbc/data_dictionary.py b/sqlspec/adapters/adbc/data_dictionary.py index 594ae2c84..9b76fa560 100644 --- a/sqlspec/adapters/adbc/data_dictionary.py +++ b/sqlspec/adapters/adbc/data_dictionary.py @@ -57,7 +57,6 @@ __all__ = ("AdbcDataDictionary",) logger = get_logger("sqlspec.adapters.adbc") -_pa: Any = import_optional("pyarrow") _NATIVE_TABLE_TYPES: Final = frozenset({"table", "base table"}) _NATIVE_FALLBACK_ERRORS: Final = (AdbcNotSupportedError, AdbcOperationalError, AttributeError) @@ -346,18 +345,24 @@ def get_indexes( operation = "by_schema" if table is None else "by_table" return driver.select(self._get_query(dialect, "indexes", operation), schema_type=IndexMetadata) - if table is None: + if dialect == "snowflake": + return [] + + try: + if table is None: + return driver.select( + self._get_query(dialect, "indexes", "by_schema"), schema_name=schema_name, schema_type=IndexMetadata + ) + + table_name = self._resolve_identifier(dialect, table) return driver.select( - self._get_query(dialect, "indexes", "by_schema"), schema_name=schema_name, schema_type=IndexMetadata + self._get_query(dialect, "indexes", "by_table"), + schema_name=schema_name, + table_name=table_name, + schema_type=IndexMetadata, ) - - table_name = self._resolve_identifier(dialect, table) - return driver.select( - self._get_query(dialect, "indexes", "by_table"), - schema_name=schema_name, - table_name=table_name, - schema_type=IndexMetadata, - ) + except SQLFileNotFoundError: + return [] def get_foreign_keys( self, driver: "AdbcDriver", table: "str | None" = None, schema: "str | None" = None @@ -821,10 +826,12 @@ def _native_get_foreign_keys( self, driver: "AdbcDriver", dialect: str, table_name: "str | None", schema_name: "str | None" ) -> "list[ForeignKeyMetadata]": rows = self._native_get_objects(driver, dialect, "all", schema_name, table_name) - foreign_keys = _normalize_native_foreign_keys(rows, table_name_exact=table_name) - if not foreign_keys: + if not rows: raise _NativeMetadataIncompleteError - return foreign_keys + for _, _, table in _iter_object_tables(rows): + if table.get("table_constraints") is None: + raise _NativeMetadataIncompleteError + return _normalize_native_foreign_keys(rows, table_name_exact=table_name) class _NativeMetadataIncompleteError(Exception): @@ -994,7 +1001,9 @@ def _normalize_native_constraints( def _arrow_type_to_sql(data_type: Any) -> str: ensure_pyarrow() - types = _pa.types + pa = import_optional("pyarrow") + assert pa is not None + types = pa.types if types.is_boolean(data_type): return "BOOLEAN" if types.is_int8(data_type) or types.is_int16(data_type) or types.is_uint8(data_type) or types.is_uint16(data_type): diff --git a/sqlspec/adapters/adbc/driver.py b/sqlspec/adapters/adbc/driver.py index c3f6cb028..8ef83f0ea 100644 --- a/sqlspec/adapters/adbc/driver.py +++ b/sqlspec/adapters/adbc/driver.py @@ -5,12 +5,14 @@ """ import contextlib +import sys from pathlib import Path from time import perf_counter from typing import TYPE_CHECKING, Any, Literal, cast, final -from sqlspec.adapters.adbc._typing import AdbcCursor, AdbcNativeError, AdbcSessionContext +from sqlspec.adapters.adbc._typing import AdbcCursor, AdbcNativeError, AdbcNotSupportedError, AdbcSessionContext from sqlspec.adapters.adbc.core import ( + _transform_duckdb_copy_parameters, apply_driver_features, cleanup_duckdb_profile_path, collect_rows, @@ -67,7 +69,6 @@ __all__ = ("AdbcCursor", "AdbcDriver", "AdbcExceptionHandler", "AdbcSessionContext") logger = get_logger("sqlspec.adapters.adbc") -_pa: Any = import_optional("pyarrow") _DIALECTS_WITHOUT_SAVEPOINTS = frozenset({"bigquery", "duckdb", "snowflake"}) _MULTI_ROW_BIND_UNSUPPORTED = "Binding multiple rows at once is not supported" @@ -97,14 +98,16 @@ def _handle_exception(self, exc_type: "type[BaseException] | None", exc_val: "Ba class AdbcSelectStreamSource: """Native ADBC chunk source backed by a ``RecordBatchReader``.""" - __slots__ = ("_cursor_manager", "_driver", "_parameters", "_reader", "_sql") + __slots__ = ("_chunk_size", "_driver", "_owner", "_parameters", "_reader", "_remainder", "_sql") - def __init__(self, driver: "AdbcDriver", sql: str, parameters: Any) -> None: + def __init__(self, driver: "AdbcDriver", sql: str, parameters: Any, chunk_size: int = 1000) -> None: self._driver = driver self._sql = sql self._parameters = parameters - self._cursor_manager: AdbcCursor | None = None + self._chunk_size = chunk_size + self._owner: _AdbcArrowStreamOwner | None = None self._reader: Any = None + self._remainder: list[dict[str, Any]] = [] def start(self) -> None: cursor_manager = self._driver.with_cursor(self._driver.connection) @@ -113,30 +116,62 @@ def start(self) -> None: try: cursor = cursor_manager.__enter__() handler = self._driver.handle_database_exceptions() - with handler: - execute_parameters = normalize_postgres_empty_parameters(self._driver._dialect_name, self._parameters) - cursor.execute(self._sql, parameters=execute_parameters) - reader = _fetch_record_batch(cursor) + if sys.exc_info()[1] is None: + with handler: + execute_parameters = normalize_postgres_empty_parameters( + self._driver._dialect_name, self._parameters + ) + cursor.execute(self._sql, parameters=execute_parameters) + reader = _fetch_record_batch(cursor) + else: + + def _stream_op() -> Any: + execute_parameters = normalize_postgres_empty_parameters( + self._driver._dialect_name, self._parameters + ) + cursor.execute(self._sql, parameters=execute_parameters) + return _fetch_record_batch(cursor) + + reader = self._driver._run_with_exception_handler(handler, _stream_op) self._driver._check_pending_exception(handler) except BaseException: if cursor is not None: - handle_postgres_rollback(self._driver._dialect_name, cursor, logger) - with contextlib.suppress(Exception): + with contextlib.suppress(BaseException): + handle_postgres_rollback(self._driver._dialect_name, cursor, logger) + with contextlib.suppress(BaseException): cursor_manager.__exit__(None, None, None) raise - self._cursor_manager = cursor_manager - self._reader = iter(reader) + self._owner = _AdbcArrowStreamOwner(reader, cursor) + try: + self._reader = iter(reader) + except BaseException: + with contextlib.suppress(BaseException): + self.close() + raise def fetch_chunk(self) -> "list[dict[str, Any]]": + chunk_size = self._chunk_size + results: list[dict[str, Any]] = [] + if self._remainder: + if len(self._remainder) >= chunk_size: + chunk = self._remainder[:chunk_size] + self._remainder = self._remainder[chunk_size:] + return chunk + results.extend(self._remainder) + self._remainder = [] + reader = self._reader if reader is None: - return [] - while True: + return results + + while len(results) < chunk_size: try: batch = next(reader) except StopIteration: - return [] + self._reader = None + break + if bool(self._driver.driver_features.get("enable_lowercase_column_names", False)) and hasattr( batch, "schema" ): @@ -144,22 +179,32 @@ def fetch_chunk(self) -> "list[dict[str, Any]]": normalized_names = normalize_column_names(batch_names, self._driver.driver_features) if normalized_names != batch_names and hasattr(batch, "rename_columns"): batch = batch.rename_columns(normalized_names) + rows = arrow_table_to_pylist( batch, decode_arrow_extension_types=bool( self._driver.driver_features.get("enable_arrow_extension_types", True) ), ) - if rows: - return rows + needed = chunk_size - len(results) + if len(rows) <= needed: + results.extend(rows) + else: + results.extend(rows[:needed]) + self._remainder = rows[needed:] + break + + return results def close(self, error: bool = False) -> None: self._reader = None - cursor_manager = self._cursor_manager - self._cursor_manager = None - if cursor_manager is not None: - with contextlib.suppress(Exception): - cursor_manager.__exit__(None, None, None) + self._remainder = [] + if self._owner is not None: + if error: + with contextlib.suppress(BaseException): + self._owner.close() + else: + self._owner.close() @final @@ -175,9 +220,12 @@ class AdbcDriver(SyncDriverAdapterBase): "_data_dictionary", "_detected_dialect", "_dialect_name", + "_has_streams", "_is_flightsql", "_is_postgres", "_json_serializer", + "_previous_autocommit", + "_stream_closers", "_transaction_active", "dialect", ) @@ -209,7 +257,7 @@ def __init__( explicit_lowercase = bool(driver_features and "enable_lowercase_column_names" in driver_features) statement_config, normalized_features = apply_driver_features(statement_config, driver_features) - if not explicit_lowercase and (dialect or self._detected_dialect) in {"oracle", "db2"}: + if not explicit_lowercase and (dialect or self._detected_dialect) in {"oracle", "db2", "snowflake"}: normalized_features["enable_lowercase_column_names"] = True super().__init__(connection=connection, statement_config=statement_config, driver_features=normalized_features) self.dialect = statement_config.dialect @@ -219,6 +267,27 @@ def __init__( self._data_dictionary: AdbcDataDictionary | None = None self._column_name_cache: dict[int, tuple[Any, list[str]]] = {} self._transaction_active = False + self._previous_autocommit: bool | None = None + self._has_streams = False + self._stream_closers: list[Callable[[], None]] = [] + + @property + def has_streams(self) -> bool: + """Return whether this session created any lazy native stream.""" + return self._has_streams + + def close_streams(self) -> None: + """Close every registered native owner before returning a session connection.""" + closers, self._stream_closers = self._stream_closers, [] + first_error: BaseException | None = None + for close in closers: + try: + close() + except BaseException as exc: + if first_error is None: + first_error = exc + if first_error is not None: + raise first_error @property def data_dictionary(self) -> "AdbcDataDictionary": @@ -423,37 +492,48 @@ def dispatch_execute_script(self, cursor: "AdbcRawCursor", statement: "SQL") -> def dispatch_select_stream(self, statement: "SQL", chunk_size: int) -> "SyncRowStream[dict[str, Any]] | None": """Return a native ADBC stream backed by ``fetch_record_batch()``.""" - _ = chunk_size + self._has_streams = True if not statement.returns_rows(): return None sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) - return SyncRowStream(AdbcSelectStreamSource(self, sql, prepared_parameters)) + source = AdbcSelectStreamSource(self, sql, prepared_parameters, chunk_size=chunk_size) + stream: SyncRowStream[dict[str, Any]] = SyncRowStream(source) + self._stream_closers.extend((source.close, stream.close)) + return stream def begin(self) -> None: - """Begin database transaction. - - ADBC connections operate with autocommit disabled and manage the - active transaction internally, so no explicit statement is issued. - """ + """Begin a transaction, disabling native autocommit when its state is available.""" + if self._transaction_active: + return + native = getattr(self.connection, "adbc_connection", None) + get_option = getattr(native, "get_option", None) + mode = None + if callable(get_option): + try: + mode = get_option("adbc.connection.autocommit") + except (AdbcNotSupportedError, NotImplementedError): + pass + except AdbcNativeError as exc: + raise create_mapped_exception(exc) from exc + if mode == "true" and native is not None: + try: + native.set_autocommit(False) + except BaseException as exc: + with contextlib.suppress(BaseException): + native.set_autocommit(True) + if isinstance(exc, AdbcNativeError): + raise create_mapped_exception(exc) from exc + raise + self._previous_autocommit = True self._transaction_active = True def commit(self) -> None: """Commit database transaction.""" - try: - self.connection.commit() - except AdbcNativeError as e: - raise create_mapped_exception(e) from e - finally: - self._transaction_active = False + self._finish_transaction(commit=True) def rollback(self) -> None: """Rollback database transaction.""" - try: - self.connection.rollback() - except AdbcNativeError as e: - raise create_mapped_exception(e) from e - finally: - self._transaction_active = False + self._finish_transaction(commit=False) def create_savepoint(self, name: str) -> None: """Create a savepoint within the current transaction. @@ -589,9 +669,11 @@ def select_to_arrow( arrow_result: ArrowResult | None = None if return_format in {"reader", "batches"}: + if return_format == "reader": + self._has_streams = True cursor_manager = self.with_cursor(self.connection) cursor = None - reader: object | None = None + reader: Any = None with exc_handler: cursor = cursor_manager.__enter__() try: @@ -612,9 +694,11 @@ def select_to_arrow( ) else: reader = fetch_record_batch() - except Exception: - handle_postgres_rollback(self._dialect_name, cursor, logger) - cursor_manager.__exit__(None, None, None) + except BaseException: + with contextlib.suppress(BaseException): + handle_postgres_rollback(self._dialect_name, cursor, logger) + with contextlib.suppress(BaseException): + cursor_manager.__exit__(None, None, None) cursor = None raise @@ -622,7 +706,7 @@ def select_to_arrow( if cursor is not None: with contextlib.suppress(Exception): cursor_manager.__exit__(None, None, None) - raise exc_handler.pending_exception from None + self._check_pending_exception(exc_handler) if arrow_result is not None: return arrow_result if cursor is None or reader is None: @@ -631,18 +715,39 @@ def select_to_arrow( cursor_manager.__exit__(None, None, None) msg = "ADBC did not return an Arrow reader." raise SQLSpecError(msg) + owner = _AdbcArrowStreamOwner(reader, cursor) + if return_format == "reader": + self._stream_closers.append(owner.close) try: - return build_arrow_result_from_reader( + if bool(self.driver_features.get("enable_lowercase_column_names", False)) and hasattr(reader, "schema"): + raw_names = list(reader.schema.names) + normalized_names = normalize_column_names(raw_names, self.driver_features) + pa = import_optional("pyarrow") + if normalized_names != raw_names and pa is not None: + new_schema = pa.schema([ + field.with_name(name) for field, name in zip(reader.schema, normalized_names, strict=True) + ]) + source_reader = reader + + def _lower_batches() -> Any: + for batch in source_reader: + yield batch.rename_columns(normalized_names) + + reader = pa.RecordBatchReader.from_batches(new_schema, _lower_batches()) + result = build_arrow_result_from_reader( prepared_statement, - arrow_reader_with_deferred_close(reader, cursor.close), + arrow_reader_with_deferred_close(reader, owner.finish), return_format=return_format, batch_size=batch_size, arrow_schema=arrow_schema, ) + if return_format == "batches": + owner.close() except BaseException: - with contextlib.suppress(Exception): - cursor_manager.__exit__(None, None, None) + with contextlib.suppress(BaseException): + owner.close() raise + return result with self.with_cursor(self.connection) as cursor, exc_handler: if cursor is None: @@ -654,8 +759,9 @@ def select_to_arrow( try: cursor.execute(sql, parameters=execute_parameters) arrow_table = cursor.fetch_arrow_table() - except Exception: - handle_postgres_rollback(self._dialect_name, cursor, logger) + except BaseException: + with contextlib.suppress(BaseException): + handle_postgres_rollback(self._dialect_name, cursor, logger) raise if bool(self.driver_features.get("enable_lowercase_column_names", False)) and hasattr( @@ -674,8 +780,7 @@ def select_to_arrow( arrow_schema=arrow_schema, ) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + self._check_pending_exception(exc_handler) if arrow_result is None: msg = "Unreachable" @@ -750,7 +855,7 @@ def load_from_arrow( ingest_mode = "replace" if overwrite else "create_append" target_table, ingest_kwargs = _parse_adbc_ingest_target(table, mode=ingest_mode) - if isinstance(source_data, _pa.RecordBatchReader): + if hasattr(source_data, "read_next_batch"): row_count = 0 exc_handler = self.handle_database_exceptions() with self.with_cursor(self.connection) as cursor, exc_handler: @@ -762,8 +867,7 @@ def load_from_arrow( if isinstance(ingest_result, int) and ingest_result > 0: row_count = ingest_result - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + self._check_pending_exception(exc_handler) telemetry_payload: StorageTelemetry = { "rows_processed": row_count, "bytes_processed": 0, @@ -782,8 +886,7 @@ def load_from_arrow( handle_postgres_rollback(self._dialect_name, cursor, logger) raise - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + self._check_pending_exception(exc_handler) telemetry_payload = self._ingest_telemetry(arrow_table) telemetry_payload["destination"] = table self._attach_partition_telemetry(telemetry_payload, partitioner) @@ -874,6 +977,13 @@ def prepare_driver_parameters( def _compiled_sql( self, statement: "SQL", statement_config: "StatementConfig", flatten_single_parameters: bool = False ) -> "tuple[str, object]": + if ( + self._dialect_name == "duckdb" + and statement_config.parameter_config.ast_transformer is None + and statement.raw_sql.lstrip().upper().startswith(("COPY", "INSERT")) + ): + param_config = statement_config.parameter_config.replace(ast_transformer=_transform_duckdb_copy_parameters) + statement_config = statement_config.replace(parameter_config=param_config) compiled_sql, prepared_parameters = super()._compiled_sql( statement, statement_config, flatten_single_parameters=flatten_single_parameters ) @@ -881,6 +991,42 @@ def _compiled_sql( compiled_sql, prepared_parameters, is_many=statement.is_many, dialect=self._dialect_name ) + def _finish_transaction(self, *, commit: bool) -> None: + restore = self._previous_autocommit + primary: BaseException | None = None + try: + connection = self.connection.adbc_connection if restore is True else self.connection + if commit: + connection.commit() + else: + connection.rollback() + except BaseException as exc: + primary = exc + completed = primary is None + if restore is True and not completed: + with contextlib.suppress(BaseException): + self.connection.adbc_connection.rollback() + completed = True + if restore is True: + self._previous_autocommit = True + self._transaction_active = True + if completed: + try: + self.connection.adbc_connection.set_autocommit(True) + except BaseException as exc: + if primary is None: + primary = exc + else: + self._previous_autocommit = None + self._transaction_active = False + else: + self._previous_autocommit = None + self._transaction_active = False + if isinstance(primary, AdbcNativeError): + raise create_mapped_exception(primary) from primary + if primary is not None: + raise primary + def _connection_in_transaction(self) -> bool: """Check if connection is in transaction. @@ -1005,6 +1151,40 @@ def _parse_adbc_ingest_target( return table, kwargs +class _AdbcArrowStreamOwner: + """Close actual native owners independently of lazy wrapper iteration.""" + + __slots__ = ("_close_error", "_closed", "_cursor", "_reader") + + def __init__(self, reader: Any, cursor: "AdbcRawCursor") -> None: + self._reader = reader + self._cursor = cursor + self._closed = False + self._close_error: BaseException | None = None + + def close(self) -> None: + if not self._closed: + self._closed = True + try: + close_reader = getattr(self._reader, "close", None) + if callable(close_reader): + close_reader() + except BaseException as exc: + self._close_error = exc + try: + self._cursor.close() + except BaseException as exc: + if self._close_error is None: + self._close_error = exc + if self._close_error is not None: + raise self._close_error + + def finish(self) -> None: + """Finalize lazy iteration without replacing its exception or raising from GC.""" + with contextlib.suppress(BaseException): + self.close() + + def _fetch_record_batch(cursor: Any) -> Any: fetch_record_batch = getattr(cursor, "fetch_record_batch", None) if fetch_record_batch is None: diff --git a/sqlspec/adapters/adbc/litestar/__init__.py b/sqlspec/adapters/adbc/litestar/__init__.py index 25bc6130d..d6f1050f7 100644 --- a/sqlspec/adapters/adbc/litestar/__init__.py +++ b/sqlspec/adapters/adbc/litestar/__init__.py @@ -1,5 +1,5 @@ """Litestar integration for ADBC adapter.""" -from sqlspec.adapters.adbc.litestar.store import ADBCStore +from sqlspec.adapters.adbc.litestar.store import AdbcStore -__all__ = ("ADBCStore",) +__all__ = ("AdbcStore",) diff --git a/sqlspec/adapters/adbc/litestar/store.py b/sqlspec/adapters/adbc/litestar/store.py index c2951e289..f1659617d 100644 --- a/sqlspec/adapters/adbc/litestar/store.py +++ b/sqlspec/adapters/adbc/litestar/store.py @@ -13,22 +13,25 @@ - Snowflake: BINARY data type, TIMESTAMP WITH TIME ZONE, ? parameters, MERGE """ +import contextlib from datetime import UTC, datetime, timedelta from typing import TYPE_CHECKING -from sqlspec.adapters.adbc.core import is_postgres_dialect +from sqlspec.adapters.adbc.core import _DB2_INDEX_EXISTS_SQL, _DB2_TABLE_EXISTS_SQL, is_postgres_dialect +from sqlspec.adapters.adbc.core import split_db2_table_name as _split_db2_table_name +from sqlspec.exceptions import IntegrityError from sqlspec.extensions.litestar.store import BaseSQLSpecStore from sqlspec.utils.sync_tools import async_ -from sqlspec.utils.text import normalize_identifier, split_qualified_identifier +from sqlspec.utils.text import normalize_identifier if TYPE_CHECKING: from sqlspec.adapters.adbc.config import AdbcConfig -__all__ = ("ADBCStore",) +__all__ = ("AdbcStore",) -class ADBCStore(BaseSQLSpecStore["AdbcConfig"]): +class AdbcStore(BaseSQLSpecStore["AdbcConfig"]): """ADBC session store using synchronous ADBC driver. Implements server-side session storage for Litestar using ADBC @@ -333,9 +336,14 @@ def _create_table(self) -> None: END; """.strip() with self._config.provide_session() as driver: - driver.execute(table_block) - driver.execute(index_block) - driver.commit() + try: + driver.execute(table_block) + driver.execute(index_block) + driver.commit() + except Exception: + with contextlib.suppress(Exception): + driver.rollback() + raise self._log_table_created() return @@ -343,18 +351,30 @@ def _create_table(self) -> None: schema_name, catalog_table = _split_db2_table_name(self._table_name) index_name = normalize_identifier(f"idx_{self._table_name}_expires_at", "db2") with self._config.provide_session() as driver: - if driver.select_one_or_none(_DB2_TABLE_EXISTS_SQL, (schema_name, catalog_table)) is None: - driver.execute(self._table_ddl().strip()) - if driver.select_one_or_none(_DB2_INDEX_EXISTS_SQL, (schema_name, index_name)) is None: - driver.execute(f"CREATE INDEX idx_{self._table_name}_expires_at ON {self._table_name}(expires_at)") - driver.commit() + try: + if driver.select_one_or_none(_DB2_TABLE_EXISTS_SQL, (schema_name, catalog_table)) is None: + driver.execute(self._table_ddl().strip()) + if driver.select_one_or_none(_DB2_INDEX_EXISTS_SQL, (schema_name, index_name)) is None: + driver.execute( + f"CREATE INDEX idx_{self._table_name}_expires_at ON {self._table_name}(expires_at)" + ) + driver.commit() + except Exception: + with contextlib.suppress(Exception): + driver.rollback() + raise self._log_table_created() return sql_text = self._table_ddl() with self._config.provide_session() as driver: - driver.execute_script(sql_text) - driver.commit() + try: + driver.execute_script(sql_text) + driver.commit() + except Exception: + with contextlib.suppress(Exception): + driver.rollback() + raise self._log_table_created() def _drop_table_sql(self) -> "list[str]": @@ -450,24 +470,6 @@ def _set(self, key: str, value: "str | bytes", expires_in: "int | timedelta | No INSERT OR REPLACE INTO {self._table_name} (session_id, data, expires_at) VALUES ({p1}, {p2}, {p3}) """ - elif is_postgres_dialect(dialect): - if expires_at is None: - sql = f""" - INSERT INTO {self._table_name} (session_id, data, expires_at) - VALUES ({p1}, {p2}, NULL) - ON CONFLICT (session_id) DO UPDATE - SET data = EXCLUDED.data, expires_at = NULL - """ - with self._config.provide_session() as driver: - driver.execute(sql, key, data) - driver.commit() - return - sql = f""" - INSERT INTO {self._table_name} (session_id, data, expires_at) - VALUES ({p1}, {p2}, {p3}) - ON CONFLICT (session_id) DO UPDATE - SET data = EXCLUDED.data, expires_at = EXCLUDED.expires_at - """ else: if expires_at is None: sql = f""" @@ -477,8 +479,13 @@ def _set(self, key: str, value: "str | bytes", expires_in: "int | timedelta | No SET data = EXCLUDED.data, expires_at = NULL """ with self._config.provide_session() as driver: - driver.execute(sql, key, data) - driver.commit() + try: + driver.execute(sql, key, data) + driver.commit() + except Exception: + with contextlib.suppress(Exception): + driver.rollback() + raise return sql = f""" INSERT INTO {self._table_name} (session_id, data, expires_at) @@ -494,38 +501,59 @@ def _set(self, key: str, value: "str | bytes", expires_in: "int | timedelta | No """ elif dialect in {"bigquery", "snowflake", "oracle", "db2", "mssql"}: with self._config.provide_session() as driver: - check_sql = f"SELECT COUNT(*) as count FROM {self._table_name} WHERE session_id = {p1}" - result = driver.select_one(check_sql, key) - exists = result and result.get("count", 0) > 0 - - if exists: - if expires_at is None: + try: + check_sql = f"SELECT COUNT(*) as count FROM {self._table_name} WHERE session_id = {p1}" + result = driver.select_one(check_sql, key) + exists = result and result.get("count", 0) > 0 + + if exists: + if expires_at is None: + sql = f""" + UPDATE {self._table_name} + SET data = {p1}, expires_at = NULL + WHERE session_id = {p2} + """ + driver.execute(sql, data, key) + else: + sql = f""" + UPDATE {self._table_name} + SET data = {p1}, expires_at = {p2} + WHERE session_id = {p3} + """ + driver.execute(sql, data, expires_at, key) + elif expires_at is None: sql = f""" - UPDATE {self._table_name} - SET data = {p1}, expires_at = NULL - WHERE session_id = {p2} + INSERT INTO {self._table_name} (session_id, data, expires_at) + VALUES ({p1}, {p2}, NULL) """ - driver.execute(sql, data, key) + try: + driver.execute(sql, key, data) + except IntegrityError: + sql = f""" + UPDATE {self._table_name} + SET data = {p1}, expires_at = NULL + WHERE session_id = {p2} + """ + driver.execute(sql, data, key) else: sql = f""" - UPDATE {self._table_name} - SET data = {p1}, expires_at = {p2} - WHERE session_id = {p3} + INSERT INTO {self._table_name} (session_id, data, expires_at) + VALUES ({p1}, {p2}, {p3}) """ - driver.execute(sql, data, expires_at, key) - elif expires_at is None: - sql = f""" - INSERT INTO {self._table_name} (session_id, data, expires_at) - VALUES ({p1}, {p2}, NULL) - """ - driver.execute(sql, key, data) - else: - sql = f""" - INSERT INTO {self._table_name} (session_id, data, expires_at) - VALUES ({p1}, {p2}, {p3}) - """ - driver.execute(sql, key, data, expires_at) - driver.commit() + try: + driver.execute(sql, key, data, expires_at) + except IntegrityError: + sql = f""" + UPDATE {self._table_name} + SET data = {p1}, expires_at = {p2} + WHERE session_id = {p3} + """ + driver.execute(sql, data, expires_at, key) + driver.commit() + except Exception: + with contextlib.suppress(Exception): + driver.rollback() + raise return else: sql = f""" @@ -536,8 +564,13 @@ def _set(self, key: str, value: "str | bytes", expires_in: "int | timedelta | No """ with self._config.provide_session() as driver: - driver.execute(sql, key, data, expires_at) - driver.commit() + try: + driver.execute(sql, key, data, expires_at) + driver.commit() + except Exception: + with contextlib.suppress(Exception): + driver.rollback() + raise def _delete(self, key: str) -> None: """Synchronous implementation of delete using ADBC driver.""" @@ -545,8 +578,13 @@ def _delete(self, key: str) -> None: sql = f"DELETE FROM {self._table_name} WHERE session_id = {p1}" with self._config.provide_session() as driver: - driver.execute(sql, key) - driver.commit() + try: + driver.execute(sql, key) + driver.commit() + except Exception: + with contextlib.suppress(Exception): + driver.rollback() + raise def _delete_all(self) -> None: """Synchronous implementation of delete_all using ADBC driver.""" @@ -554,8 +592,13 @@ def _delete_all(self) -> None: sql = f"DELETE FROM {self._table_name}" with self._config.provide_session() as driver: - driver.execute(sql) - driver.commit() + try: + driver.execute(sql) + driver.commit() + except Exception: + with contextlib.suppress(Exception): + driver.rollback() + raise def _exists(self, key: str) -> bool: """Synchronous implementation of exists using ADBC driver.""" @@ -614,36 +657,31 @@ def _delete_expired(self) -> int: delete_sql = f"DELETE FROM {self._table_name} WHERE expires_at <= {current_ts}" with self._config.provide_session() as driver: - result = driver.select_one(count_sql) - count = result.get("count", 0) if result else 0 + try: + result = driver.select_one(count_sql) + count = result.get("count", 0) if result else 0 + + if count > 0: + driver.execute(delete_sql) + driver.commit() + self._log_delete_expired(count) + except Exception: + with contextlib.suppress(Exception): + driver.rollback() + raise + else: + return count + with self._config.provide_session() as driver: + try: + exec_result = driver.execute(sql) + driver.commit() + count = exec_result.rows_affected if count > 0: - driver.execute(delete_sql) - driver.commit() self._log_delete_expired(count) - + except Exception: + with contextlib.suppress(Exception): + driver.rollback() + raise + else: return count - - with self._config.provide_session() as driver: - exec_result = driver.execute(sql) - driver.commit() - count = exec_result.rows_affected - if count > 0: - self._log_delete_expired(count) - return count - - -_DB2_TABLE_EXISTS_SQL = ( - "SELECT 1 FROM SYSCAT.TABLES WHERE TABSCHEMA = COALESCE(CAST(? AS VARCHAR(128)), CURRENT SCHEMA) AND TABNAME = ?" -) -_DB2_INDEX_EXISTS_SQL = ( - "SELECT 1 FROM SYSCAT.INDEXES WHERE INDSCHEMA = COALESCE(CAST(? AS VARCHAR(128)), CURRENT SCHEMA) AND INDNAME = ?" -) -_MIN_QUALIFIED_PARTS = 2 - - -def _split_db2_table_name(table_name: str) -> "tuple[str | None, str]": - parts = [part.strip() for part in split_qualified_identifier(table_name, quote_chars='"') if part.strip()] - if len(parts) >= _MIN_QUALIFIED_PARTS: - return normalize_identifier(parts[-2], "db2"), normalize_identifier(parts[-1], "db2") - return None, normalize_identifier(parts[0] if parts else table_name, "db2") diff --git a/sqlspec/adapters/adbc/pool.py b/sqlspec/adapters/adbc/pool.py new file mode 100644 index 000000000..764b3b1e9 --- /dev/null +++ b/sqlspec/adapters/adbc/pool.py @@ -0,0 +1,237 @@ +"""Exclusive PostgreSQL ADBC connection leases.""" + +import contextlib +import math +import os +import threading +import time +from typing import TYPE_CHECKING + +from sqlspec.exceptions import ConnectionTimeoutError, DatabaseConnectionError, ImproperConfigurationError + +if TYPE_CHECKING: + from collections.abc import Callable + + from sqlspec.adapters.adbc._typing import AdbcConnection, AdbcRawCursor + +__all__ = ("AdbcConnectionPool",) + + +class AdbcConnectionPool: + """Own one original seed and bounded exclusive native clone leases. + + Native creation, cloning and cleanup run outside the bookkeeping condition. + Failed cleanup retains ownership until a subsequent close retries it. + """ + + __slots__ = ( + "_clones", + "_condition", + "_create_connection", + "_max_size", + "_pid", + "_reservations", + "_seed", + "_seed_lock", + "_state", + "_timeout", + ) + + def __init__( + self, create_connection: "Callable[[], AdbcConnection]", *, max_size: object = 4, timeout: float = 30.0 + ) -> None: + """Configure a lazy pool without opening any native connections. + + Raw capacity validation runs before compiled integer argument coercion. + """ + if isinstance(max_size, bool) or not isinstance(max_size, int) or max_size <= 0: + msg = "ADBC pool max_size must be a positive integer." + raise ImproperConfigurationError(msg) + if not math.isfinite(timeout) or timeout < 0: + msg = "ADBC pool timeout must be finite and nonnegative." + raise ImproperConfigurationError(msg) + self._create_connection = create_connection + self._max_size = max_size + self._timeout = timeout + self._pid = os.getpid() + self._condition = threading.Condition(threading.Lock()) + self._seed_lock = threading.Lock() + self._seed: AdbcConnection | None = None + self._clones: dict[int, tuple[AdbcConnection, str]] = {} + self._reservations = 0 + self._state = "OPEN" + + def acquire(self) -> "AdbcConnection": + """Acquire an exclusive clone or wait up to the configured timeout.""" + self.check_pid() + deadline = time.monotonic() + self._timeout + with self._condition: + while True: + if self._state != "OPEN": + msg = "ADBC pool is closed to new acquisitions." + raise DatabaseConnectionError(msg) + for key, (connection, state) in self._clones.items(): + if state == "idle": + self._clones[key] = (connection, "active") + return connection + if len(self._clones) + self._reservations < self._max_size: + self._reservations += 1 + break + remaining = deadline - time.monotonic() + if remaining <= 0: + msg = "Timed out waiting for an exclusive ADBC connection." + raise ConnectionTimeoutError(msg) + self._condition.wait(min(remaining, threading.TIMEOUT_MAX)) + try: + with self._seed_lock: + if self._seed is None: + self._seed = self._create_connection() + connection = self._seed.adbc_clone() + except BaseException: + with self._condition: + self._reservations -= 1 + self._condition.notify_all() + with contextlib.suppress(BaseException): + self._finish_close() + raise + with self._condition: + self._reservations -= 1 + accepted = self._state == "OPEN" + self._clones[id(connection)] = (connection, "active" if accepted else "closing") + self._condition.notify_all() + if not accepted: + self._close_clone(connection) + self._finish_close() + msg = "ADBC pool closed during connection acquisition." + raise DatabaseConnectionError(msg) + return connection + + def release(self, connection: "AdbcConnection", *, discard: bool = False) -> None: + """Reset a successful lease for reuse, or evict it after failure/shutdown.""" + self.check_pid() + with self._condition: + entry = self._clones.get(id(connection)) + if entry is None or entry[0] is not connection or entry[1] != "active": + msg = "Connection is not an active lease of this ADBC pool." + raise DatabaseConnectionError(msg) + reset = self._state == "OPEN" and not discard + self._clones[id(connection)] = (connection, "resetting" if reset else "closing") + try: + if reset: + self._reset_connection(connection) + elif discard: + connection.rollback() + except BaseException: + with self._condition: + self._clones[id(connection)] = (connection, "closing") + with contextlib.suppress(BaseException): + self._close_clone(connection) + self._finish_close() + raise + with self._condition: + if reset and self._state == "OPEN": + self._clones[id(connection)] = (connection, "idle") + self._condition.notify_all() + return + self._clones[id(connection)] = (connection, "closing") + self._close_clone(connection) + self._finish_close() + + def close(self) -> bool: + """Begin terminal drain and retry failed cleanup without closing active leases. + + Returns: + True only when every owned clone and the original seed are closed. + """ + self.check_pid() + with self._condition: + if self._state == "CLOSED": + return True + self._state = "DRAINING" + closing = [] + for key, (connection, state) in self._clones.items(): + if state in {"idle", "quarantined"}: + self._clones[key] = (connection, "closing") + closing.append(connection) + self._condition.notify_all() + first_error: BaseException | None = None + for connection in closing: + try: + self._close_clone(connection) + except BaseException as exc: + if first_error is None: + first_error = exc + try: + self._finish_close() + except BaseException as exc: + if first_error is None: + first_error = exc + if first_error is not None: + raise first_error + with self._condition: + return self._state == "CLOSED" + + @property + def is_closed(self) -> bool: + """Return whether every owned native resource has been closed.""" + self.check_pid() + with self._condition: + return self._state == "CLOSED" + + def check_pid(self) -> None: + """Reject inherited pool use before acquiring any lock or native handle.""" + if os.getpid() != self._pid: + msg = "ADBC pools cannot be used after fork; create a new configuration in the child." + raise DatabaseConnectionError(msg) + + def _reset_connection(self, connection: "AdbcConnection") -> None: + cursor: AdbcRawCursor | None = None + first_error: BaseException | None = None + try: + connection.rollback() + connection.adbc_connection.set_autocommit(True) + cursor = connection.cursor() + cursor.execute("DISCARD ALL") + except BaseException as exc: + first_error = exc + finally: + if cursor is not None: + try: + cursor.close() + except BaseException as exc: + if first_error is None: + first_error = exc + try: + connection.adbc_connection.set_autocommit(False) + except BaseException as exc: + if first_error is None: + first_error = exc + if first_error is not None: + raise first_error + + def _close_clone(self, connection: "AdbcConnection") -> None: + try: + connection.close() + except BaseException: + with self._condition: + self._clones[id(connection)] = (connection, "quarantined") + self._condition.notify_all() + raise + with self._condition: + del self._clones[id(connection)] + self._condition.notify_all() + + def _finish_close(self) -> None: + with self._condition: + if self._state != "DRAINING" or self._clones or self._reservations: + return + with self._seed_lock: + with self._condition: + if self._state != "DRAINING" or self._clones or self._reservations: + return + if self._seed is not None: + self._seed.close() + self._seed = None + with self._condition: + self._state = "CLOSED" + self._condition.notify_all() diff --git a/sqlspec/adapters/adbc/type_converter.py b/sqlspec/adapters/adbc/type_converter.py index 8183897fc..63d1b2ee1 100644 --- a/sqlspec/adapters/adbc/type_converter.py +++ b/sqlspec/adapters/adbc/type_converter.py @@ -88,6 +88,10 @@ def convert_sequence( return converted +_CONVERTER_CACHE: Final[dict[tuple[str, int], ADBCOutputConverter]] = {} +_CONVERTER_CACHE_MAX_SIZE: Final[int] = 128 + + def get_adbc_type_converter( dialect: str, cache_size: int = 5000, *, json_serializer: "Callable[[Any], str] | None" = None ) -> ADBCOutputConverter: @@ -102,4 +106,12 @@ def get_adbc_type_converter( Configured ADBCOutputConverter instance. """ _ = cache_size - return ADBCOutputConverter(dialect, json_serializer=json_serializer or to_json) + serializer = json_serializer or to_json + cache_key = (dialect, id(serializer)) + converter = _CONVERTER_CACHE.get(cache_key) + if converter is None: + converter = ADBCOutputConverter(dialect, json_serializer=serializer) + if len(_CONVERTER_CACHE) >= _CONVERTER_CACHE_MAX_SIZE: + del _CONVERTER_CACHE[next(iter(_CONVERTER_CACHE))] + _CONVERTER_CACHE[cache_key] = converter + return converter diff --git a/sqlspec/adapters/aiomysql/_typing.py b/sqlspec/adapters/aiomysql/_typing.py index c101c8f48..911dcf71c 100644 --- a/sqlspec/adapters/aiomysql/_typing.py +++ b/sqlspec/adapters/aiomysql/_typing.py @@ -61,8 +61,7 @@ class AiomysqlFieldTypeProtocol(Protocol): AiomysqlPool: TypeAlias = _AiomysqlPool AiomysqlPymysqlError: TypeAlias = _PymysqlError AiomysqlPymysqlMySQLError: TypeAlias = _PymysqlMySQLError - -if not TYPE_CHECKING: +else: AiomysqlConnection = Connection AiomysqlModule = _aiomysql AiomysqlRawCursor = _AiomysqlCursor diff --git a/sqlspec/adapters/aiomysql/adk/store.py b/sqlspec/adapters/aiomysql/adk/store.py index 3a9016b6e..d0f4458d6 100644 --- a/sqlspec/adapters/aiomysql/adk/store.py +++ b/sqlspec/adapters/aiomysql/adk/store.py @@ -5,7 +5,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, NotRequired, cast from sqlspec.adapters.aiomysql._typing import AiomysqlCursor, AiomysqlProgrammingError, AiomysqlRawCursor -from sqlspec.adapters.aiomysql.core import resolve_rowcount +from sqlspec.adapters.aiomysql.core import resolve_many_rowcount, resolve_rowcount from sqlspec.config import ADKConfig from sqlspec.extensions.adk import BaseAsyncADKStore, StoredEvent, StoredSession, normalize_session_list_options from sqlspec.extensions.adk.memory.store import BaseAsyncADKMemoryStore @@ -566,46 +566,52 @@ async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: " ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) """ - async with self._config.provide_connection() as conn: - async with AiomysqlCursor(conn, cursor_class=AiomysqlRawCursor) as cursor: - for entry in entries: - content_json = self._json_serializer(entry["content_json"]) - meta = entry.get("metadata_json") - metadata_json = self._json_serializer(meta) if meta is not None else None - params: tuple[Any, ...] - if self._owner_id_column_name: - params = ( - entry["id"], - entry["session_id"], - entry["app_name"], - entry["user_id"], - entry.get("scope", "user"), - entry["event_id"], - entry.get("author"), - owner_id, - entry["timestamp"], - content_json, - entry["content_text"], - metadata_json, - entry["inserted_at"], - ) - else: - params = ( - entry["id"], - entry["session_id"], - entry["app_name"], - entry["user_id"], - entry.get("scope", "user"), - entry["event_id"], - entry.get("author"), - entry["timestamp"], - content_json, - entry["content_text"], - metadata_json, - entry["inserted_at"], - ) - await cursor.execute(sql, params) - inserted_count += resolve_rowcount(cursor) + params_list: list[tuple[Any, ...]] = [] + for entry in entries: + content_json = self._json_serializer(entry["content_json"]) + meta = entry.get("metadata_json") + metadata_json = self._json_serializer(meta) if meta is not None else None + if self._owner_id_column_name: + params_list.append(( + entry["id"], + entry["session_id"], + entry["app_name"], + entry["user_id"], + entry.get("scope", "user"), + entry["event_id"], + entry.get("author"), + owner_id, + entry["timestamp"], + content_json, + entry["content_text"], + metadata_json, + entry["inserted_at"], + )) + else: + params_list.append(( + entry["id"], + entry["session_id"], + entry["app_name"], + entry["user_id"], + entry.get("scope", "user"), + entry["event_id"], + entry.get("author"), + entry["timestamp"], + content_json, + entry["content_text"], + metadata_json, + entry["inserted_at"], + )) + + if not params_list: + return 0 + + async with ( + self._config.provide_connection() as conn, + AiomysqlCursor(conn, cursor_class=AiomysqlRawCursor) as cursor, + ): + await cursor.executemany(sql, params_list) + inserted_count = resolve_many_rowcount(cursor, params_list, fallback_count=len(params_list)) await conn.commit() return inserted_count diff --git a/sqlspec/adapters/aiomysql/config.py b/sqlspec/adapters/aiomysql/config.py index 701380930..e5655fbdb 100644 --- a/sqlspec/adapters/aiomysql/config.py +++ b/sqlspec/adapters/aiomysql/config.py @@ -171,6 +171,13 @@ def build_connection_config( class _AiomysqlSessionFactory(AsyncPoolSessionFactory): + """Session factory for aiomysql connections. + + The ``_contexts`` dictionary maps connection IDs to pool acquire contexts + to track and release connections safely across concurrent session checkouts + from the same factory instance. + """ + __slots__ = ("_contexts",) def __init__(self, config: "AiomysqlConfig") -> None: @@ -328,7 +335,11 @@ async def create_connection(self) -> AiomysqlConnection: An aiomysql connection instance. """ connection = await aiomysql.connect(**self._connection_kwargs()) - await self._ensure_connection(connection) + try: + await self._ensure_connection(connection) + except BaseException: + connection.close() + raise return connection async def provide_pool(self, *args: Any, **kwargs: Any) -> "AiomysqlPool": diff --git a/sqlspec/adapters/aiomysql/core.py b/sqlspec/adapters/aiomysql/core.py index 005efd2ee..99ec7f546 100644 --- a/sqlspec/adapters/aiomysql/core.py +++ b/sqlspec/adapters/aiomysql/core.py @@ -17,6 +17,7 @@ ForeignKeyViolationError, IntegrityError, NotNullViolationError, + OperationalError, PermissionDeniedError, QueryTimeoutError, SQLParsingError, @@ -28,7 +29,7 @@ from sqlspec.utils.serializers import from_json, to_json from sqlspec.utils.text import quote_backtick_identifier, split_qualified_identifier from sqlspec.utils.type_converters import build_uuid_coercions -from sqlspec.utils.type_guards import has_cursor_metadata, has_lastrowid, has_rowcount, has_sqlstate, has_type_code +from sqlspec.utils.type_guards import has_lastrowid, has_rowcount, has_sqlstate, has_type_code if TYPE_CHECKING: from sqlspec.core.parameters import ParameterValidator @@ -45,8 +46,6 @@ "collect_stream_rows", "create_mapped_exception", "default_statement_config", - "detect_json_columns", - "detect_json_columns_from_description", "driver_profile", "encode_records_for_local_infile", "escape_literal_percent", @@ -54,7 +53,6 @@ "normalize_execute_many_parameters", "normalize_execute_parameters", "normalize_lastrowid", - "resolve_column_names", "resolve_many_rowcount", "resolve_row_plan", "resolve_rowcount", @@ -80,7 +78,10 @@ MYSQL_SYNTAX_ERROR_MIN: Final[int] = 1064 MYSQL_SYNTAX_ERROR_MAX_EXCLUSIVE: Final[int] = 1100 -_MYSQL_MIGRATION_ERROR_CODES: Final[frozenset[int]] = frozenset((1061, 1091)) +_MYSQL_DDL_ERROR_DISPATCH: dict[int, tuple[type[SQLSpecError], str]] = { + 1061: (OperationalError, "duplicate key name"), + 1091: (OperationalError, "cannot drop column or key"), +} _MYSQL_SQLSTATE_EXACT_DISPATCH: Final[dict[str, tuple[type[SQLSpecError], str]]] = { "23505": (UniqueViolationError, "unique constraint violation"), "23503": (ForeignKeyViolationError, "foreign key constraint violation"), @@ -207,7 +208,7 @@ async def _start(self) -> None: def build_profile() -> "DriverParameterProfile": """Create the aiomysql driver parameter profile.""" - coercions: dict[type, Callable[[Any], Any]] = {bool: bool_to_int, **build_uuid_coercions()} + coercions: dict[type, Callable[[Any], Any]] = {bool: int, **build_uuid_coercions()} return DriverParameterProfile( name="aiomysql", default_style=ParameterStyle.QMARK, @@ -261,14 +262,14 @@ def format_identifier(identifier: str) -> str: msg = "Table name must not be empty" raise SQLSpecError(msg) parts = split_qualified_identifier(cleaned, quote_chars="`", allow_bracket_quotes=False) - formatted = ".".join(quote_backtick_identifier(part) for part in parts) + formatted = ".".join([quote_backtick_identifier(part) for part in parts]) return formatted or quote_backtick_identifier(cleaned) def build_insert_statement(table: str, columns: "list[str]") -> str: """Construct an INSERT statement with positional placeholders.""" - column_clause = ", ".join(quote_backtick_identifier(column) for column in columns) - placeholders = ", ".join("%s" for _ in columns) + column_clause = ", ".join([quote_backtick_identifier(column) for column in columns]) + placeholders = ", ".join(["%s" for _ in columns]) return f"INSERT INTO {format_identifier(table)} ({column_clause}) VALUES ({placeholders})" @@ -301,7 +302,11 @@ def can_use_local_infile_records(arrow_table: Any, records: "Sequence[Sequence[A all_safe_types = False if all_safe_types: return True - return not any(isinstance(value, (bytes, bytearray, memoryview, timedelta)) for row in records for value in row) + for row in records: + for value in row: + if isinstance(value, (bytes, bytearray, memoryview, timedelta)): + return False + return True def encode_records_for_local_infile(records: "list[tuple[Any, ...]]") -> bytes: @@ -345,13 +350,6 @@ def normalize_execute_many_parameters(parameters: Any) -> Any: return parameters -def resolve_column_names(description: "Sequence[Any] | None") -> "list[str]": - """Extract column names from cursor description.""" - if not description: - return [] - return [col[0] for col in description] - - def resolve_row_plan( description: "Sequence[Any] | None", json_type_codes: "set[int]" ) -> "tuple[list[str], list[int] | None]": @@ -361,7 +359,7 @@ def resolve_row_plan( column_names: list[str] = [] if not json_type_codes: - column_names.extend(column[0] for column in description) + column_names.extend([column[0] for column in description]) return column_names, None json_indexes: list[int] = [] @@ -377,24 +375,6 @@ def resolve_row_plan( return column_names, json_indexes or None -def detect_json_columns_from_description( - description: "Sequence[Any] | None", json_type_codes: "set[int]" -) -> "list[int]": - """Identify JSON column indexes from pre-fetched cursor description metadata.""" - return resolve_row_plan(description, json_type_codes)[1] or [] - - -def detect_json_columns( - cursor: Any, json_type_codes: "set[int]", description: "Sequence[Any] | None" = None -) -> "list[int]": - """Identify JSON column indexes from cursor metadata.""" - if description is None: - if not has_cursor_metadata(cursor): - return [] - description = cursor.description - return detect_json_columns_from_description(description, json_type_codes) - - def collect_stream_rows( rows: "Sequence[Any]", row_plan: "tuple[list[str], list[int] | None]", @@ -476,13 +456,9 @@ def normalize_lastrowid(cursor: Any) -> int | None: return last_id if isinstance(last_id, int) else None -def bool_to_int(value: bool) -> int: - """Coerce boolean values to integer for MySQL parameter compatibility.""" - return int(value) - - -def create_mapped_exception(error: Any, *, logger: Any | None = None) -> "SQLSpecError | bool": +def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpecError: """Map MySQL driver exceptions to SQLSpec errors.""" + _ = logger error_code = getattr(error, "errno", None) if error_code is None and isinstance(error, Exception) and error.args: value = error.args[0] @@ -491,10 +467,9 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> "SQLSpe sqlstate = error.sqlstate if has_sqlstate(error) else None sqlstate_prefix = sqlstate[:2] if isinstance(sqlstate, str) and sqlstate else None - if error_code in _MYSQL_MIGRATION_ERROR_CODES: - if logger is not None: - logger.warning("MySQL expected migration error (ignoring): %s", error) - return True + dispatch = _MYSQL_DDL_ERROR_DISPATCH.get(error_code) if error_code is not None else None + if dispatch is not None: + return _create_mysql_error(error, sqlstate, error_code, dispatch[0], dispatch[1]) dispatch = _MYSQL_SQLSTATE_EXACT_DISPATCH.get(sqlstate) if sqlstate is not None else None if dispatch is not None: @@ -550,7 +525,9 @@ def escape_literal_percent(sql: str, parameters: Any, validator: "ParameterValid if not parameters or "%" not in sql: return sql keep = {info.position for info in validator.extract_parameters(sql)} - return "".join("%%" if char == "%" and index not in keep else char for index, char in enumerate(sql)) + if len(keep) == sql.count("%"): + return sql + return "".join(["%%" if char == "%" and index not in keep else char for index, char in enumerate(sql)]) def build_load_data_statement( @@ -560,10 +537,10 @@ def build_load_data_statement( table_sql = format_identifier(table) if escape_percent: table_sql = table_sql.replace("%", "%%") - column_list = ", ".join( + column_list = ", ".join([ quote_backtick_identifier(column).replace("%", "%%") if escape_percent else quote_backtick_identifier(column) for column in columns - ) + ]) return ( f"LOAD DATA LOCAL INFILE {placeholder} INTO TABLE {table_sql} " "CHARACTER SET utf8mb4 FIELDS TERMINATED BY '\\t' ESCAPED BY '\\\\' " diff --git a/sqlspec/adapters/aiomysql/driver.py b/sqlspec/adapters/aiomysql/driver.py index f6299971d..e6734526d 100644 --- a/sqlspec/adapters/aiomysql/driver.py +++ b/sqlspec/adapters/aiomysql/driver.py @@ -19,7 +19,6 @@ AiomysqlCursor, AiomysqlFieldType, AiomysqlPymysqlError, - AiomysqlPymysqlMySQLError, AiomysqlSessionContext, ) from sqlspec.adapters.aiomysql.core import ( @@ -44,7 +43,6 @@ from sqlspec.adapters.aiomysql.data_dictionary import AiomysqlDataDictionary from sqlspec.core import ArrowResult, get_cache_config, register_driver_profile from sqlspec.driver import AsyncDriverAdapterBase, AsyncRowStream, BaseAsyncExceptionHandler -from sqlspec.exceptions import SQLSpecError from sqlspec.utils.logging import get_logger from sqlspec.utils.serializers import from_json from sqlspec.utils.type_guards import supports_json_type @@ -98,10 +96,7 @@ def _handle_exception(self, exc_type: "type[BaseException] | None", exc_val: "Ba if exc_type is None: return False if issubclass(exc_type, AiomysqlPymysqlError): - result = create_mapped_exception(exc_val, logger=logger) - if result is True: - return True - self.pending_exception = cast("Exception", result) + self.pending_exception = create_mapped_exception(exc_val) return True return False @@ -156,8 +151,8 @@ async def dispatch_execute(self, cursor: Any, statement: "SQL") -> "ExecutionRes Returns: ExecutionResult: Statement execution results with data or row counts """ - sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) - sql = escape_literal_percent(sql, prepared_parameters, self.statement_config.parameter_validator) + sql, prepared_parameters = self._compiled_sql(statement, statement.statement_config) + sql = escape_literal_percent(sql, prepared_parameters, statement.statement_config.parameter_validator) await cursor.execute(sql, normalize_execute_parameters(prepared_parameters)) description = getattr(cursor, "description", None) or None @@ -193,15 +188,17 @@ async def dispatch_execute_many(self, cursor: Any, statement: "SQL") -> "Executi Returns: ExecutionResult: Batch execution results """ - sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) + sql, prepared_parameters = self._compiled_sql(statement, statement.statement_config) match = AIOMYSQL_INSERT_VALUES_PATTERN.match(sql) if match: sql = ( - escape_literal_percent(match.group(1), prepared_parameters, self.statement_config.parameter_validator) + escape_literal_percent( + match.group(1), prepared_parameters, statement.statement_config.parameter_validator + ) + sql[match.end(1) :] ) else: - sql = escape_literal_percent(sql, prepared_parameters, self.statement_config.parameter_validator) + sql = escape_literal_percent(sql, prepared_parameters, statement.statement_config.parameter_validator) prepared_parameters = normalize_execute_many_parameters(prepared_parameters) parameter_count = len(prepared_parameters) if isinstance(prepared_parameters, Sized) else None @@ -221,25 +218,7 @@ async def dispatch_execute_script(self, cursor: Any, statement: "SQL") -> "Execu Returns: ExecutionResult: Script execution results with statement count """ - sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) - statements = self.split_script_statements(sql, statement.statement_config, strip_trailing_semicolon=True) - - successful_count = 0 - last_cursor = cursor - normalized_params = normalize_execute_parameters(prepared_parameters) - - for stmt in statements: - stmt_sql = ( - escape_literal_percent(stmt, prepared_parameters, self.statement_config.parameter_validator) - if normalized_params is not None - else stmt - ) - await cursor.execute(stmt_sql, normalized_params) - successful_count += 1 - - return self.create_execution_result( - last_cursor, statement_count=len(statements), successful_statements=successful_count, is_script_result=True - ) + return await super().dispatch_execute_script(cursor, statement) async def begin(self) -> None: """Begin a database transaction. @@ -249,12 +228,8 @@ async def begin(self) -> None: """ try: await self.connection.begin() - except AiomysqlPymysqlMySQLError as e: - mapped = create_mapped_exception(e, logger=logger) - if isinstance(mapped, Exception): - raise mapped from e - msg = f"Failed to begin MySQL transaction: {e}" - raise SQLSpecError(msg) from e + except AiomysqlPymysqlError as e: + raise create_mapped_exception(e) from e async def commit(self) -> None: """Commit the current transaction. @@ -264,12 +239,8 @@ async def commit(self) -> None: """ try: await self.connection.commit() - except AiomysqlPymysqlMySQLError as e: - mapped = create_mapped_exception(e, logger=logger) - if isinstance(mapped, Exception): - raise mapped from e - msg = f"Failed to commit MySQL transaction: {e}" - raise SQLSpecError(msg) from e + except AiomysqlPymysqlError as e: + raise create_mapped_exception(e) from e async def rollback(self) -> None: """Rollback the current transaction. @@ -279,12 +250,8 @@ async def rollback(self) -> None: """ try: await self.connection.rollback() - except AiomysqlPymysqlMySQLError as e: - mapped = create_mapped_exception(e, logger=logger) - if isinstance(mapped, Exception): - raise mapped from e - msg = f"Failed to rollback MySQL transaction: {e}" - raise SQLSpecError(msg) from e + except AiomysqlPymysqlError as e: + raise create_mapped_exception(e) from e def with_cursor(self, connection: "AiomysqlConnection") -> "AiomysqlCursor": """Create cursor context manager for the connection. @@ -301,8 +268,8 @@ def dispatch_select_stream(self, statement: "SQL", chunk_size: int) -> "AsyncRow """Return a native aiomysql row stream backed by an unbuffered ``SSCursor``.""" if not statement.returns_rows(): return None - sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) - sql = escape_literal_percent(sql, prepared_parameters, self.statement_config.parameter_validator) + sql, prepared_parameters = self._compiled_sql(statement, statement.statement_config) + sql = escape_literal_percent(sql, prepared_parameters, statement.statement_config.parameter_validator) return AsyncRowStream( AiomysqlStreamSource(self, sql, prepared_parameters, chunk_size, AIOMYSQL_JSON_TYPE_CODES) ) @@ -392,11 +359,9 @@ async def load_from_arrow( ) if use_infile: payload = encode_records_for_local_infile(records) - fd, tmp_name = tempfile.mkstemp(suffix=".tsv") + tmp_name = await asyncio.to_thread(_write_local_infile_tempfile, payload) try: - with os.fdopen(fd, "wb") as tmp: - tmp.write(payload) - load_sql = build_load_data_statement(table, columns) + load_sql = build_load_data_statement(table, columns, escape_percent=True) exc_handler = self.handle_database_exceptions() await self._run_with_exception_handler( exc_handler, self._execute_cursor_statement, load_sql, (tmp_name,) @@ -422,22 +387,6 @@ async def load_from_arrow( self._attach_partition_telemetry(telemetry_payload, partitioner) return self._storage_job(telemetry_payload, telemetry) - async def load_from_storage( - self, - table: str, - source: "StorageDestination", - *, - file_format: "StorageFormat", - partitioner: "dict[str, object] | None" = None, - overwrite: bool = False, - **kwargs: Any, - ) -> "StorageBridgeJob": - """Load staged artifacts from storage into MySQL.""" - - return await super().load_from_storage( - table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs - ) - def collect_rows(self, cursor: Any, fetched: "list[Any]") -> "tuple[list[Any], list[str], int]": """Collect aiomysql rows for the direct execution path.""" description = cursor.description or None @@ -452,10 +401,18 @@ def resolve_rowcount(self, cursor: Any) -> int: async def _execute_cache_hit( self, sql: str, params: "tuple[Any, ...] | list[Any] | dict[str, Any]", cached: "CachedQuery" ) -> "SQLResult": - statement = self._cached_statement( - sql, params, cached, params, params_are_simple=True, compiled_sql=cached.compiled_sql - ) - return await self._execute_cached_statement(statement) + if ( + "%" in cached.compiled_sql + and escape_literal_percent( + cached.compiled_sql, params or (None,), self.statement_config.parameter_validator + ) + != cached.compiled_sql + ): + statement = self._cached_statement( + sql, params, cached, params, params_are_simple=True, compiled_sql=cached.compiled_sql + ) + return await self._execute_cached_statement(statement) + return await super()._execute_cache_hit(sql, params, cached) async def _execute_cursor_statement(self, statement: str, parameters: Any = None) -> None: async with self.with_cursor(self.connection) as cursor: @@ -477,6 +434,14 @@ def _connection_in_transaction(self) -> bool: return bool(self.connection.get_transaction_status()) +def _write_local_infile_tempfile(payload: bytes) -> str: + """Write LOCAL INFILE payload to a temporary file.""" + fd, tmp_name = tempfile.mkstemp(suffix=".tsv") + with os.fdopen(fd, "wb") as tmp: + tmp.write(payload) + return tmp_name + + def _resolve_column_types(description: Any) -> "dict[str, str] | None": """Map MySQL cursor column FIELD_TYPE codes to neutral Arrow type tokens. diff --git a/sqlspec/adapters/aiomysql/litestar/store.py b/sqlspec/adapters/aiomysql/litestar/store.py index 8f8b449f7..b2b154786 100644 --- a/sqlspec/adapters/aiomysql/litestar/store.py +++ b/sqlspec/adapters/aiomysql/litestar/store.py @@ -140,10 +140,10 @@ async def set(self, key: str, value: "str | bytes", expires_in: "int | timedelta sql = f""" INSERT INTO {self._table_name} (session_id, data, expires_at) - VALUES (%s, %s, %s) AS new + VALUES (%s, %s, %s) ON DUPLICATE KEY UPDATE - data = new.data, - expires_at = new.expires_at, + data = VALUES(data), + expires_at = VALUES(expires_at), updated_at = UTC_TIMESTAMP(6) """ @@ -301,12 +301,9 @@ def _drop_table_sql(self) -> "list[str]": """Get MySQL/MariaDB DROP TABLE SQL statements. Returns: - List of SQL statements to drop indexes and table. + List of SQL statements to drop table. """ - return [ - f"DROP INDEX idx_{self._table_name}_expires_at ON {self._table_name}", - f"DROP TABLE IF EXISTS {self._table_name}", - ] + return [f"DROP TABLE IF EXISTS {self._table_name}"] def _is_mysql_table_missing(exc: BaseException) -> bool: diff --git a/sqlspec/adapters/aiosqlite/adk/store.py b/sqlspec/adapters/aiosqlite/adk/store.py index d37e75016..8997b932e 100644 --- a/sqlspec/adapters/aiosqlite/adk/store.py +++ b/sqlspec/adapters/aiosqlite/adk/store.py @@ -1,8 +1,11 @@ """Aiosqlite async ADK store for Google Agent Development Kit session/event storage.""" +import array import contextlib import inspect +import math import re +import weakref from collections.abc import Mapping from datetime import UTC, datetime, timedelta from typing import TYPE_CHECKING, Any, Final, Literal, NotRequired, cast @@ -74,7 +77,7 @@ class AiosqliteADKStore(BaseAsyncADKStore["AiosqliteConfig"]): config: AiosqliteConfig with extension_config["adk"] settings. """ - __slots__ = ("_configured_connection_ids", "_pragma_overrides") + __slots__ = ("_configured_connections", "_pragma_overrides") def __init__(self, config: "AiosqliteConfig") -> None: """Initialize Aiosqlite ADK store. @@ -84,7 +87,7 @@ def __init__(self, config: "AiosqliteConfig") -> None: """ super().__init__(config) self._pragma_overrides = _pragma_overrides(config) - self._configured_connection_ids: set[int] = set() + self._configured_connections: weakref.WeakSet[Any] = weakref.WeakSet() async def create_tables(self) -> None: """Create both sessions and events tables if they don't exist.""" @@ -474,7 +477,11 @@ async def get_events( params.append(_datetime_to_julian(after_timestamp)) where_clause = " AND ".join(where_clauses) - limit_clause = f" LIMIT {limit}" if limit else "" + if limit is not None: + limit_clause = " LIMIT ?" + params.append(int(limit)) + else: + limit_clause = "" sql = f""" SELECT id, app_name, user_id, session_id, invocation_id, timestamp, event_data @@ -688,7 +695,7 @@ async def set_metadata(self, key: str, value: str) -> None: async def _ensure_pragmas(self, connection: Any) -> None: """Apply PRAGMA optimization profile once per connection.""" - if id(connection) not in self._configured_connection_ids: + if connection not in self._configured_connections: await self._apply_pragmas(connection) async def _apply_pragmas(self, connection: Any) -> None: @@ -703,7 +710,7 @@ async def _apply_pragmas(self, connection: Any) -> None: await _close_cursor(await connection.execute("PRAGMA journal_size_limit = 67108864")) for pragma_name, pragma_value in self._pragma_overrides: await _close_cursor(await connection.execute(f"PRAGMA {pragma_name} = {pragma_value}")) - self._configured_connection_ids.add(id(connection)) + self._configured_connections.add(connection) async def _sessions_table_ddl(self) -> str: """Get SQLite CREATE TABLE SQL for sessions. @@ -823,7 +830,7 @@ class AiosqliteADKMemoryStore(BaseAsyncADKMemoryStore["AiosqliteConfig"]): config: AiosqliteConfig with extension_config["adk"] settings. """ - __slots__ = ("_configured_connection_ids", "_fts_options", "_pragma_overrides") + __slots__ = ("_configured_connections", "_fts_options", "_pragma_overrides") def __init__(self, config: "AiosqliteConfig") -> None: """Initialize Aiosqlite ADK memory store. @@ -834,7 +841,7 @@ def __init__(self, config: "AiosqliteConfig") -> None: super().__init__(config) self._fts_options = _fts_options(config) self._pragma_overrides = _pragma_overrides(config) - self._configured_connection_ids: set[int] = set() + self._configured_connections: weakref.WeakSet[Any] = weakref.WeakSet() async def create_tables(self) -> None: """Create the memory table and indexes if they don't exist. @@ -873,8 +880,8 @@ async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: " INSERT OR IGNORE INTO {self._memory_table} (id, session_id, app_name, user_id, scope, event_id, author, {self._owner_id_column_name}, timestamp, content_json, - content_text, metadata_json, inserted_at) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + content_text, metadata_json, inserted_at, embedding) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """ for entry in entries: scope = entry.get("scope", "user") @@ -893,13 +900,14 @@ async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: " entry["content_text"], to_json(metadata_json) if metadata_json is not None else None, _datetime_to_julian(entry["inserted_at"]), + _pack_embedding(entry.get("embedding")), )) else: sql = f""" INSERT OR IGNORE INTO {self._memory_table} (id, session_id, app_name, user_id, scope, event_id, author, - timestamp, content_json, content_text, metadata_json, inserted_at) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + timestamp, content_json, content_text, metadata_json, inserted_at, embedding) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """ for entry in entries: scope = entry.get("scope", "user") @@ -917,6 +925,7 @@ async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: " entry["content_text"], to_json(metadata_json) if metadata_json is not None else None, _datetime_to_julian(entry["inserted_at"]), + _pack_embedding(entry.get("embedding")), )) cursor = await conn.executemany(sql, params_list) try: @@ -935,12 +944,13 @@ async def search_entries( scope_filter: Literal["all", "user", "app"] = "all", embedding: "Sequence[float] | None" = None, ) -> "list[StoredMemory]": - """Search memory entries by text query.""" + """Search memory entries by text query or vector embedding.""" if not self._enabled: msg = "Memory store is disabled" raise RuntimeError(msg) - if not query or not query.strip(): + has_query = bool(query and query.strip()) + if not has_query and embedding is None: return [] limit_value = limit if limit is not None else self._max_results @@ -949,7 +959,7 @@ async def search_entries( async with self._config.provide_connection() as conn: await self._ensure_foreign_keys(conn) - if self._use_fts: + if has_query and self._use_fts: where_scope, scope_params = _build_sqlite_scope_clause("m.", app_name, user_id, scope_filter) fts_sql = f""" SELECT m.* FROM {self._memory_table} AS m @@ -958,7 +968,8 @@ async def search_entries( ORDER BY m.timestamp DESC LIMIT ? """ - fts_params = (*scope_params, query, limit_value) + fts_limit = limit_value if embedding is None else max(limit_value * 5, 50) + fts_params = (*scope_params, query, fts_limit) try: cursor = await conn.execute(fts_sql, fts_params) try: @@ -976,13 +987,23 @@ async def search_entries( if rows is None: where_scope, scope_params = _build_sqlite_scope_clause("", app_name, user_id, scope_filter) - sql = f""" - SELECT * FROM {self._memory_table} - WHERE {where_scope} AND content_text LIKE ? - ORDER BY timestamp DESC - LIMIT ? - """ - params = (*scope_params, f"%{query}%", limit_value) + fetch_limit = limit_value if embedding is None else max(limit_value * 5, 50) + if has_query: + sql = f""" + SELECT * FROM {self._memory_table} + WHERE {where_scope} AND content_text LIKE ? + ORDER BY timestamp DESC + LIMIT ? + """ + params = (*scope_params, f"%{query}%", fetch_limit) + else: + sql = f""" + SELECT * FROM {self._memory_table} + WHERE {where_scope} AND embedding IS NOT NULL + ORDER BY timestamp DESC + LIMIT ? + """ + params = (*scope_params, fetch_limit) try: cursor = await conn.execute(sql, params) try: @@ -1002,8 +1023,19 @@ async def search_entries( raw["inserted_at"] = _julian_to_datetime(raw["inserted_at"]) raw["content_json"] = from_json(raw["content_json"]) raw["metadata_json"] = from_json(raw["metadata_json"]) if raw["metadata_json"] else None - raw["embedding"] = None + raw["embedding"] = _unpack_embedding(raw.get("embedding")) records.append(cast("StoredMemory", raw)) + + if embedding is not None and records: + query_vec = [float(x) for x in embedding] + + def _score(r: "StoredMemory") -> float: + emb = r.get("embedding") + return _cosine_similarity(query_vec, emb) if emb is not None else -1.0 + + records.sort(key=_score, reverse=True) + records = records[:limit_value] + return records async def delete_entries_by_session(self, session_id: str) -> int: @@ -1062,7 +1094,7 @@ async def delete_entries_older_than( raise async def _ensure_foreign_keys(self, connection: Any) -> None: - if id(connection) not in self._configured_connection_ids: + if connection not in self._configured_connections: await self._enable_foreign_keys(connection) async def _enable_foreign_keys(self, connection: Any) -> None: @@ -1070,7 +1102,7 @@ async def _enable_foreign_keys(self, connection: Any) -> None: await _close_cursor(await connection.execute("PRAGMA foreign_keys = ON")) for pragma_name, pragma_value in self._pragma_overrides: await _close_cursor(await connection.execute(f"PRAGMA {pragma_name} = {pragma_value}")) - self._configured_connection_ids.add(id(connection)) + self._configured_connections.add(connection) async def _memory_table_ddl(self) -> str: """Get SQLite CREATE TABLE SQL for memory entries. @@ -1121,7 +1153,8 @@ async def _memory_table_ddl(self) -> str: content_json TEXT NOT NULL, content_text TEXT NOT NULL, metadata_json TEXT, - inserted_at REAL NOT NULL + inserted_at REAL NOT NULL, + embedding BLOB ); CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_app_scope_user_time @@ -1243,6 +1276,34 @@ def _build_sqlite_scope_clause( return f"{prefix}app_name = ? AND {prefix}scope = 'app'", (app_name,) +def _pack_embedding(raw: Any) -> bytes | None: + """Pack float sequence into float32 binary BLOB.""" + if raw is None: + return None + return array.array("f", [float(x) for x in raw]).tobytes() + + +def _unpack_embedding(blob: Any) -> list[float] | None: + """Unpack float32 binary BLOB into float list.""" + if blob is None: + return None + return list(array.array("f", blob)) + + +def _cosine_similarity(v1: "Sequence[float]", v2: "Sequence[float]") -> float: + """Compute cosine similarity between two float vectors.""" + dot = 0.0 + norm1 = 0.0 + norm2 = 0.0 + for a, b in zip(v1, v2, strict=False): + dot += a * b + norm1 += a * a + norm2 += b * b + if norm1 == 0.0 or norm2 == 0.0: + return 0.0 + return float(dot / (math.sqrt(norm1) * math.sqrt(norm2))) + + def _session_list_query( session_table: str, app_name: str, diff --git a/sqlspec/adapters/aiosqlite/config.py b/sqlspec/adapters/aiosqlite/config.py index b8c04305c..a727d086a 100644 --- a/sqlspec/adapters/aiosqlite/config.py +++ b/sqlspec/adapters/aiosqlite/config.py @@ -17,7 +17,6 @@ build_connection_config, default_statement_config, render_pragmas, - require_python_version, ) from sqlspec.adapters.aiosqlite.driver import AiosqliteDriver, AiosqliteExceptionHandler from sqlspec.adapters.aiosqlite.pool import ( @@ -101,7 +100,7 @@ class AiosqlitePoolParams(AiosqliteConnectionParams): enable_optimizations: NotRequired[bool] """Apply WAL journal mode, ``synchronous=NORMAL``, and busy-timeout PRAGMAs on connect.""" enable_foreign_keys: NotRequired[bool] - """Execute ``PRAGMA foreign_keys = ON`` on every pooled connection (defaults to ``True``).""" + """Execute ``PRAGMA foreign_keys = ON`` on every pooled connection (defaults to ``False``).""" extra: NotRequired["dict[str, Any]"] """Additional keyword arguments forwarded directly to ``aiosqlite.connect``.""" @@ -408,13 +407,21 @@ async def _create_pool(self) -> AiosqliteConnectionPool: Returns: AiosqliteConnectionPool: The connection pool instance. """ - pool_size = self.connection_config.get("pool_size") or 5 + pool_size = self.connection_config.get("pool_size") + if pool_size is None: + pool_size = 5 min_size = self.connection_config.get("min_size") if min_size is None: min_size = 0 - connect_timeout = self.connection_config.get("connect_timeout") or 30.0 - idle_timeout = self.connection_config.get("idle_timeout") or 24 * 60 * 60 - operation_timeout = self.connection_config.get("operation_timeout") or 10.0 + connect_timeout = self.connection_config.get("connect_timeout") + if connect_timeout is None: + connect_timeout = 30.0 + idle_timeout = self.connection_config.get("idle_timeout") + if idle_timeout is None: + idle_timeout = 24 * 60 * 60 + operation_timeout = self.connection_config.get("operation_timeout") + if operation_timeout is None: + operation_timeout = 10.0 health_check_interval = self.connection_config.get("health_check_interval") if health_check_interval is None: health_check_interval = 30.0 @@ -518,8 +525,6 @@ def _build_runtime_setup(features: "dict[str, Any]") -> "dict[str, Any] | None": ) window_functions = runtime_setup.get("custom_window_functions", ()) _validate_entries(window_functions, ("name", "narg", "window_class"), "custom_window_functions") - if window_functions: - require_python_version("Custom SQLite window functions", (3, 11)) interval = runtime_setup.get("progress_handler_interval") if interval is not None and (not isinstance(interval, int) or isinstance(interval, bool) or interval < 1): diff --git a/sqlspec/adapters/aiosqlite/core.py b/sqlspec/adapters/aiosqlite/core.py index 6441691ad..25bb989df 100644 --- a/sqlspec/adapters/aiosqlite/core.py +++ b/sqlspec/adapters/aiosqlite/core.py @@ -6,7 +6,7 @@ from collections.abc import Mapping from datetime import UTC, date, datetime, timedelta from decimal import Decimal -from typing import TYPE_CHECKING, Any, TypeVar, cast +from typing import TYPE_CHECKING, Any, Final, TypeVar, cast from sqlglot import exp @@ -47,6 +47,7 @@ "SECONDS_PER_DAY", "SQLITE_CONNECT_SUPPORTS_AUTOCOMMIT", "AiosqliteStreamSource", + "_filter_connection_params", "apply_driver_features", "apply_extension_pragmas", "build_connection_config", @@ -95,8 +96,8 @@ SQLITE_CONSTRAINT_PRIMARYKEY_CODE = 1555 SQLITE_CONSTRAINT_ROWID_CODE = 2579 SQLITE_CONSTRAINT_FOREIGNKEY_CODE = 787 -SQLITE_CONSTRAINT_NOTNULL_CODE = 1811 -SQLITE_CONSTRAINT_CHECK_CODE = 531 +SQLITE_CONSTRAINT_NOTNULL_CODE = 1299 +SQLITE_CONSTRAINT_CHECK_CODE = 275 SQLITE_CONSTRAINT_DATATYPE_CODE = 3091 SQLITE_CONSTRAINT_CODE = 19 SQLITE_CANTOPEN_CODE = 14 @@ -108,6 +109,10 @@ SQLITE_INTERRUPT_CODE = 9 SQLITE_PERM_CODE = 3 SQLITE_READONLY_CODE = 8 +SQLITE_TOOBIG_CODE = 18 +SQLITE_AUTH_CODE = 23 +SQLITE_RANGE_CODE = 25 +SQLITE_NOTADB_CODE = 26 SQLITE_CONNECT_SUPPORTS_AUTOCOMMIT = sys.version_info >= (3, 12) SQLITE_DATABASE_LIST_MIN_COLUMNS = 2 SQLITE_TABLE_LIST_MIN_COLUMNS = 5 @@ -154,8 +159,8 @@ async def end_transaction( raw_conn = getattr(connection, "_conn", None) if isinstance(raw_conn, sqlite3.Connection) and callable(getattr(connection, "_execute", None)): conn_dict = getattr(connection, "__dict__", None) - has_method_override = isinstance(conn_dict, dict) and any( - key in conn_dict for key in ("commit", "rollback", "execute") + has_method_override = isinstance(conn_dict, dict) and ( + "commit" in conn_dict or "rollback" in conn_dict or "execute" in conn_dict ) if not has_method_override: await run_on_worker_thread( @@ -392,6 +397,37 @@ def normalize_execute_many_parameters(parameters: Any) -> Any: return parameters +_EXCLUDED_CONNECTION_KEYS: Final[frozenset[str]] = frozenset({ + "enable_foreign_keys", + "enable_optimizations", + "pool_size", + "min_size", + "connect_timeout", + "idle_timeout", + "operation_timeout", + "health_check_interval", + "extra", + "pool_min_size", + "pool_max_size", + "pool_timeout", + "pool_recycle_seconds", + "db", + "path", + "file", +}) + + +def _filter_connection_params(mapping: Mapping[str, Any]) -> dict[str, Any]: + """Filter connection parameters to only those accepted by SQLite/aiosqlite.""" + return { + key: value + for key, value in mapping.items() + if key not in _EXCLUDED_CONNECTION_KEYS + and (value is not None or key == "isolation_level") + and (key != "autocommit" or SQLITE_CONNECT_SUPPORTS_AUTOCOMMIT) + } + + def build_connection_config(connection_config: "Mapping[str, Any]") -> "dict[str, Any]": """Build connection configuration for pool creation. @@ -401,35 +437,7 @@ def build_connection_config(connection_config: "Mapping[str, Any]") -> "dict[str Returns: Dictionary with connection parameters. """ - excluded_keys = { - "enable_foreign_keys", - "enable_optimizations", - "pool_size", - "min_size", - "connect_timeout", - "idle_timeout", - "operation_timeout", - "health_check_interval", - "extra", - "pool_min_size", - "pool_max_size", - "pool_timeout", - "pool_recycle_seconds", - "db", - "path", - "file", - } - - def _filter_params(mapping: Mapping[str, Any]) -> dict[str, Any]: - return { - key: value - for key, value in mapping.items() - if key not in excluded_keys - and (value is not None or key == "isolation_level") - and (key != "autocommit" or SQLITE_CONNECT_SUPPORTS_AUTOCOMMIT) - } - - connection_parameters = _filter_params(connection_config) + connection_parameters = _filter_connection_params(connection_config) if "database" not in connection_parameters: database = ( connection_config.get("database") @@ -441,7 +449,7 @@ def _filter_params(mapping: Mapping[str, Any]) -> dict[str, Any]: connection_parameters["database"] = database extra = connection_config.get("extra") if isinstance(extra, Mapping): - connection_parameters.update(_filter_params(extra)) + connection_parameters.update(_filter_connection_params(extra)) return connection_parameters @@ -459,7 +467,7 @@ def create_mapped_exception(error: BaseException, *, logger: Any | None = None) Returns: A SQLSpec exception that wraps the original error """ - del logger + _ = logger if has_sqlite_error(error): error_code = error.sqlite_errorcode error_name = error.sqlite_errorname @@ -477,13 +485,19 @@ def create_mapped_exception(error: BaseException, *, logger: Any | None = None) return _create_aiosqlite_error(error, error_code, DeadlockError, "database busy") if primary_code == SQLITE_LOCKED_CODE or (isinstance(error_name, str) and error_name.startswith("SQLITE_LOCKED")): return _create_aiosqlite_error(error, error_code, DeadlockError, "database locked") - if "locked" in error_msg or "busy" in error_msg: - return _create_aiosqlite_error(error, error_code or 0, DeadlockError, "database locked") + if error_code is None and ( + "database is locked" in error_msg + or "database table is locked" in error_msg + or "database locked" in error_msg + or "database is busy" in error_msg + or "database busy" in error_msg + ): + return _create_aiosqlite_error(error, 0, DeadlockError, "database locked") if primary_code == SQLITE_INTERRUPT_CODE or error_name == "SQLITE_INTERRUPT": return _create_aiosqlite_error(error, error_code, OperationCancelledError, "query interrupted") - if "interrupt" in error_msg: - return _create_aiosqlite_error(error, error_code or 0, OperationCancelledError, "query interrupted") + if error_code is None and ("query was interrupted" in error_msg or "interrupted" in error_msg): + return _create_aiosqlite_error(error, 0, OperationCancelledError, "query interrupted") if primary_code == SQLITE_PERM_CODE or error_name == "SQLITE_PERM": return _create_aiosqlite_error(error, error_code, PermissionDeniedError, "permission denied") @@ -491,11 +505,20 @@ def create_mapped_exception(error: BaseException, *, logger: Any | None = None) isinstance(error_name, str) and error_name.startswith("SQLITE_READONLY") ): return _create_aiosqlite_error(error, error_code, PermissionDeniedError, "database is read-only") - if "permission denied" in error_msg or "readonly" in error_msg: - return _create_aiosqlite_error(error, error_code or 0, PermissionDeniedError, "permission denied") - - if primary_code == SQLITE_CANTOPEN_CODE or ( - isinstance(error_name, str) and error_name.startswith("SQLITE_CANTOPEN") + if error_code is None and ( + "attempt to write a readonly database" in error_msg + or "readonly database" in error_msg + or "permission denied" in error_msg + or "access permission denied" in error_msg + ): + return _create_aiosqlite_error(error, 0, PermissionDeniedError, "permission denied") + + if primary_code == SQLITE_AUTH_CODE or error_name == "SQLITE_AUTH": + return _create_aiosqlite_error(error, error_code, PermissionDeniedError, "authorization denied") + if primary_code in (SQLITE_RANGE_CODE, SQLITE_TOOBIG_CODE) or error_name in ("SQLITE_RANGE", "SQLITE_TOOBIG"): + return _create_aiosqlite_error(error, error_code, DataError, "data range or size exceeded") + if primary_code in (SQLITE_CANTOPEN_CODE, SQLITE_NOTADB_CODE) or ( + isinstance(error_name, str) and error_name.startswith(("SQLITE_CANTOPEN", "SQLITE_NOTADB")) ): return _create_aiosqlite_error(error, error_code, DatabaseConnectionError, "connection error") if primary_code in (SQLITE_IOERR_CODE, SQLITE_CORRUPT_CODE) or ( @@ -506,6 +529,10 @@ def create_mapped_exception(error: BaseException, *, logger: Any | None = None) return _create_aiosqlite_error(error, error_code, DataError, "data error") if "syntax" in error_msg: return _create_aiosqlite_error(error, error_code, SQLParsingError, "SQL syntax error") + if isinstance(error, sqlite3.ProgrammingError): + return _create_aiosqlite_error(error, error_code, DataError, "programming error") + if isinstance(error, sqlite3.InterfaceError): + return _create_aiosqlite_error(error, error_code, OperationalError, "interface error") if primary_code == 1 or isinstance(error, sqlite3.OperationalError): return _create_aiosqlite_error(error, error_code, OperationalError, "operational error") if isinstance(error, sqlite3.DataError): diff --git a/sqlspec/adapters/aiosqlite/data_dictionary.py b/sqlspec/adapters/aiosqlite/data_dictionary.py index b954d5022..54a6420e1 100644 --- a/sqlspec/adapters/aiosqlite/data_dictionary.py +++ b/sqlspec/adapters/aiosqlite/data_dictionary.py @@ -4,6 +4,7 @@ from mypy_extensions import mypyc_attr +from sqlspec.adapters.aiosqlite._typing import aiosqlite_sqlite_module as sqlite3 from sqlspec.data_dictionary import ( ColumnMetadata, DDLResult, @@ -26,10 +27,11 @@ sqlite_schema_ddl_results, ) from sqlspec.driver import AsyncDataDictionaryBase +from sqlspec.exceptions import SQLFileNotFoundError if TYPE_CHECKING: from sqlspec.adapters.aiosqlite.driver import AiosqliteDriver - from sqlspec.data_dictionary._types import DialectConfig + from sqlspec.data_dictionary import DialectConfig __all__ = ("AiosqliteDataDictionary",) @@ -41,6 +43,7 @@ class AiosqliteDataDictionary(AsyncDataDictionaryBase): """SQLite-specific async data dictionary.""" dialect: ClassVar[str] = "sqlite" + _shared_version_info: ClassVar["VersionInfo | None"] = None def __init__(self) -> None: super().__init__() @@ -68,11 +71,16 @@ async def get_version(self, driver: "AiosqliteDriver") -> "VersionInfo | None": Returns: SQLite version information or None if detection fails. """ + if AiosqliteDataDictionary._shared_version_info is not None: + return AiosqliteDataDictionary._shared_version_info + driver_id = id(driver) if driver_id in self._version_fetch_attempted: return self._version_cache.get(driver_id) - version_value = await driver.select_value_or_none(self.get_query("version", "current")) + version_value = getattr(sqlite3, "sqlite_version", None) + if not version_value: + version_value = await driver.select_value_or_none(self.get_query("version", "current")) if not version_value: self._log_version_unavailable(type(self).dialect, "missing") self.cache_version(driver_id, None) @@ -86,6 +94,7 @@ async def get_version(self, driver: "AiosqliteDriver") -> "VersionInfo | None": self._log_version_detected(type(self).dialect, version_info) self.cache_version(driver_id, version_info) + AiosqliteDataDictionary._shared_version_info = version_info return version_info async def get_feature_flag(self, driver: "AiosqliteDriver", feature: str) -> bool: @@ -142,10 +151,13 @@ async def get_columns( rows = await driver.select(query_text, table_name=table, schema_name=schema_name) if any(str(row.get("extra") or "").endswith("generated") for row in rows): table_sql = await self._get_table_sql(driver, schema_name, table) + parsed_expressions: dict[str, dict[str, str]] = {} for row in rows: if str(row.get("extra") or "").endswith("generated"): - expressions = sqlite_generated_expressions(table_sql.get(str(row.get("table_name")))) - expression = expressions.get(str(row.get("column_name"))) + table_name_str = str(row.get("table_name")) + if table_name_str not in parsed_expressions: + parsed_expressions[table_name_str] = sqlite_generated_expressions(table_sql.get(table_name_str)) + expression = parsed_expressions[table_name_str].get(str(row.get("column_name"))) if expression is not None: row["generation_expression"] = expression return cast("list[ColumnMetadata]", rows) @@ -278,7 +290,7 @@ def _get_domain_query_text(self, domain: str, query_name: str) -> str: query_text = get_data_dictionary_loader().get_domain_query_text(type(self).dialect, domain, query_name) if query_text is None: msg = f"Missing SQLite data-dictionary query: {domain}/{query_name}" - raise RuntimeError(msg) + raise SQLFileNotFoundError(msg) return query_text diff --git a/sqlspec/adapters/aiosqlite/driver.py b/sqlspec/adapters/aiosqlite/driver.py index 98a1f4697..97240fc14 100644 --- a/sqlspec/adapters/aiosqlite/driver.py +++ b/sqlspec/adapters/aiosqlite/driver.py @@ -38,7 +38,7 @@ from collections.abc import Mapping, Sequence from sqlspec.adapters.aiosqlite._typing import AiosqliteConnection - from sqlspec.core import SQL, SQLResult, StatementConfig + from sqlspec.core import SQL, CachedQuery, SQLResult, StatementConfig from sqlspec.core.compiler import OperationType from sqlspec.driver import ExecutionResult from sqlspec.storage import StorageBridgeJob, StorageDestination, StorageFormat, StorageTelemetry @@ -161,23 +161,11 @@ async def dispatch_execute_many(self, cursor: "AiosqliteRawCursor", statement: " async def dispatch_execute_script(self, cursor: "AiosqliteRawCursor", statement: "SQL") -> "ExecutionResult": """Execute SQL script.""" self._rowid_target_cache.clear() - sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) - statements = self.split_script_statements(sql, statement.statement_config, strip_trailing_semicolon=True) - - successful_count = 0 - last_cursor = cursor - try: - for stmt in statements: - await cursor.execute(stmt, normalize_execute_parameters(prepared_parameters)) - successful_count += 1 + return await super().dispatch_execute_script(cursor, statement) finally: self._rowid_target_cache.clear() - return self.create_execution_result( - last_cursor, statement_count=len(statements), successful_statements=successful_count, is_script_result=True - ) - async def begin(self, mode: "Literal['DEFERRED', 'IMMEDIATE', 'EXCLUSIVE'] | None" = None) -> None: """Begin a database transaction. @@ -307,20 +295,19 @@ async def load_from_arrow( try: if owns_transaction: await self.begin() - if overwrite: - statement = f"DELETE FROM {format_identifier(table)}" - async with self.with_cursor(self.connection) as cursor: + async with self.with_cursor(self.connection) as cursor: + if overwrite: + statement = f"DELETE FROM {format_identifier(table)}" await cursor.execute(statement) - for batch in arrow_table.to_batches(max_chunksize=batch_size): - pydict = batch.to_pydict() - records = list(zip(*(pydict[col] for col in columns), strict=False)) - if records: - prepared_records = ( - self.prepare_driver_parameters(records, self.statement_config, is_many=True) - if needs_prep - else records - ) - async with self.with_cursor(self.connection) as cursor: + for batch in arrow_table.to_batches(max_chunksize=batch_size): + pydict = batch.to_pydict() + records = list(zip(*(pydict[col] for col in columns), strict=False)) + if records: + prepared_records = ( + self.prepare_driver_parameters(records, self.statement_config, is_many=True) + if needs_prep + else records + ) await cursor.executemany(insert_sql, cast("Any", prepared_records)) if owns_transaction: await self.commit() @@ -336,22 +323,6 @@ async def load_from_arrow( self._attach_partition_telemetry(telemetry_payload, partitioner) return self._storage_job(telemetry_payload, telemetry) - async def load_from_storage( - self, - table: str, - source: "StorageDestination", - *, - file_format: "StorageFormat", - partitioner: "dict[str, object] | None" = None, - overwrite: bool = False, - **kwargs: Any, - ) -> "StorageBridgeJob": - """Load staged artifacts from storage into SQLite.""" - - return await super().load_from_storage( - table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs - ) - def collect_rows(self, cursor: Any, fetched: "list[Any]") -> "tuple[list[Any], list[str], int]": """Collect aiosqlite rows for the direct execution path.""" return collect_rows(fetched, cursor.description) @@ -360,53 +331,47 @@ def resolve_rowcount(self, cursor: Any) -> int: """Resolve rowcount from aiosqlite cursor for the direct execution path.""" return resolve_rowcount(cursor) - async def _execute_cache_hit( - self, sql: str, params: "tuple[Any, ...] | list[Any] | dict[str, Any]", cached: Any + async def _execute_cache_hit_with_cursor( + self, sql: str, params: "tuple[Any, ...] | list[Any] | dict[str, Any]", cached: "CachedQuery" ) -> "SQLResult": - """Execute cached queries through the async cursor fast path.""" + """Execute cached queries through the aiosqlite worker thread fast path.""" + wire_sql = self._format_wire_sql(cached.compiled_sql) normalized_parameters = normalize_execute_parameters(params) - exc_handler = self.handle_database_exceptions() - result: SQLResult | None = None self._invalidate_rowid_target_cache(cached.operation_type) - async with exc_handler: - if cached.operation_profile.returns_rows: - fetched_data, description, _affected_rows, last_inserted_id = await run_on_worker_thread( - self.connection, - _execute_fetchall_with_metadata, - self.connection, - cached.compiled_sql, - normalized_parameters, - cached.operation_type, - cached.processed_state.parsed_expression, - self._rowid_target_cache, - ) - data, column_names, row_count = collect_rows(fetched_data, description) - result = self.build_cached_select_result( - sql, - params, - cached, - selected_data=data, - column_names=column_names, - data_row_count=row_count, - row_format=resolve_row_format(data), - last_inserted_id=last_inserted_id, - ) - else: - affected_rows, last_inserted_id = await run_on_worker_thread( - self.connection, - _execute_and_resolve_metadata, - self.connection, - cached.compiled_sql, - normalized_parameters, - cached.operation_type, - cached.processed_state.parsed_expression, - self._rowid_target_cache, - ) - result = DMLResult(cached.operation_type, affected_rows, last_inserted_id) - - self._check_pending_exception(exc_handler) - assert result is not None - return result + if cached.operation_profile.returns_rows: + fetched_data, description, _affected_rows, last_inserted_id = await run_on_worker_thread( + self.connection, + _execute_fetchall_with_metadata, + self.connection, + wire_sql, + normalized_parameters, + cached.operation_type, + cached.processed_state.parsed_expression, + self._rowid_target_cache, + ) + data, column_names, row_count = collect_rows(fetched_data, description) + return self.build_cached_select_result( + sql, + params, + cached, + selected_data=data, + column_names=column_names, + data_row_count=row_count, + row_format=resolve_row_format(data), + last_inserted_id=last_inserted_id, + compiled_sql=wire_sql, + ) + affected_rows, last_inserted_id = await run_on_worker_thread( + self.connection, + _execute_and_resolve_metadata, + self.connection, + wire_sql, + normalized_parameters, + cached.operation_type, + cached.processed_state.parsed_expression, + self._rowid_target_cache, + ) + return DMLResult(cached.operation_type, affected_rows, last_inserted_id) def _invalidate_rowid_target_cache(self, operation_type: "OperationType") -> None: if operation_type not in {"SELECT", "INSERT", "UPDATE", "DELETE"}: diff --git a/sqlspec/adapters/aiosqlite/litestar/store.py b/sqlspec/adapters/aiosqlite/litestar/store.py index 254461451..873455fa8 100644 --- a/sqlspec/adapters/aiosqlite/litestar/store.py +++ b/sqlspec/adapters/aiosqlite/litestar/store.py @@ -2,6 +2,7 @@ import contextlib import inspect +import weakref from datetime import UTC, datetime, timedelta from typing import TYPE_CHECKING, Any, NotRequired @@ -50,7 +51,7 @@ class AiosqliteStore(BaseSQLSpecStore["AiosqliteConfig"]): config: AiosqliteConfig instance. """ - __slots__ = ("_configured_connection_ids", "_pragma_statements") + __slots__ = ("_configured_connections", "_pragma_statements") extension_config_options = BaseSQLSpecStore.extension_config_options | frozenset({ "pragma_overrides", "pragma_profile", @@ -64,7 +65,7 @@ def __init__(self, config: "AiosqliteConfig") -> None: """ super().__init__(config) self._pragma_statements = extension_pragma_statements(config, "litestar") - self._configured_connection_ids: set[int] = set() + self._configured_connections: weakref.WeakSet[Any] = weakref.WeakSet() async def create_table(self) -> None: """Create the session table if it doesn't exist.""" @@ -74,7 +75,7 @@ async def create_table(self) -> None: sql = self._table_ddl() async with self._config.provide_session() as driver: await apply_extension_pragmas(driver.connection, self._pragma_statements) - self._configured_connection_ids.add(id(driver.connection)) + self._configured_connections.add(driver.connection) await driver.execute_script(sql) await driver.commit() self._log_table_created() @@ -83,7 +84,7 @@ async def create_table(self) -> None: async def prepare_schema_async(self, driver: Any) -> None: """Apply configured SQLite PRAGMAs before migration DDL generation.""" await apply_extension_pragmas(driver.connection, self._pragma_statements) - self._configured_connection_ids.add(id(driver.connection)) + self._configured_connections.add(driver.connection) async def get(self, key: str, renew_for: "int | timedelta | None" = None) -> "bytes | None": """Get a session value by key. @@ -251,9 +252,9 @@ async def delete_expired(self) -> int: return count async def _ensure_pragmas(self, connection: Any) -> None: - if self._pragma_statements and id(connection) not in self._configured_connection_ids: + if self._pragma_statements and connection not in self._configured_connections: await apply_extension_pragmas(connection, self._pragma_statements) - self._configured_connection_ids.add(id(connection)) + self._configured_connections.add(connection) def _table_ddl(self) -> str: """Get SQLite CREATE TABLE SQL. diff --git a/sqlspec/adapters/aiosqlite/pool.py b/sqlspec/adapters/aiosqlite/pool.py index 2c166cc9f..d69530830 100644 --- a/sqlspec/adapters/aiosqlite/pool.py +++ b/sqlspec/adapters/aiosqlite/pool.py @@ -419,8 +419,7 @@ async def release(self, connection: AiosqlitePoolConnection) -> None: try: raw_conn = connection.connection if getattr(raw_conn, "in_transaction", False): - with suppress(Exception): - await end_transaction(raw_conn, commit=False) + await end_transaction(raw_conn, commit=False) connection.idle_since = time.time() self._queue.put_nowait(connection) except Exception as e: diff --git a/sqlspec/adapters/aiosqlite/type_converter.py b/sqlspec/adapters/aiosqlite/type_converter.py index 23fb2d079..e83d86713 100644 --- a/sqlspec/adapters/aiosqlite/type_converter.py +++ b/sqlspec/adapters/aiosqlite/type_converter.py @@ -13,15 +13,12 @@ from typing import TYPE_CHECKING, Any from sqlspec.adapters.aiosqlite._typing import aiosqlite_sqlite_module as sqlite3 -from sqlspec.utils.logging import get_logger if TYPE_CHECKING: from collections.abc import Callable __all__ = ("json_adapter", "json_converter", "register_type_handlers", "unregister_type_handlers") -logger = get_logger(__name__) - DEFAULT_JSON_TYPE = "JSON" diff --git a/sqlspec/adapters/arrow_odbc/_typing.py b/sqlspec/adapters/arrow_odbc/_typing.py index 0be32aa33..e012b40f4 100644 --- a/sqlspec/adapters/arrow_odbc/_typing.py +++ b/sqlspec/adapters/arrow_odbc/_typing.py @@ -1,15 +1,11 @@ -# pyright: reportAttributeAccessIssue=false """arrow-odbc adapter type definitions and mypyc-excluded context managers.""" from typing import TYPE_CHECKING, Any -import arrow_odbc as _arrow_odbc # pyright: ignore[reportMissingImports] +import arrow_odbc as _arrow_odbc from arrow_odbc import Error as ArrowOdbcError -from arrow_odbc import ( - TextEncoding, # pyright: ignore[reportMissingImports] - enable_odbc_connection_pooling, # pyright: ignore[reportMissingImports] -) -from arrow_odbc import connect as arrow_odbc_connect # pyright: ignore[reportMissingImports] +from arrow_odbc import TextEncoding, enable_odbc_connection_pooling +from arrow_odbc import connect as arrow_odbc_connect if TYPE_CHECKING: from collections.abc import Callable diff --git a/sqlspec/adapters/arrow_odbc/adk/store.py b/sqlspec/adapters/arrow_odbc/adk/store.py index ced882b1b..25988fad2 100644 --- a/sqlspec/adapters/arrow_odbc/adk/store.py +++ b/sqlspec/adapters/arrow_odbc/adk/store.py @@ -1,6 +1,5 @@ """arrow-odbc ADK stores for Google Agent Development Kit session storage.""" -import contextlib from datetime import UTC, date, datetime from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, NotRequired, cast @@ -11,15 +10,18 @@ extract_native_error_number, split_db2_name, ) +from sqlspec.adapters.arrow_odbc.core import atomic_store_write as _atomic_write +from sqlspec.adapters.arrow_odbc.core import quote_bracket_identifier as _quote_identifier +from sqlspec.adapters.arrow_odbc.core import strip_bracket_quotes as _strip_tsql_quotes from sqlspec.config import ADKConfig -from sqlspec.exceptions import SQLSpecError +from sqlspec.exceptions import ImproperConfigurationError, SQLSpecError from sqlspec.extensions.adk import BaseSyncADKStore, StoredEvent, StoredSession, normalize_session_list_options from sqlspec.extensions.adk.memory import BaseSyncADKMemoryStore, StoredMemory from sqlspec.utils.serializers import from_json, to_json from sqlspec.utils.text import split_qualified_identifier if TYPE_CHECKING: - from collections.abc import Callable, Iterator, Sequence + from collections.abc import Callable, Sequence from datetime import timedelta from sqlspec.adapters.arrow_odbc.config import ArrowOdbcConfig @@ -54,6 +56,8 @@ class ArrowOdbcADKStore(BaseSyncADKStore["ArrowOdbcConfig"]): connector_name: ClassVar[str] = "arrow_odbc" __slots__ = ("_sql",) + _max_table_name_parts = 2 + def __init__(self, config: "ArrowOdbcConfig") -> None: super().__init__(config) self._sql = _adk_sql_for(config) @@ -432,6 +436,8 @@ class ArrowOdbcADKMemoryStore(BaseSyncADKMemoryStore["ArrowOdbcConfig"]): __slots__ = ("_sql",) + _max_table_name_parts = 2 + def __init__(self, config: "ArrowOdbcConfig") -> None: super().__init__(config) self._sql = _adk_sql_for(config) @@ -501,10 +507,17 @@ def search_entries( scope_filter: Literal["all", "user", "app"] = "all", embedding: "Sequence[float] | None" = None, ) -> "list[StoredMemory]": - """Search memory entries with LIKE matching.""" + """Search memory entries with LIKE matching. + + Vector embeddings are not supported in ArrowOdbcADKMemoryStore; passing + an embedding raises NotImplementedError. + """ if not self._enabled: msg = "Memory store is disabled" raise RuntimeError(msg) + if embedding is not None: + msg = "Vector embeddings are not supported in ArrowOdbcADKMemoryStore; only LIKE search is supported." + raise NotImplementedError(msg) if not query or not query.strip(): return [] effective_limit = max(0, int(limit if limit is not None else self._max_results)) @@ -945,19 +958,6 @@ def _is_table_missing(exc: BaseException) -> bool: return "invalid object name" in text or "42s02" in text or "(208)" in text -def _quote_identifier(identifier: str) -> str: - return f"[{identifier.replace(']', ']]')}]" - - -def _strip_tsql_quotes(part: str) -> str: - stripped = part.strip() - if len(stripped) >= _QUALIFIED_PARTS_MIN and stripped[0] == "[" and stripped[-1] == "]": - return stripped[1:-1].replace("]]", "]") - if len(stripped) >= _QUALIFIED_PARTS_MIN and stripped[0] == '"' and stripped[-1] == '"': - return stripped[1:-1].replace('""', '"') - return stripped - - def _split_tsql_table(table: str) -> "tuple[str, str]": parts = split_qualified_identifier(table.strip()) if len(parts) >= _QUALIFIED_PARTS_MIN: @@ -977,12 +977,6 @@ def _constraint_ref(prefix: str, table: str, suffix: str) -> str: return _quote_identifier(f"{prefix}_{bare_table}_{suffix}") -def _bare_db2_table(table: str) -> str: - parts = split_qualified_identifier(table.strip()) - raw = parts[-1] if parts else table - return _strip_tsql_quotes(raw) - - def _raise_session_not_found(session_id: str) -> None: msg = f"Session {session_id} not found during append_event_and_update_state." raise ValueError(msg) @@ -1027,29 +1021,6 @@ def _session_list_query( return sql, tuple(params) -@contextlib.contextmanager -def _atomic_write(driver: Any) -> "Iterator[None]": - """Wrap multi-statement store mutations in an explicit transaction when supported. - - Yields: - None while the atomic write block executes. - """ - in_tx = getattr(driver, "_connection_in_transaction", None) - if callable(in_tx) and in_tx() is True: - yield - return - if getattr(driver, "_connection_autocommit", True) is False: - driver.begin() - try: - yield - except BaseException: - with contextlib.suppress(Exception): - driver.rollback() - raise - else: - driver.commit() - - class _TsqlAdkSql: """SQL Server statements for the ADK stores.""" @@ -1275,8 +1246,8 @@ def events_query( return sql, tuple(params) def session_index_specs(self, session_table: str, events_table: str) -> "list[tuple[str, str, str]]": - bare_session = _bare_db2_table(session_table) - bare_events = _bare_db2_table(events_table) + bare_session = split_db2_name(session_table)[1].lower() + bare_events = split_db2_name(events_table)[1].lower() return [ (f"idx_{bare_session}_app_user", session_table, "app_name, user_id"), (f"idx_{bare_session}_update_time", session_table, "update_time"), @@ -1289,7 +1260,7 @@ def session_index_specs(self, session_table: str, events_table: str) -> "list[tu def sessions_table_ddl(self, table: str, owner_id_column_ddl: "str | None") -> str: owner_line = f", {owner_id_column_ddl}" if owner_id_column_ddl else "" - bare_table = _bare_db2_table(table) + bare_table = split_db2_name(table)[1].lower() return ( f"CREATE TABLE {table} (id VARCHAR(128) NOT NULL, app_name VARCHAR(128) NOT NULL, " f"user_id VARCHAR(128) NOT NULL{owner_line}, state {DB2_JSON_COLUMN_TYPE} NOT NULL, " @@ -1299,7 +1270,7 @@ def sessions_table_ddl(self, table: str, owner_id_column_ddl: "str | None") -> s ) def events_table_ddl(self, table: str, session_table: str) -> str: - bare_table = _bare_db2_table(table) + bare_table = split_db2_name(table)[1].lower() return ( f"CREATE TABLE {table} (id VARCHAR(128) NOT NULL, app_name VARCHAR(128) NOT NULL, " "user_id VARCHAR(128) NOT NULL, session_id VARCHAR(128) NOT NULL, " @@ -1309,7 +1280,7 @@ def events_table_ddl(self, table: str, session_table: str) -> str: ) def app_states_table_ddl(self, table: str) -> str: - bare_table = _bare_db2_table(table) + bare_table = split_db2_name(table)[1].lower() return ( f"CREATE TABLE {table} (app_name VARCHAR(128) NOT NULL, state {DB2_JSON_COLUMN_TYPE} NOT NULL, " "update_time TIMESTAMP NOT NULL DEFAULT CURRENT TIMESTAMP, " @@ -1317,7 +1288,7 @@ def app_states_table_ddl(self, table: str) -> str: ) def user_states_table_ddl(self, table: str) -> str: - bare_table = _bare_db2_table(table) + bare_table = split_db2_name(table)[1].lower() return ( f"CREATE TABLE {table} (app_name VARCHAR(128) NOT NULL, user_id VARCHAR(128) NOT NULL, " f"state {DB2_JSON_COLUMN_TYPE} NOT NULL, update_time TIMESTAMP NOT NULL DEFAULT CURRENT TIMESTAMP, " @@ -1325,7 +1296,7 @@ def user_states_table_ddl(self, table: str) -> str: ) def metadata_table_ddl(self, table: str) -> str: - bare_table = _bare_db2_table(table) + bare_table = split_db2_name(table)[1].lower() return ( f'CREATE TABLE {table} ("KEY" VARCHAR(128) NOT NULL, value VARCHAR(512) NOT NULL, ' f'CONSTRAINT pk_{bare_table}_key PRIMARY KEY ("KEY"))' @@ -1336,7 +1307,7 @@ def drop_table_sql(self, table: str) -> str: def memory_table_ddl(self, table: str, owner_id_column_ddl: "str | None") -> str: owner_line = f", {owner_id_column_ddl}" if owner_id_column_ddl else "" - bare_table = _bare_db2_table(table) + bare_table = split_db2_name(table)[1].lower() return ( f"CREATE TABLE {table} (id VARCHAR(128) NOT NULL, session_id VARCHAR(128) NOT NULL, " "app_name VARCHAR(128) NOT NULL, user_id VARCHAR(128) NOT NULL, " @@ -1348,7 +1319,7 @@ def memory_table_ddl(self, table: str, owner_id_column_ddl: "str | None") -> str ) def memory_index_specs(self, table: str) -> "list[tuple[str, str, str]]": - bare_table = _bare_db2_table(table) + bare_table = split_db2_name(table)[1].lower() return [ (f"idx_{bare_table}_scope", table, "app_name, user_id, scope, timestamp"), (f"idx_{bare_table}_session", table, "session_id, timestamp"), @@ -1374,6 +1345,11 @@ def memory_search_sql( def _adk_sql_for(config: "ArrowOdbcConfig") -> "_TsqlAdkSql | _Db2AdkSql": """Return the statement provider for the config's resolved dialect.""" - if str(config.statement_config.dialect).lower() == "db2": + dialect_attr = getattr(config.statement_config, "dialect", None) + dialect = str(dialect_attr).lower() if isinstance(dialect_attr, str) else "mssql" + if dialect == "db2": return _Db2AdkSql() - return _TsqlAdkSql() + if dialect in {"mssql", "tsql"}: + return _TsqlAdkSql() + msg = f"ArrowOdbcADKStore requires dialect 'mssql' or 'db2' (received {dialect!r})." + raise ImproperConfigurationError(msg) diff --git a/sqlspec/adapters/arrow_odbc/config.py b/sqlspec/adapters/arrow_odbc/config.py index 50bad63dd..5ed6174fd 100644 --- a/sqlspec/adapters/arrow_odbc/config.py +++ b/sqlspec/adapters/arrow_odbc/config.py @@ -1,7 +1,7 @@ """arrow-odbc database configuration.""" import threading -from typing import TYPE_CHECKING, Any, ClassVar, NotRequired, TypedDict, cast +from typing import TYPE_CHECKING, Any, ClassVar, Literal, NotRequired, TypedDict, cast from sqlspec.adapters.arrow_odbc._typing import ( ArrowOdbcConnection, @@ -86,7 +86,7 @@ class ArrowOdbcDriverFeatures(TypedDict): """``arrow-odbc`` driver feature flags controlling Arrow batch sizing, dialect hints, and codecs.""" chunk_size: NotRequired[int] - """Maximum rows per Arrow ``RecordBatch`` fetched from ODBC cursors (defaults to ``65535``).""" + """Maximum rows per Arrow ``RecordBatch`` fetched from ODBC cursors (defaults to ``65536``).""" max_bytes_per_batch: NotRequired[int] """Upper memory budget in bytes per Arrow ``RecordBatch`` buffer.""" max_text_size: NotRequired[int] @@ -99,7 +99,7 @@ class ArrowOdbcDriverFeatures(TypedDict): """Per-statement ODBC query execution timeout in seconds.""" payload_text_encoding: NotRequired["TextEncoding"] """Wire text encoding override (``'auto'``, ``'utf8'``, or ``'utf16'``) passed to ``arrow-odbc``.""" - enable_driver_pooling: NotRequired[bool] + enable_pooling: NotRequired[bool] """Enable process-global ODBC driver-manager connection pooling prior to first connect (defaults to ``False``).""" connection_string: NotRequired[str] """Resolved connection string inspected for SQL dialect auto-detection.""" @@ -115,6 +115,8 @@ class ArrowOdbcDriverFeatures(TypedDict): """Deserializer parsing JSON text payloads (defaults to ``sqlspec.utils.serializers.from_json``).""" enable_events: NotRequired[bool] """Enable SQLSpec table-backed event channel support.""" + events_backend: NotRequired[Literal["poll_queue"]] + """Event queue backend name (``'poll_queue'``).""" on_connection_create: "NotRequired[Callable[[ArrowOdbcConnection], None]]" """Callback invoked with the raw ``arrow_odbc.Connection`` immediately after connect.""" @@ -239,7 +241,7 @@ def create_connection(self) -> "ArrowOdbcConnection": if self.connection_instance is not None: return cast("ArrowOdbcConnection", self.connection_instance) try: - _apply_driver_pooling(bool(self.driver_features.get("enable_driver_pooling", False))) + _apply_driver_pooling(bool(self.driver_features.get("enable_pooling", False))) dialect = _resolve_connection_dialect_hint(self.driver_features) connection_string, connect_kwargs = build_connection_config(self.connection_config, dialect=dialect) connection = cast("ArrowOdbcConnection", arrow_odbc_connect(connection_string, **connect_kwargs)) diff --git a/sqlspec/adapters/arrow_odbc/core.py b/sqlspec/adapters/arrow_odbc/core.py index 8ade8209d..c991a393b 100644 --- a/sqlspec/adapters/arrow_odbc/core.py +++ b/sqlspec/adapters/arrow_odbc/core.py @@ -1,5 +1,6 @@ """arrow-odbc adapter core helpers.""" +import contextlib import re from datetime import UTC, datetime from typing import TYPE_CHECKING, Any, Final @@ -24,12 +25,13 @@ TransactionError, UniqueViolationError, ) +from sqlspec.typing import ArrowException from sqlspec.utils.config_tools import parse_odbc_connection_string from sqlspec.utils.serializers import from_json, to_json from sqlspec.utils.type_converters import build_uuid_coercions if TYPE_CHECKING: - from collections.abc import Callable, Mapping + from collections.abc import Callable, Iterator, Mapping from sqlspec.core import StatementConfig @@ -37,6 +39,7 @@ "DB2_INDEX_EXISTS_SQL", "DB2_TABLE_EXISTS_SQL", "apply_driver_features", + "atomic_store_write", "build_connection_config", "build_profile", "build_statement_config", @@ -46,8 +49,10 @@ "driver_profile", "extract_native_error_number", "normalize_column_names", + "quote_bracket_identifier", "resolve_dialect_from_dbms_name", "split_db2_name", + "strip_bracket_quotes", ) DB2_TABLE_EXISTS_SQL: Final[str] = ( @@ -176,7 +181,9 @@ def resolve_dialect_from_dbms_name(dbms_name: str | None) -> str: def create_mapped_exception(error: BaseException, *, logger: Any | None = None) -> SQLSpecError: """Map an arrow-odbc exception to SQLSpec's exception hierarchy.""" - del logger + _ = logger + if isinstance(error, ArrowException): + return DataError(f"PyArrow data error. Original error: {error}") message = str(error) if _is_sql_server_diagnostic(message): error_number = extract_native_error_number(error) @@ -440,8 +447,9 @@ def _match_dialect_from_dbms_name(dbms_name: str | None) -> str | None: return None lowered = driver_name.lower() for dialect, patterns in _DIALECT_PATTERNS: - if any(pattern in lowered for pattern in patterns): - return dialect + for pattern in patterns: + if pattern in lowered: + return dialect return None @@ -657,3 +665,52 @@ def _format_connection_value(option_name: str, value: Any) -> str: driver_profile = build_profile() default_statement_config = build_statement_config() + + +_MIN_DELIMITER_LEN: Final = 2 + + +def quote_bracket_identifier(identifier: str) -> str: + """Quote an identifier using SQL Server / T-SQL bracket syntax.""" + return f"[{identifier.replace(']', ']]')}]" + + +def strip_bracket_quotes(part: str) -> str: + """Strip T-SQL bracket or double-quote enclosing delimiters from an identifier.""" + stripped = part.strip() + if len(stripped) >= _MIN_DELIMITER_LEN and stripped[0] == "[" and stripped[-1] == "]": + return stripped[1:-1].replace("]]", "]") + if len(stripped) >= _MIN_DELIMITER_LEN and stripped[0] == '"' and stripped[-1] == '"': + return stripped[1:-1].replace('""', '"') + return stripped + + +@contextlib.contextmanager +def atomic_store_write(driver: Any) -> "Iterator[None]": + """Wrap multi-statement store mutations in an explicit transaction when supported. + + Yields: + None while the atomic write block executes. + """ + in_tx = getattr(driver, "_connection_in_transaction", None) + if callable(in_tx) and in_tx() is True: + yield + return + + dialect = getattr(driver, "_dialect", "") + autocommit = getattr(driver, "_connection_autocommit", True) + can_begin = not (autocommit and dialect in {"db2", "oracle"}) + + if can_begin and hasattr(driver, "begin"): + driver.begin() + + try: + yield + except BaseException: + if hasattr(driver, "rollback"): + with contextlib.suppress(Exception): + driver.rollback() + raise + else: + if hasattr(driver, "commit"): + driver.commit() diff --git a/sqlspec/adapters/arrow_odbc/data_dictionary.py b/sqlspec/adapters/arrow_odbc/data_dictionary.py index 7b503551b..cf16cee5c 100644 --- a/sqlspec/adapters/arrow_odbc/data_dictionary.py +++ b/sqlspec/adapters/arrow_odbc/data_dictionary.py @@ -30,13 +30,9 @@ if TYPE_CHECKING: from collections.abc import Sequence - import pyarrow as pa - from sqlspec.adapters.arrow_odbc.driver import ArrowOdbcDriver from sqlspec.core import SQL - from sqlspec.data_dictionary._types import DialectConfig -else: - pa = import_optional("pyarrow") + from sqlspec.data_dictionary import DialectConfig __all__ = ("ArrowOdbcDataDictionary",) @@ -243,17 +239,12 @@ def _probe_columns(self, driver: "ArrowOdbcDriver", table: str, schema: "str | N ) probe_sql = f"SELECT * FROM {qualified} WHERE 1=0" try: - reader = driver._read_arrow_batches( # pyright: ignore[reportPrivateUsage] - probe_sql, None, 1, fetch_concurrently=False - ) + reader = driver._read_arrow_batches(probe_sql, None, 1, fetch_concurrently=False) except Exception: return [] columns: list[ColumnMetadata] = [] raw_names = [field.name for field in reader.schema] - normalized_names = normalize_column_names( - raw_names, - driver._lowercase_column_names, # pyright: ignore[reportPrivateUsage] - ) + normalized_names = normalize_column_names(raw_names, driver._lowercase_column_names) for position, (field, column_name) in enumerate(zip(reader.schema, normalized_names, strict=False), start=1): entry: ColumnMetadata = { "table_name": table, @@ -299,7 +290,7 @@ def _capability_for_domain(self, domain: str) -> MetadataCapability: source=MetadataSource.DRIVER_METADATA, warnings=(_ODBC_COLUMN_PARTIAL_WARNING,), ) - if domain in {"tables", "indexes", "foreign_keys"} and self._has_query(domain, "by_schema"): + if self._has_query(domain, "by_schema"): return MetadataCapability( domain=domain, support=MetadataSupport.SUPPORTED, @@ -324,6 +315,8 @@ def _has_query(self, domain: str, operation: str) -> bool: def _arrow_type_to_sql(data_type: Any) -> str: ensure_pyarrow() + pa = import_optional("pyarrow") + assert pa is not None types = pa.types if types.is_boolean(data_type): return "BOOLEAN" diff --git a/sqlspec/adapters/arrow_odbc/driver.py b/sqlspec/adapters/arrow_odbc/driver.py index fa25a1361..fe3b08a7e 100644 --- a/sqlspec/adapters/arrow_odbc/driver.py +++ b/sqlspec/adapters/arrow_odbc/driver.py @@ -2,6 +2,7 @@ import contextlib import re +import sys from collections.abc import Iterable, Mapping, Sequence from datetime import UTC, datetime from itertools import chain @@ -28,20 +29,16 @@ from sqlspec.core.parameters._validator import ParameterValidator from sqlspec.driver import BaseSyncExceptionHandler, SyncDriverAdapterBase, SyncRowStream, validate_savepoint_name from sqlspec.exceptions import DataError, ImproperConfigurationError, SQLSpecError -from sqlspec.typing import import_optional +from sqlspec.typing import ArrowException from sqlspec.utils.module_loader import ensure_pyarrow from sqlspec.utils.text import quote_identifier, split_qualified_identifier if TYPE_CHECKING: - import pyarrow as pa - from sqlspec.builder import QueryBuilder from sqlspec.core import ArrowResult, ParameterProfile, Statement, StatementConfig, StatementFilter from sqlspec.driver import ExecutionResult from sqlspec.storage import StorageBridgeJob, StorageDestination, StorageFormat, StorageTelemetry from sqlspec.typing import ArrowRecordBatch, ArrowRecordBatchReader, ArrowReturnFormat, StatementParameters -else: - pa = import_optional("pyarrow") __all__ = ("ArrowOdbcCursor", "ArrowOdbcDriver", "ArrowOdbcExceptionHandler", "resolve_dialect_from_dbms_name") @@ -54,7 +51,7 @@ class ArrowOdbcExceptionHandler(BaseSyncExceptionHandler): def _handle_exception(self, exc_type: "type[BaseException] | None", exc_val: "BaseException") -> bool: if exc_type is None: return False - if isinstance(exc_val, ArrowOdbcError) or (pa is not None and isinstance(exc_val, pa.ArrowException)): + if isinstance(exc_val, (ArrowOdbcError, ArrowException)): self.pending_exception = create_mapped_exception(exc_val) return True return False @@ -77,10 +74,20 @@ def __init__( def start(self) -> None: handler = self._driver.handle_database_exceptions() - with handler: - raw_reader = self._driver._read_arrow_batches(self._sql, self._parameters, self._chunk_size) - self._raw_reader = raw_reader - self._reader = iter(self._driver._normalize_reader(_to_pyarrow_reader(raw_reader))) + if sys.exc_info()[1] is None: + with handler: + raw_reader = self._driver._read_arrow_batches(self._sql, self._parameters, self._chunk_size) + self._raw_reader = raw_reader + self._reader = iter(self._driver._normalize_reader(_to_pyarrow_reader(raw_reader))) + else: + + def _start_op() -> Any: + raw_reader = self._driver._read_arrow_batches(self._sql, self._parameters, self._chunk_size) + self._raw_reader = raw_reader + self._reader = iter(self._driver._normalize_reader(_to_pyarrow_reader(raw_reader))) + return self._reader + + self._driver._run_with_exception_handler(handler, _start_op) self._driver._check_pending_exception(handler) def fetch_chunk(self) -> "list[dict[str, Any]]": @@ -91,8 +98,11 @@ def fetch_chunk(self) -> "list[dict[str, Any]]": batch: Any = None handler = self._driver.handle_database_exceptions() try: - with handler: - batch = next(reader) + if sys.exc_info()[1] is None: + with handler: + batch = next(reader) + else: + batch = self._driver._run_with_exception_handler(handler, next, reader) except StopIteration: return [] self._driver._check_pending_exception(handler) @@ -300,9 +310,6 @@ def begin(self) -> None: mapped = create_mapped_exception(exc) msg = f"Failed to begin transaction: {mapped}" raise type(mapped)(msg) from exc - except Exception as exc: - msg = f"Failed to begin transaction: {exc}" - raise SQLSpecError(msg) from exc self._transaction_active = True def commit(self) -> None: @@ -390,23 +397,49 @@ def select_to_arrow( arrow_result: ArrowResult | None = None exc_handler = self.handle_database_exceptions() - with exc_handler, self.with_cursor(self.connection): - reader = self._read_arrow_batches( - sql, - _odbc_parameters(prepared_parameters, naive_utc_datetimes=self._dialect == "db2"), - resolved_batch_size, - ) - if return_format in {"reader", "batches"}: - arrow_reader = self._normalize_reader(_to_pyarrow_reader(reader)) - arrow_result = build_arrow_result_from_reader( - prepared_statement, - arrow_reader, - return_format=return_format, - batch_size=resolved_batch_size, - arrow_schema=arrow_schema, - ) + with self.with_cursor(self.connection): + if sys.exc_info()[1] is None: + with exc_handler: + reader = self._read_arrow_batches( + sql, + _odbc_parameters(prepared_parameters, naive_utc_datetimes=self._dialect == "db2"), + resolved_batch_size, + ) + if return_format in {"reader", "batches"}: + arrow_reader = self._normalize_reader(_to_pyarrow_reader(reader)) + arrow_result = build_arrow_result_from_reader( + prepared_statement, + arrow_reader, + return_format=return_format, + batch_size=resolved_batch_size, + arrow_schema=arrow_schema, + ) + else: + table = self._normalize_table(_reader_to_table(reader)) else: - table = self._normalize_table(_reader_to_table(reader)) + + def _select_op() -> tuple[Any, Any]: + r = self._read_arrow_batches( + sql, + _odbc_parameters(prepared_parameters, naive_utc_datetimes=self._dialect == "db2"), + resolved_batch_size, + ) + if return_format in {"reader", "batches"}: + ar = self._normalize_reader(_to_pyarrow_reader(r)) + res = build_arrow_result_from_reader( + prepared_statement, + ar, + return_format=return_format, + batch_size=resolved_batch_size, + arrow_schema=arrow_schema, + ) + return None, res + tbl = self._normalize_table(_reader_to_table(r)) + return tbl, None + + res = self._run_with_exception_handler(exc_handler, _select_op) + if res is not None: + table, arrow_result = res self._check_pending_exception(exc_handler) if arrow_result is not None: @@ -453,15 +486,26 @@ def bulk_insert_arrow(self, target_table: str, source: Any, *, chunk_size: int | resolved_chunk_size = chunk_size or self._chunk_size() exc_handler = self.handle_database_exceptions() inserted = False - with exc_handler, self.with_cursor(self.connection): - if isinstance(source, pa.Table) and hasattr(self.connection, "from_table_to_db"): - self.connection.from_table_to_db(source=source, target=target_table, chunk_size=resolved_chunk_size) - inserted = True - else: - reader = _table_to_reader(source, resolved_chunk_size) if isinstance(source, pa.Table) else source + with self.with_cursor(self.connection): + + def _insert_op() -> bool: + is_table = hasattr(source, "to_batches") and not hasattr(source, "read_next_batch") + if is_table and hasattr(self.connection, "from_table_to_db"): + self.connection.from_table_to_db(source=source, target=target_table, chunk_size=resolved_chunk_size) + return True + reader = _table_to_reader(source, resolved_chunk_size) if is_table else source if hasattr(self.connection, "insert_into_table"): self.connection.insert_into_table(reader=reader, table=target_table, chunk_size=resolved_chunk_size) - inserted = True + return True + return False + + if sys.exc_info()[1] is None: + with exc_handler: + inserted = _insert_op() + else: + ins = self._run_with_exception_handler(exc_handler, _insert_op) + if ins is not None: + inserted = ins self._check_pending_exception(exc_handler) if inserted: @@ -510,7 +554,7 @@ def load_from_arrow( if overwrite: target = _db2_table_reference(table) if self._dialect == "db2" else _quote_odbc_table(table) self.execute(f"DELETE FROM {target}") - if isinstance(raw_source, pa.RecordBatchReader) and hasattr(self.connection, "insert_into_table"): + if hasattr(raw_source, "read_next_batch") and hasattr(self.connection, "insert_into_table"): self.bulk_insert_arrow(table, raw_source) telemetry_payload: StorageTelemetry = { "rows_processed": 0, @@ -564,20 +608,53 @@ def _read_arrow_batches( *, fetch_concurrently: bool | None = None, ) -> Any: - kwargs: dict[str, Any] = { - "query": sql, - "batch_size": batch_size, - "parameters": parameters, - "max_bytes_per_batch": self._max_batch_bytes, - "max_text_size": self._max_text_size_val, - "max_binary_size": self._max_binary_size_val, - "fetch_concurrently": self._use_concurrent_fetch if fetch_concurrently is None else fetch_concurrently, - } - if self._query_timeout_sec_val is not None: - kwargs["query_timeout_sec"] = self._query_timeout_sec_val - if self._payload_text_encoding is not None: - kwargs["payload_text_encoding"] = self._payload_text_encoding - return self.connection.read_arrow_batches(**kwargs) + use_concurrent = self._use_concurrent_fetch if fetch_concurrently is None else fetch_concurrently + timeout = self._query_timeout_sec_val + encoding = self._payload_text_encoding + + if timeout is not None and encoding is not None: + return self.connection.read_arrow_batches( + query=sql, + batch_size=batch_size, + parameters=parameters, + max_bytes_per_batch=self._max_batch_bytes, + max_text_size=self._max_text_size_val, + max_binary_size=self._max_binary_size_val, + fetch_concurrently=use_concurrent, + query_timeout_sec=timeout, + payload_text_encoding=encoding, + ) + if timeout is not None: + return self.connection.read_arrow_batches( + query=sql, + batch_size=batch_size, + parameters=parameters, + max_bytes_per_batch=self._max_batch_bytes, + max_text_size=self._max_text_size_val, + max_binary_size=self._max_binary_size_val, + fetch_concurrently=use_concurrent, + query_timeout_sec=timeout, + ) + if encoding is not None: + return self.connection.read_arrow_batches( + query=sql, + batch_size=batch_size, + parameters=parameters, + max_bytes_per_batch=self._max_batch_bytes, + max_text_size=self._max_text_size_val, + max_binary_size=self._max_binary_size_val, + fetch_concurrently=use_concurrent, + payload_text_encoding=encoding, + ) + return self.connection.read_arrow_batches( + query=sql, + batch_size=batch_size, + parameters=parameters, + max_bytes_per_batch=self._max_batch_bytes, + max_text_size=self._max_text_size_val, + max_binary_size=self._max_binary_size_val, + fetch_concurrently=use_concurrent, + ) def _chunk_size(self) -> int: return self._chunk_size_val @@ -598,6 +675,9 @@ def _normalize_reader(self, reader: "ArrowRecordBatchReader") -> "ArrowRecordBat renamed = schema for index, name in enumerate(names): renamed = renamed.set(index, renamed.field(index).with_name(name)) + ensure_pyarrow() + import pyarrow as pa + return pa.RecordBatchReader.from_batches( renamed, (pa.RecordBatch.from_arrays(batch.columns, schema=renamed) for batch in reader) ) @@ -746,6 +826,7 @@ def _odbc_parameters(parameters: Any, *, naive_utc_datetimes: bool = False) -> " def _reader_to_table(reader: Any) -> Any: ensure_pyarrow() + import pyarrow as pa if isinstance(reader, pa.Table): return reader @@ -765,6 +846,7 @@ def _reader_to_table(reader: Any) -> Any: def _to_pyarrow_reader(reader: object) -> "ArrowRecordBatchReader": ensure_pyarrow() + import pyarrow as pa if isinstance(reader, pa.RecordBatchReader): return reader @@ -786,6 +868,7 @@ def _to_pyarrow_reader(reader: object) -> "ArrowRecordBatchReader": def _table_to_reader(table: Any, chunk_size: int) -> Any: ensure_pyarrow() + import pyarrow as pa return pa.RecordBatchReader.from_batches(table.schema, table.to_batches(max_chunksize=chunk_size)) diff --git a/sqlspec/adapters/arrow_odbc/events/store.py b/sqlspec/adapters/arrow_odbc/events/store.py index dfbd4597c..0a55b5be2 100644 --- a/sqlspec/adapters/arrow_odbc/events/store.py +++ b/sqlspec/adapters/arrow_odbc/events/store.py @@ -3,7 +3,8 @@ from typing import Final from sqlspec.adapters.arrow_odbc.config import ArrowOdbcConfig -from sqlspec.adapters.arrow_odbc.core import split_db2_name +from sqlspec.adapters.arrow_odbc.core import quote_bracket_identifier, split_db2_name, strip_bracket_quotes +from sqlspec.exceptions import ImproperConfigurationError from sqlspec.extensions.events import BaseEventQueueStore from sqlspec.utils.text import split_qualified_identifier @@ -33,7 +34,14 @@ class ArrowOdbcEventQueueStore(BaseEventQueueStore[ArrowOdbcConfig]): def __init__(self, config: ArrowOdbcConfig) -> None: super().__init__(config) - self._db2 = str(config.statement_config.dialect).lower() == "db2" + dialect = str(config.statement_config.dialect).lower() + if dialect == "db2": + self._db2 = True + elif dialect in {"mssql", "tsql"}: + self._db2 = False + else: + msg = f"ArrowOdbcEventQueueStore requires dialect 'mssql' or 'db2' (received {dialect!r})." + raise ImproperConfigurationError(msg) def _column_types(self) -> tuple[str, str, str]: return "NVARCHAR(MAX)", "NVARCHAR(MAX)", "DATETIME2(6)" @@ -57,9 +65,7 @@ def _table_ddl(self) -> str: def _index_name(self) -> str: if self._db2: return super()._index_name() - parts = [ - _strip_identifier_quotes(part) for part in split_qualified_identifier(self.table_name, quote_chars='"') - ] + parts = [strip_bracket_quotes(part) for part in split_qualified_identifier(self.table_name, quote_chars='"')] flat = "_".join(part for part in parts if part) return f"idx_{flat}_channel_status" @@ -87,17 +93,8 @@ def _wrap_drop_statement(self, statement: str) -> str: return f"IF OBJECT_ID(N'{_object_name(self.table_name)}', N'U') IS NOT NULL {statement};" -def _strip_identifier_quotes(part: str) -> str: - stripped = part.strip() - if len(stripped) >= _QUALIFIED_IDENTIFIER_MIN_PARTS and stripped.startswith("[") and stripped.endswith("]"): - return stripped[1:-1].replace("]]", "]") - if len(stripped) >= _QUALIFIED_IDENTIFIER_MIN_PARTS and stripped.startswith('"') and stripped.endswith('"'): - return stripped[1:-1].replace('""', '"') - return stripped - - def _split_table_name(table_name: str) -> tuple[str, str]: - parts = [_strip_identifier_quotes(part) for part in split_qualified_identifier(table_name, quote_chars='"')] + parts = [strip_bracket_quotes(part) for part in split_qualified_identifier(table_name, quote_chars='"')] if len(parts) < _QUALIFIED_IDENTIFIER_MIN_PARTS: return "dbo", parts[0] if parts else table_name schema_name = ".".join(parts[:-1]) @@ -106,8 +103,4 @@ def _split_table_name(table_name: str) -> tuple[str, str]: def _object_name(table_name: str) -> str: schema_name, bare_table_name = _split_table_name(table_name) - return f"{_quote_bracket_identifier(schema_name)}.{_quote_bracket_identifier(bare_table_name)}" - - -def _quote_bracket_identifier(identifier: str) -> str: - return f"[{identifier.replace(']', ']]')}]" + return f"{quote_bracket_identifier(schema_name)}.{quote_bracket_identifier(bare_table_name)}" diff --git a/sqlspec/adapters/arrow_odbc/litestar/store.py b/sqlspec/adapters/arrow_odbc/litestar/store.py index f1ecdb753..10d899d3a 100644 --- a/sqlspec/adapters/arrow_odbc/litestar/store.py +++ b/sqlspec/adapters/arrow_odbc/litestar/store.py @@ -1,7 +1,6 @@ """arrow-odbc Litestar Store implementation.""" import base64 -import contextlib from datetime import UTC, datetime, timedelta from typing import TYPE_CHECKING, Any @@ -11,13 +10,14 @@ db2_timestamp_text, split_db2_name, ) +from sqlspec.adapters.arrow_odbc.core import atomic_store_write as _atomic_write +from sqlspec.adapters.arrow_odbc.core import strip_bracket_quotes as _strip_tsql_quotes +from sqlspec.exceptions import ImproperConfigurationError from sqlspec.extensions.litestar.store import BaseSQLSpecStore from sqlspec.utils.sync_tools import async_ from sqlspec.utils.text import split_qualified_identifier if TYPE_CHECKING: - from collections.abc import Iterator - from sqlspec.adapters.arrow_odbc.config import ArrowOdbcConfig __all__ = ("ArrowOdbcStore",) @@ -37,10 +37,15 @@ class ArrowOdbcStore(BaseSQLSpecStore["ArrowOdbcConfig"]): def __init__(self, config: "ArrowOdbcConfig") -> None: super().__init__(config) - dialect = str(config.statement_config.dialect).lower() - self._sql: _TsqlSessionSql | _Db2SessionSql = ( - _Db2SessionSql(self._table_name) if dialect == "db2" else _TsqlSessionSql(self._table_name) - ) + dialect_attr = getattr(config.statement_config, "dialect", None) + dialect = str(dialect_attr).lower() if isinstance(dialect_attr, str) else "mssql" + if dialect == "db2": + self._sql: _TsqlSessionSql | _Db2SessionSql = _Db2SessionSql(self._table_name) + elif dialect in {"mssql", "tsql"}: + self._sql = _TsqlSessionSql(self._table_name) + else: + msg = f"ArrowOdbcStore requires dialect 'mssql' or 'db2' (received {dialect!r})." + raise ImproperConfigurationError(msg) @property def _chunk_table_name(self) -> str: @@ -430,32 +435,6 @@ def _chunk_table_ddl(self) -> str: ) -@contextlib.contextmanager -def _atomic_write(driver: Any) -> "Iterator[None]": - if driver._connection_in_transaction(): # pyright: ignore[reportPrivateUsage] - yield - return - if not getattr(driver, "_connection_autocommit", True): - driver.begin() - try: - yield - except Exception: - with contextlib.suppress(Exception): - driver.rollback() - raise - else: - driver.commit() - - -def _strip_tsql_quotes(part: str) -> str: - stripped = part.strip() - if len(stripped) >= _QUALIFIED_PARTS_MIN and stripped.startswith("[") and stripped.endswith("]"): - return stripped[1:-1].replace("]]", "]") - if len(stripped) >= _QUALIFIED_PARTS_MIN and stripped.startswith('"') and stripped.endswith('"'): - return stripped[1:-1].replace('""', '"') - return stripped - - def _split_tsql_table_name(table_name: str) -> "tuple[str, str]": parts = [_strip_tsql_quotes(part) for part in split_qualified_identifier(table_name, quote_chars='"')] if len(parts) < _QUALIFIED_PARTS_MIN: diff --git a/sqlspec/adapters/arrow_odbc/type_converter.py b/sqlspec/adapters/arrow_odbc/type_converter.py index b239327da..b648c21de 100644 --- a/sqlspec/adapters/arrow_odbc/type_converter.py +++ b/sqlspec/adapters/arrow_odbc/type_converter.py @@ -6,9 +6,7 @@ from sqlspec.utils.module_loader import ensure_pyarrow if TYPE_CHECKING: - import pyarrow as pa -else: - pa = import_optional("pyarrow") + from pyarrow import DataType __all__ = ("odbc_type_to_arrow",) @@ -66,7 +64,7 @@ _ARROW_TYPE_CACHE: "dict[tuple[str, tuple[Any, ...], tuple[tuple[str, Any], ...]], Any]" = {} -def odbc_type_to_arrow(sql_type: str, *, precision: int | None = None, scale: int | None = None) -> "pa.DataType": +def odbc_type_to_arrow(sql_type: str, *, precision: int | None = None, scale: int | None = None) -> "DataType": """Resolve an ODBC SQL type name to an Arrow data type.""" normalized_type = sql_type.lower().split("(", 1)[0].strip() if normalized_type in {"decimal", "numeric"} and precision is not None and scale is not None: @@ -82,14 +80,16 @@ def odbc_type_to_arrow(sql_type: str, *, precision: int | None = None, scale: in return _arrow_type(name, args, kwargs) -def _arrow_type(name: str, args: tuple[Any, ...] = (), kwargs: dict[str, Any] | None = None) -> "pa.DataType": +def _arrow_type(name: str, args: tuple[Any, ...] = (), kwargs: dict[str, Any] | None = None) -> "DataType": kwargs_key = tuple(sorted((kwargs or {}).items())) cache_key = (name, args, kwargs_key) cached = _ARROW_TYPE_CACHE.get(cache_key) if cached is not None: - return cast("pa.DataType", cached) + return cast("DataType", cached) ensure_pyarrow() + pa = import_optional("pyarrow") + assert pa is not None resolved = getattr(pa, name)(*args, **(kwargs or {})) _ARROW_TYPE_CACHE[cache_key] = resolved - return cast("pa.DataType", resolved) + return cast("DataType", resolved) diff --git a/sqlspec/adapters/asyncmy/_typing.py b/sqlspec/adapters/asyncmy/_typing.py index 0ebf568b0..1c4f81617 100644 --- a/sqlspec/adapters/asyncmy/_typing.py +++ b/sqlspec/adapters/asyncmy/_typing.py @@ -25,13 +25,9 @@ from sqlspec.exceptions import SQLSpecError if TYPE_CHECKING: - from collections.abc import Awaitable, Callable, Iterator - from types import TracebackType + from collections.abc import Iterator from typing import Protocol, TypeAlias - from sqlspec.adapters.asyncmy.driver import AsyncmyDriver - from sqlspec.core import StatementConfig - class AsyncmyMySQLError(Exception): errno: int @@ -68,8 +64,7 @@ class AsyncmyFieldTypeProtocol(Protocol): AsyncmyModule: TypeAlias = AsyncmyModuleProtocol AsyncmyPool: TypeAlias = _AsyncmyPool AsyncmyRawCursor: TypeAlias = _AsyncmyCursor - -if not TYPE_CHECKING: +else: AsyncmyConnection = Connection AsyncmyDictCursor = _AsyncmyDictCursor AsyncmyError = _asyncmy_errors.Error @@ -93,12 +88,9 @@ class AsyncmyFieldTypeProtocol(Protocol): "AsyncmyProgrammingError", "AsyncmyRawCursor", "AsyncmySSCursor", - "AsyncmySessionContext", "asyncmy_local_infile", ) -_resolved_driver_type: "type[AsyncmyDriver] | None" = None - class AsyncmyCursor: """Context manager for AsyncMy cursor operations. @@ -133,60 +125,6 @@ async def __aexit__(self, *_: object) -> None: await raw_cursor.__aexit__(None, None, None) -class AsyncmySessionContext: - """Async context manager for AsyncMy sessions. - - This class is intentionally excluded from mypyc compilation to avoid ABI - boundary issues. It receives callables from uncompiled config classes and - instantiates compiled Driver objects, acting as a bridge between compiled - and uncompiled code. - - Uses callable-based connection management to decouple from config implementation. - """ - - __slots__ = ( - "_acquire_connection", - "_connection", - "_driver", - "_driver_features", - "_prepare_driver", - "_release_connection", - "_statement_config", - ) - - def __init__( - self, - acquire_connection: "Callable[[], Awaitable[AsyncmyConnection]]", - release_connection: "Callable[..., Any]", - statement_config: "StatementConfig", - driver_features: "dict[str, Any]", - prepare_driver: "Callable[[AsyncmyDriver], AsyncmyDriver]", - ) -> None: - self._acquire_connection = acquire_connection - self._release_connection = release_connection - self._statement_config = statement_config - self._driver_features = driver_features - self._prepare_driver = prepare_driver - self._connection: AsyncmyConnection | None = None - self._driver: AsyncmyDriver | None = None - - async def __aenter__(self) -> "AsyncmyDriver": - driver_cls = _get_asyncmy_driver_type() - self._connection = await self._acquire_connection() - self._driver = driver_cls( - connection=self._connection, statement_config=self._statement_config, driver_features=self._driver_features - ) - return self._prepare_driver(self._driver) - - async def __aexit__( - self, exc_type: "type[BaseException] | None", exc_val: "BaseException | None", exc_tb: "TracebackType | None" - ) -> "bool | None": - if self._connection is not None: - await self._release_connection(self._connection, exc_type=exc_type, exc_val=exc_val, exc_tb=exc_tb) - self._connection = None - return None - - @contextlib.contextmanager def asyncmy_local_infile(connection: "AsyncmyConnection", filename: str) -> "Iterator[None]": """Scope the asyncmy 0.2.13/0.2.14 filename handoff fix to one native load. @@ -256,12 +194,3 @@ async def _read_load_local_packet(self, first_packet: Any) -> None: raise SQLSpecError(msg) raw_self: Any = self raw_self._read_ok_packet(packet) - - -def _get_asyncmy_driver_type() -> "type[AsyncmyDriver]": - global _resolved_driver_type - if _resolved_driver_type is None: - from sqlspec.adapters.asyncmy.driver import AsyncmyDriver as _Driver - - _resolved_driver_type = _Driver - return _resolved_driver_type diff --git a/sqlspec/adapters/asyncmy/adk/store.py b/sqlspec/adapters/asyncmy/adk/store.py index 3b335dc4b..c2ed507cd 100644 --- a/sqlspec/adapters/asyncmy/adk/store.py +++ b/sqlspec/adapters/asyncmy/adk/store.py @@ -5,7 +5,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, NotRequired, cast from sqlspec.adapters.asyncmy._typing import AsyncmyCursor, AsyncmyProgrammingError, AsyncmyRawCursor -from sqlspec.adapters.asyncmy.core import resolve_rowcount +from sqlspec.adapters.asyncmy.core import resolve_many_rowcount, resolve_rowcount from sqlspec.config import ADKConfig from sqlspec.extensions.adk import BaseAsyncADKStore, StoredEvent, StoredSession, normalize_session_list_options from sqlspec.extensions.adk.memory.store import BaseAsyncADKMemoryStore @@ -565,46 +565,52 @@ async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: " ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) """ - async with self._config.provide_connection() as conn: - async with AsyncmyCursor(conn, cursor_class=AsyncmyRawCursor) as cursor: - for entry in entries: - content_json = self._json_serializer(entry["content_json"]) - meta = entry.get("metadata_json") - metadata_json = self._json_serializer(meta) if meta is not None else None - params: tuple[Any, ...] - if self._owner_id_column_name: - params = ( - entry["id"], - entry["session_id"], - entry["app_name"], - entry["user_id"], - entry.get("scope", "user"), - entry["event_id"], - entry.get("author"), - owner_id, - entry["timestamp"], - content_json, - entry["content_text"], - metadata_json, - entry["inserted_at"], - ) - else: - params = ( - entry["id"], - entry["session_id"], - entry["app_name"], - entry["user_id"], - entry.get("scope", "user"), - entry["event_id"], - entry.get("author"), - entry["timestamp"], - content_json, - entry["content_text"], - metadata_json, - entry["inserted_at"], - ) - await cursor.execute(sql, params) - inserted_count += resolve_rowcount(cursor) + params_list: list[tuple[Any, ...]] = [] + for entry in entries: + content_json = self._json_serializer(entry["content_json"]) + meta = entry.get("metadata_json") + metadata_json = self._json_serializer(meta) if meta is not None else None + if self._owner_id_column_name: + params_list.append(( + entry["id"], + entry["session_id"], + entry["app_name"], + entry["user_id"], + entry.get("scope", "user"), + entry["event_id"], + entry.get("author"), + owner_id, + entry["timestamp"], + content_json, + entry["content_text"], + metadata_json, + entry["inserted_at"], + )) + else: + params_list.append(( + entry["id"], + entry["session_id"], + entry["app_name"], + entry["user_id"], + entry.get("scope", "user"), + entry["event_id"], + entry.get("author"), + entry["timestamp"], + content_json, + entry["content_text"], + metadata_json, + entry["inserted_at"], + )) + + if not params_list: + return 0 + + async with ( + self._config.provide_connection() as conn, + AsyncmyCursor(conn, cursor_class=AsyncmyRawCursor) as cursor, + ): + await cursor.executemany(sql, params_list) + inserted_count = resolve_many_rowcount(cursor, params_list, fallback_count=len(params_list)) await conn.commit() return inserted_count @@ -637,13 +643,14 @@ async def search_entries( """ params = (*scope_params, query, limit_value) else: + escaped_query = query.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_") sql = f""" SELECT * FROM {self._memory_table} - WHERE {where_scope} AND content_text LIKE %s + WHERE {where_scope} AND content_text LIKE %s ESCAPE '\\\\' ORDER BY timestamp DESC LIMIT %s """ - params = (*scope_params, f"%{query}%", limit_value) + params = (*scope_params, f"%{escaped_query}%", limit_value) try: async with ( @@ -750,7 +757,7 @@ async def _memory_table_ddl(self) -> str: author VARCHAR(256){owner_id_line}, timestamp TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6), content_json JSON NOT NULL, - content_text TEXT NOT NULL, + content_text LONGTEXT NOT NULL, metadata_json JSON, inserted_at TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6), INDEX idx_{self._memory_table}_app_scope_user_time (app_name, scope, user_id, timestamp), diff --git a/sqlspec/adapters/asyncmy/config.py b/sqlspec/adapters/asyncmy/config.py index 69dfd4052..371fd8886 100644 --- a/sqlspec/adapters/asyncmy/config.py +++ b/sqlspec/adapters/asyncmy/config.py @@ -14,7 +14,6 @@ AsyncmyModule, AsyncmyPool, AsyncmyRawCursor, - AsyncmySessionContext, ) from sqlspec.adapters.asyncmy.core import apply_driver_features, default_statement_config from sqlspec.adapters.asyncmy.driver import AsyncmyDriver, AsyncmyExceptionHandler @@ -24,6 +23,7 @@ from sqlspec.exceptions import ImproperConfigurationError from sqlspec.extensions.events import EventRuntimeHints from sqlspec.utils.config_tools import normalize_connection_config, parse_mysql_dsn +from sqlspec.utils.module_loader import ensure_pyarrow if TYPE_CHECKING: import ssl @@ -32,6 +32,7 @@ from sqlspec.core import StatementConfig from sqlspec.observability import ObservabilityConfig + from sqlspec.storage import StorageCapabilities __all__ = ( @@ -40,6 +41,7 @@ "AsyncmyDriverFeatures", "AsyncmyPoolParams", "AsyncmySSLParams", + "AsyncmySessionContext", "build_connection_config", ) @@ -222,9 +224,13 @@ async def acquire_connection(self) -> "AsyncmyConnection": async def release_connection(self, _conn: "AsyncmyConnection", **kwargs: Any) -> None: ctx = self._contexts.pop(id(_conn), None) if ctx is not None: - if _conn is not None and bool(getattr(_conn, "server_status", 0) & 1): - with contextlib.suppress(Exception): - await _conn.rollback() + if _conn is not None: + in_transaction = False + with contextlib.suppress(AttributeError): + in_transaction = bool(_conn.server_status & 1) + if in_transaction: + with contextlib.suppress(Exception): + await _conn.rollback() await ctx.__aexit__(kwargs.get("exc_type"), kwargs.get("exc_val"), kwargs.get("exc_tb")) @@ -259,9 +265,13 @@ async def __aexit__( ) -> bool | None: conn = self._connection self._connection = None - if conn is not None and bool(getattr(conn, "server_status", 0) & 1): - with contextlib.suppress(Exception): - await conn.rollback() + if conn is not None: + in_transaction = False + with contextlib.suppress(AttributeError): + in_transaction = bool(conn.server_status & 1) + if in_transaction: + with contextlib.suppress(Exception): + await conn.rollback() if self._ctx: ctx = self._ctx self._ctx = None @@ -269,6 +279,55 @@ async def __aexit__( return None +class AsyncmySessionContext: + """Async context manager for AsyncMy sessions. + + Acquires a connection from the config, builds the driver, and releases the + connection when the session ends. + """ + + __slots__ = ( + "_acquire_connection", + "_connection", + "_driver", + "_driver_features", + "_prepare_driver", + "_release_connection", + "_statement_config", + ) + + def __init__( + self, + acquire_connection: "Callable[[], Awaitable[AsyncmyConnection]]", + release_connection: "Callable[..., Any]", + statement_config: "StatementConfig", + driver_features: "dict[str, Any]", + prepare_driver: "Callable[[AsyncmyDriver], AsyncmyDriver]", + ) -> None: + self._acquire_connection = acquire_connection + self._release_connection = release_connection + self._statement_config = statement_config + self._driver_features = driver_features + self._prepare_driver = prepare_driver + self._connection: AsyncmyConnection | None = None + self._driver: AsyncmyDriver | None = None + + async def __aenter__(self) -> "AsyncmyDriver": + self._connection = await self._acquire_connection() + self._driver = AsyncmyDriver( + connection=self._connection, statement_config=self._statement_config, driver_features=self._driver_features + ) + return self._prepare_driver(self._driver) + + async def __aexit__( + self, exc_type: "type[BaseException] | None", exc_val: "BaseException | None", exc_tb: "TracebackType | None" + ) -> "bool | None": + if self._connection is not None: + await self._release_connection(self._connection, exc_type=exc_type, exc_val=exc_val, exc_tb=exc_tb) + self._connection = None + return None + + @mypyc_attr(native_class=False) class AsyncmyConfig(AsyncDatabaseConfig[AsyncmyConnection, "AsyncmyPool", AsyncmyDriver]): """Configuration for Asyncmy database connections.""" @@ -278,10 +337,10 @@ class AsyncmyConfig(AsyncDatabaseConfig[AsyncmyConnection, "AsyncmyPool", Asyncm driver_type: ClassVar[type[AsyncmyDriver]] = AsyncmyDriver connection_type: "ClassVar[type[Any]]" = cast("type[Any]", AsyncmyConnection) supports_transactional_ddl: ClassVar[bool] = False - supports_native_arrow_export: ClassVar[bool] = True - supports_native_parquet_export: ClassVar[bool] = True - supports_native_arrow_import: ClassVar[bool] = True - supports_native_parquet_import: ClassVar[bool] = True + supports_native_arrow_export: ClassVar[bool] = False + supports_native_parquet_export: ClassVar[bool] = False + supports_native_arrow_import: ClassVar[bool] = False + supports_native_parquet_import: ClassVar[bool] = False supports_native_row_streaming: ClassVar[bool] = True type_coercion_capabilities: ClassVar[TypeCoercionCapabilities] = TypeCoercionCapabilities( datetime_binding="native", timestamp_precision="microsecond", json_columns_decoded=False, uuid_binding="text" @@ -320,8 +379,7 @@ def __init__( connection_config = build_connection_config(connection_config) statement_config = statement_config or default_statement_config - statement_config, driver_features = apply_driver_features(statement_config, driver_features) - features_dict = dict(driver_features) if driver_features else {} + statement_config, features_dict = apply_driver_features(statement_config, driver_features) self._user_connection_hook: Callable[[AsyncmyConnection], Awaitable[None]] | None = features_dict.pop( "on_connection_create", None ) @@ -356,7 +414,11 @@ async def create_connection(self) -> AsyncmyConnection: pool_kwargs, connection_kwargs = _split_pool_config(self.connection_config) connection_kwargs.update({key: value for key, value in pool_kwargs.items() if key in _ASYNCMY_SHARED_POOL_KEYS}) connection = await asyncmy.connect(**connection_kwargs) - await self._ensure_connection(connection) + try: + await self._ensure_connection(connection) + except BaseException: + connection.close() + raise return connection async def provide_pool(self, *args: Any, **kwargs: Any) -> "AsyncmyPool": @@ -398,6 +460,15 @@ def get_event_runtime_hints(self) -> "EventRuntimeHints": return EventRuntimeHints(poll_interval=0.25, lease_seconds=5, select_for_update=True, skip_locked=True) + def _build_storage_capabilities(self) -> "StorageCapabilities": + capabilities = super()._build_storage_capabilities() + arrow_ready = self._dependency_available(ensure_pyarrow) + capabilities["arrow_import_enabled"] = arrow_ready + capabilities["arrow_export_enabled"] = arrow_ready + capabilities["parquet_import_enabled"] = arrow_ready + capabilities["parquet_export_enabled"] = arrow_ready + return capabilities + async def _create_pool(self) -> "AsyncmyPool": """Create the actual async connection pool. diff --git a/sqlspec/adapters/asyncmy/core.py b/sqlspec/adapters/asyncmy/core.py index 14bbed72e..602da097c 100644 --- a/sqlspec/adapters/asyncmy/core.py +++ b/sqlspec/adapters/asyncmy/core.py @@ -17,6 +17,7 @@ ForeignKeyViolationError, IntegrityError, NotNullViolationError, + OperationalError, PermissionDeniedError, QueryTimeoutError, SQLParsingError, @@ -28,7 +29,7 @@ from sqlspec.utils.serializers import from_json, to_json from sqlspec.utils.text import quote_backtick_identifier, split_qualified_identifier from sqlspec.utils.type_converters import build_uuid_coercions -from sqlspec.utils.type_guards import has_cursor_metadata, has_lastrowid, has_rowcount, has_sqlstate, has_type_code +from sqlspec.utils.type_guards import has_lastrowid, has_rowcount, has_sqlstate, has_type_code if TYPE_CHECKING: from sqlspec.core.parameters import ParameterValidator @@ -45,8 +46,6 @@ "collect_stream_rows", "create_mapped_exception", "default_statement_config", - "detect_json_columns", - "detect_json_columns_from_description", "driver_profile", "encode_records_for_local_infile", "escape_literal_percent", @@ -54,7 +53,6 @@ "normalize_execute_many_parameters", "normalize_execute_parameters", "normalize_lastrowid", - "resolve_column_names", "resolve_many_rowcount", "resolve_row_plan", "resolve_rowcount", @@ -80,13 +78,11 @@ MYSQL_SYNTAX_ERROR_MIN: Final[int] = 1064 MYSQL_SYNTAX_ERROR_MAX_EXCLUSIVE: Final[int] = 1100 -_MYSQL_MIGRATION_ERROR_CODES: Final[frozenset[int]] = frozenset((1061, 1091)) -_MYSQL_SQLSTATE_EXACT_DISPATCH: Final[dict[str, tuple[type[SQLSpecError], str]]] = { - "23505": (UniqueViolationError, "unique constraint violation"), - "23503": (ForeignKeyViolationError, "foreign key constraint violation"), - "23502": (NotNullViolationError, "not-null constraint violation"), - "23514": (CheckViolationError, "check constraint violation"), +_MYSQL_DDL_ERROR_DISPATCH: dict[int, tuple[type[SQLSpecError], str]] = { + 1061: (OperationalError, "duplicate key name"), + 1091: (OperationalError, "cannot drop column or key"), } +_MYSQL_SQLSTATE_EXACT_DISPATCH: Final[dict[str, tuple[type[SQLSpecError], str]]] = {} _MYSQL_SQLSTATE_PREFIX_DISPATCH: Final[dict[str, tuple[type[SQLSpecError], str]]] = { "23": (IntegrityError, "integrity constraint violation"), "28": (PermissionDeniedError, "authorization error"), @@ -95,6 +91,10 @@ "08": (DatabaseConnectionError, "connection error"), "22": (DataError, "data error"), } +_MYSQL_SYNTAX_ERROR_DISPATCH: Final[dict[int, tuple[type[SQLSpecError], str]]] = { + 1054: (SQLParsingError, "unknown column"), + 1146: (SQLParsingError, "table does not exist"), +} _MYSQL_CONSTRAINT_ERROR_DISPATCH: Final[dict[int, tuple[type[SQLSpecError], str]]] = { MYSQL_ER_DUP_ENTRY: (UniqueViolationError, "unique constraint violation"), 1169: (UniqueViolationError, "unique constraint violation"), @@ -112,6 +112,8 @@ MYSQL_ER_DBACCESS_DENIED: (PermissionDeniedError, "access denied"), MYSQL_ER_ACCESS_DENIED: (PermissionDeniedError, "access denied"), MYSQL_ER_TABLEACCESS_DENIED: (PermissionDeniedError, "access denied"), + 1290: (PermissionDeniedError, "read-only server"), + 1836: (PermissionDeniedError, "read-only mode"), } _MYSQL_TRANSACTION_ERROR_DISPATCH: Final[dict[int, tuple[type[SQLSpecError], str]]] = { MYSQL_ER_LOCK_DEADLOCK: (DeadlockError, "deadlock detected"), @@ -121,6 +123,7 @@ 3024: (QueryTimeoutError, "query execution timeout"), } _MYSQL_CONNECTION_ERROR_DISPATCH: Final[dict[int, tuple[type[SQLSpecError], str]]] = { + 1040: (DatabaseConnectionError, "too many connections"), MYSQL_CR_SERVER_LOST: (ConnectionTimeoutError, "connection lost"), MYSQL_CR_CONNECTION_ERROR: (DatabaseConnectionError, "connection error"), MYSQL_CR_CONN_HOST_ERROR: (DatabaseConnectionError, "connection error"), @@ -191,8 +194,11 @@ async def close(self, error: bool = False) -> None: cursor = self._cursor self._cursor = None if cursor is not None: - with contextlib.suppress(Exception): + try: await cursor.close() + except BaseException: + with contextlib.suppress(Exception): + self._driver.connection.close() async def _start(self) -> None: try: @@ -207,7 +213,7 @@ async def _start(self) -> None: def build_profile() -> "DriverParameterProfile": """Create the AsyncMy driver parameter profile.""" - coercions: dict[type, Callable[[Any], Any]] = {bool: bool_to_int, **build_uuid_coercions()} + coercions: dict[type, Callable[[Any], Any]] = {bool: int, **build_uuid_coercions()} return DriverParameterProfile( name="asyncmy", default_style=ParameterStyle.QMARK, @@ -345,13 +351,6 @@ def normalize_execute_many_parameters(parameters: Any) -> Any: return parameters -def resolve_column_names(description: "Sequence[Any] | None") -> "list[str]": - """Extract column names from cursor description.""" - if not description: - return [] - return [col[0] for col in description] - - def resolve_row_plan( description: "Sequence[Any] | None", json_type_codes: "set[int]" ) -> "tuple[list[str], list[int] | None]": @@ -377,24 +376,6 @@ def resolve_row_plan( return column_names, json_indexes or None -def detect_json_columns_from_description( - description: "Sequence[Any] | None", json_type_codes: "set[int]" -) -> "list[int]": - """Identify JSON column indexes from pre-fetched cursor description metadata.""" - return resolve_row_plan(description, json_type_codes)[1] or [] - - -def detect_json_columns( - cursor: Any, json_type_codes: "set[int]", description: "Sequence[Any] | None" = None -) -> "list[int]": - """Identify JSON column indexes from cursor metadata.""" - if description is None: - if not has_cursor_metadata(cursor): - return [] - description = cursor.description - return detect_json_columns_from_description(description, json_type_codes) - - def collect_stream_rows( rows: "Sequence[Any]", row_plan: "tuple[list[str], list[int] | None]", @@ -476,13 +457,9 @@ def normalize_lastrowid(cursor: Any) -> int | None: return last_id if isinstance(last_id, int) else None -def bool_to_int(value: bool) -> int: - """Coerce boolean values to integer for MySQL parameter compatibility.""" - return int(value) - - -def create_mapped_exception(error: Any, *, logger: Any | None = None) -> "SQLSpecError | bool": +def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpecError: """Map MySQL driver exceptions to SQLSpec errors.""" + _ = logger error_code = getattr(error, "errno", None) if error_code is None and isinstance(error, Exception) and error.args: value = error.args[0] @@ -491,10 +468,9 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> "SQLSpe sqlstate = error.sqlstate if has_sqlstate(error) else None sqlstate_prefix = sqlstate[:2] if isinstance(sqlstate, str) and sqlstate else None - if error_code in _MYSQL_MIGRATION_ERROR_CODES: - if logger is not None: - logger.warning("MySQL expected migration error (ignoring): %s", error) - return True + dispatch = _MYSQL_DDL_ERROR_DISPATCH.get(error_code) if error_code is not None else None + if dispatch is not None: + return _create_mysql_error(error, sqlstate, error_code, dispatch[0], dispatch[1]) dispatch = _MYSQL_SQLSTATE_EXACT_DISPATCH.get(sqlstate) if sqlstate is not None else None if dispatch is not None: @@ -522,6 +498,10 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> "SQLSpe dispatch = _MYSQL_SQLSTATE_PREFIX_DISPATCH["40"] return _create_mysql_error(error, sqlstate, error_code, dispatch[0], dispatch[1]) + dispatch = _MYSQL_SYNTAX_ERROR_DISPATCH.get(error_code) if error_code is not None else None + if dispatch is not None: + return _create_mysql_error(error, sqlstate, error_code, dispatch[0], dispatch[1]) + if sqlstate_prefix == "42": dispatch = _MYSQL_SQLSTATE_PREFIX_DISPATCH["42"] return _create_mysql_error(error, sqlstate, error_code, dispatch[0], dispatch[1]) diff --git a/sqlspec/adapters/asyncmy/driver.py b/sqlspec/adapters/asyncmy/driver.py index ca983e94f..3d88f2b78 100644 --- a/sqlspec/adapters/asyncmy/driver.py +++ b/sqlspec/adapters/asyncmy/driver.py @@ -4,6 +4,8 @@ type coercion, error handling, and transaction management. """ +import asyncio +import shutil import tempfile from collections.abc import Sized from typing import TYPE_CHECKING, Any, Final, Literal, cast @@ -14,7 +16,6 @@ AsyncmyError, AsyncmyFieldType, AsyncmyMySQLError, - AsyncmySessionContext, asyncmy_local_infile, ) from sqlspec.adapters.asyncmy.core import ( @@ -39,7 +40,6 @@ from sqlspec.adapters.asyncmy.data_dictionary import AsyncmyDataDictionary from sqlspec.core import ArrowResult, get_cache_config, register_driver_profile from sqlspec.driver import AsyncDriverAdapterBase, AsyncRowStream, BaseAsyncExceptionHandler -from sqlspec.exceptions import SQLSpecError from sqlspec.utils.logging import get_logger from sqlspec.utils.serializers import from_json from sqlspec.utils.type_guards import supports_json_type @@ -52,7 +52,7 @@ from sqlspec.driver import CachedQuery, ExecutionResult from sqlspec.storage import StorageBridgeJob, StorageDestination, StorageFormat, StorageTelemetry -__all__ = ("AsyncmyCursor", "AsyncmyDriver", "AsyncmyExceptionHandler", "AsyncmySessionContext") +__all__ = ("AsyncmyCursor", "AsyncmyDriver", "AsyncmyExceptionHandler") logger = get_logger("sqlspec.adapters.asyncmy") @@ -93,10 +93,7 @@ def _handle_exception(self, exc_type: "type[BaseException] | None", exc_val: "Ba if exc_type is None: return False if issubclass(exc_type, cast("type[BaseException]", AsyncmyError)): - result = create_mapped_exception(exc_val, logger=logger) - if result is True: - return True - self.pending_exception = cast("Exception", result) + self.pending_exception = create_mapped_exception(exc_val) return True return False @@ -154,8 +151,8 @@ async def dispatch_execute(self, cursor: Any, statement: "SQL") -> "ExecutionRes Returns: ExecutionResult: Statement execution results with data or row counts """ - sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) - sql = escape_literal_percent(sql, prepared_parameters, self.statement_config.parameter_validator) + sql, prepared_parameters = self._compiled_sql(statement, statement.statement_config) + sql = escape_literal_percent(sql, prepared_parameters, statement.statement_config.parameter_validator) await cursor.execute(sql, normalize_execute_parameters(prepared_parameters)) description = getattr(cursor, "description", None) or None @@ -194,15 +191,17 @@ async def dispatch_execute_many(self, cursor: Any, statement: "SQL") -> "Executi Returns: ExecutionResult: Batch execution results """ - sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) + sql, prepared_parameters = self._compiled_sql(statement, statement.statement_config) match = ASYNCMY_INSERT_VALUES_PATTERN.match(sql) if match: sql = ( - escape_literal_percent(match.group(1), prepared_parameters, self.statement_config.parameter_validator) + escape_literal_percent( + match.group(1), prepared_parameters, statement.statement_config.parameter_validator + ) + sql[match.end(1) :] ) else: - sql = escape_literal_percent(sql, prepared_parameters, self.statement_config.parameter_validator) + sql = escape_literal_percent(sql, prepared_parameters, statement.statement_config.parameter_validator) prepared_parameters = normalize_execute_many_parameters(prepared_parameters) parameter_count = len(prepared_parameters) if isinstance(prepared_parameters, Sized) else None @@ -215,9 +214,6 @@ async def dispatch_execute_many(self, cursor: Any, statement: "SQL") -> "Executi async def dispatch_execute_script(self, cursor: Any, statement: "SQL") -> "ExecutionResult": """Execute SQL script with statement splitting and parameter handling. - Splits multi-statement scripts and executes each statement sequentially. - Parameters are embedded as static values for script execution compatibility. - Args: cursor: AsyncMy cursor object statement: SQL script to execute @@ -225,25 +221,7 @@ async def dispatch_execute_script(self, cursor: Any, statement: "SQL") -> "Execu Returns: ExecutionResult: Script execution results with statement count """ - sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) - statements = self.split_script_statements(sql, statement.statement_config, strip_trailing_semicolon=True) - - successful_count = 0 - last_cursor = cursor - normalized_params = normalize_execute_parameters(prepared_parameters) - - for stmt in statements: - stmt_sql = ( - escape_literal_percent(stmt, prepared_parameters, self.statement_config.parameter_validator) - if normalized_params is not None - else stmt - ) - await cursor.execute(stmt_sql, normalized_params) - successful_count += 1 - - return self.create_execution_result( - last_cursor, statement_count=len(statements), successful_statements=successful_count, is_script_result=True - ) + return await super().dispatch_execute_script(cursor, statement) async def begin(self) -> None: """Begin a database transaction. @@ -256,11 +234,7 @@ async def begin(self) -> None: try: await self.connection.begin() except AsyncmyMySQLError as e: - mapped = create_mapped_exception(e, logger=logger) - if isinstance(mapped, Exception): - raise mapped from e - msg = f"Failed to begin MySQL transaction: {e}" - raise SQLSpecError(msg) from e + raise create_mapped_exception(e) from e async def commit(self) -> None: """Commit the current transaction. @@ -271,11 +245,7 @@ async def commit(self) -> None: try: await self.connection.commit() except AsyncmyMySQLError as e: - mapped = create_mapped_exception(e, logger=logger) - if isinstance(mapped, Exception): - raise mapped from e - msg = f"Failed to commit MySQL transaction: {e}" - raise SQLSpecError(msg) from e + raise create_mapped_exception(e) from e async def rollback(self) -> None: """Rollback the current transaction. @@ -286,11 +256,7 @@ async def rollback(self) -> None: try: await self.connection.rollback() except AsyncmyMySQLError as e: - mapped = create_mapped_exception(e, logger=logger) - if isinstance(mapped, Exception): - raise mapped from e - msg = f"Failed to rollback MySQL transaction: {e}" - raise SQLSpecError(msg) from e + raise create_mapped_exception(e) from e def with_cursor(self, connection: "AsyncmyConnection") -> "AsyncmyCursor": """Create cursor context manager for the connection. @@ -307,8 +273,8 @@ def dispatch_select_stream(self, statement: "SQL", chunk_size: int) -> "AsyncRow """Return a native asyncmy row stream backed by an unbuffered ``SSCursor``.""" if not statement.returns_rows(): return None - sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) - sql = escape_literal_percent(sql, prepared_parameters, self.statement_config.parameter_validator) + sql, prepared_parameters = self._compiled_sql(statement, statement.statement_config) + sql = escape_literal_percent(sql, prepared_parameters, statement.statement_config.parameter_validator) return AsyncRowStream(AsyncmyStreamSource(self, sql, prepared_parameters, chunk_size, ASYNCMY_JSON_TYPE_CODES)) def handle_database_exceptions(self) -> "AsyncmyExceptionHandler": @@ -414,22 +380,6 @@ async def load_from_arrow( self._attach_partition_telemetry(telemetry_payload, partitioner) return self._storage_job(telemetry_payload, telemetry) - async def load_from_storage( - self, - table: str, - source: "StorageDestination", - *, - file_format: "StorageFormat", - partitioner: "dict[str, object] | None" = None, - overwrite: bool = False, - **kwargs: Any, - ) -> "StorageBridgeJob": - """Load staged artifacts from storage into MySQL.""" - - return await super().load_from_storage( - table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs - ) - def collect_rows(self, cursor: Any, fetched: "list[Any]") -> "tuple[list[Any], list[str], int]": """Collect asyncmy rows for the direct execution path.""" description = cursor.description or None @@ -444,10 +394,18 @@ def resolve_rowcount(self, cursor: Any) -> int: async def _execute_cache_hit( self, sql: str, params: "tuple[Any, ...] | list[Any] | dict[str, Any]", cached: "CachedQuery" ) -> "SQLResult": - statement = self._cached_statement( - sql, params, cached, params, params_are_simple=True, compiled_sql=cached.compiled_sql - ) - return await self._execute_cached_statement(statement) + if ( + "%" in cached.compiled_sql + and escape_literal_percent( + cached.compiled_sql, params or (None,), self.statement_config.parameter_validator + ) + != cached.compiled_sql + ): + statement = self._cached_statement( + sql, params, cached, params, params_are_simple=True, compiled_sql=cached.compiled_sql + ) + return await self._execute_cached_statement(statement) + return await super()._execute_cache_hit(sql, params, cached) async def _execute_cursor_statement(self, statement: str, parameters: Any = None) -> None: async with self.with_cursor(self.connection) as cursor: @@ -468,14 +426,38 @@ async def _execute_local_infile_cursor(self, statement: str, filename: str) -> N async def _load_from_arrow_via_local_infile( self, table: str, columns: "list[str]", records: "list[tuple[Any, ...]]" ) -> None: - with tempfile.TemporaryDirectory(prefix="sqlspec-asyncmy-") as directory: - with tempfile.NamedTemporaryFile(dir=directory, suffix=".tsv", delete=False) as payload: - payload.write(encode_records_for_local_infile(records)) - filename = payload.name + writer = asyncio.create_task(asyncio.to_thread(_write_asyncmy_local_infile_payload, records)) + pending_cancellation: asyncio.CancelledError | None = None + while not writer.done(): + try: + await asyncio.shield(writer) + except asyncio.CancelledError as exc: + pending_cancellation = exc + except BaseException: + break + try: + directory, filename = writer.result() + except BaseException: + if pending_cancellation is not None: + raise pending_cancellation + raise + try: + if pending_cancellation is not None: + raise pending_cancellation statement = build_load_data_statement(table, columns) exc_handler = self.handle_database_exceptions() await self._run_with_exception_handler(exc_handler, self._execute_local_infile_cursor, statement, filename) self._check_pending_exception(exc_handler) + finally: + cleanup = asyncio.create_task(asyncio.to_thread(shutil.rmtree, directory, ignore_errors=True)) + while not cleanup.done(): + try: + await asyncio.shield(cleanup) + except asyncio.CancelledError as exc: + pending_cancellation = exc + cleanup.result() + if pending_cancellation is not None: + raise pending_cancellation def _connection_in_transaction(self) -> bool: """Check if connection is in transaction. @@ -483,7 +465,24 @@ def _connection_in_transaction(self) -> bool: Returns: True when the server reports an open transaction. """ - return bool(getattr(self.connection, "server_status", 0) & 1) + try: + return bool(self.connection.server_status & 1) + except AttributeError: + return False + + +def _write_asyncmy_local_infile_payload(records: "list[tuple[Any, ...]]") -> tuple[str, str]: + """Write LOCAL INFILE payload to a temporary file in a dedicated directory.""" + directory = tempfile.mkdtemp(prefix="sqlspec-asyncmy-") + try: + with tempfile.NamedTemporaryFile(dir=directory, suffix=".tsv", delete=False) as payload: + payload.write(encode_records_for_local_infile(records)) + filename = payload.name + except BaseException: + shutil.rmtree(directory, ignore_errors=True) + raise + else: + return directory, filename def _resolve_column_types(description: Any) -> "dict[str, str] | None": diff --git a/sqlspec/adapters/asyncmy/litestar/store.py b/sqlspec/adapters/asyncmy/litestar/store.py index 0c0441c54..865664ee7 100644 --- a/sqlspec/adapters/asyncmy/litestar/store.py +++ b/sqlspec/adapters/asyncmy/litestar/store.py @@ -141,10 +141,10 @@ async def set(self, key: str, value: "str | bytes", expires_in: "int | timedelta sql = f""" INSERT INTO {self._table_name} (session_id, data, expires_at) - VALUES (%s, %s, %s) AS new + VALUES (%s, %s, %s) ON DUPLICATE KEY UPDATE - data = new.data, - expires_at = new.expires_at, + data = VALUES(data), + expires_at = VALUES(expires_at), updated_at = UTC_TIMESTAMP(6) """ @@ -304,12 +304,9 @@ def _drop_table_sql(self) -> "list[str]": """Get MySQL/MariaDB DROP TABLE SQL statements. Returns: - List of SQL statements to drop indexes and table. + List of SQL statements to drop table. """ - return [ - f"DROP INDEX idx_{self._table_name}_expires_at ON {self._table_name}", - f"DROP TABLE IF EXISTS {self._table_name}", - ] + return [f"DROP TABLE IF EXISTS {self._table_name}"] def _is_mysql_table_missing(exc: BaseException) -> bool: diff --git a/sqlspec/adapters/asyncpg/adk/store.py b/sqlspec/adapters/asyncpg/adk/store.py index 139e289d5..e45c50148 100644 --- a/sqlspec/adapters/asyncpg/adk/store.py +++ b/sqlspec/adapters/asyncpg/adk/store.py @@ -1,11 +1,12 @@ """AsyncPG ADK store for Google Agent Development Kit session/event storage.""" from contextlib import nullcontext -from typing import TYPE_CHECKING, Any, Final, Literal, NotRequired, cast +from typing import TYPE_CHECKING, Any, Literal, NotRequired, cast from sqlspec.adapters.asyncpg._typing import asyncpg_module as asyncpg -from sqlspec.adapters.asyncpg.core import parse_status +from sqlspec.adapters.asyncpg.core import build_postgres_storage_params_sql, parse_status from sqlspec.config import ADKConfig, AsyncConfigT +from sqlspec.exceptions import ImproperConfigurationError from sqlspec.extensions.adk import BaseAsyncADKStore, StoredEvent, StoredSession, normalize_session_list_options from sqlspec.extensions.adk.memory.store import BaseAsyncADKMemoryStore @@ -19,8 +20,6 @@ __all__ = ("AsyncpgADKConfig", "AsyncpgADKMemoryStore", "AsyncpgADKStore") -POSTGRES_TABLE_NOT_FOUND_ERROR: Final = "42P01" - class AsyncpgADKConfig(ADKConfig): """Asyncpg-specific ADK extension settings. @@ -480,7 +479,9 @@ async def _sessions_table_ddl(self) -> str: if self._owner_id_column_ddl: owner_id_line = f",\n {self._owner_id_column_ddl}" - table_options = _postgres_table_options(_adk_config(self._config)) + table_options = build_postgres_storage_params_sql( + _adk_config(self._config), default_fillfactor=80, section="adk" + ) return f""" CREATE TABLE IF NOT EXISTS {self._session_table} ( id VARCHAR(128) PRIMARY KEY, @@ -504,7 +505,9 @@ async def _sessions_table_ddl(self) -> str: async def _events_table_ddl(self) -> str: adk_config = _adk_config(self._config) - table_options = _postgres_table_options(adk_config, include_autovacuum=True) + table_options = build_postgres_storage_params_sql( + adk_config, include_autovacuum=True, default_fillfactor=80, section="adk" + ) generated_columns = "" generated_indexes = "" if adk_config.get("enable_event_generated_columns", False): @@ -544,7 +547,9 @@ async def _events_table_ddl(self) -> str: """ async def _app_states_table_ddl(self) -> str: - table_options = _postgres_table_options(_adk_config(self._config)) + table_options = build_postgres_storage_params_sql( + _adk_config(self._config), default_fillfactor=80, section="adk" + ) return f""" CREATE TABLE IF NOT EXISTS {self._app_state_table} ( app_name VARCHAR(128) PRIMARY KEY, @@ -554,7 +559,9 @@ async def _app_states_table_ddl(self) -> str: """ async def _user_states_table_ddl(self) -> str: - table_options = _postgres_table_options(_adk_config(self._config)) + table_options = build_postgres_storage_params_sql( + _adk_config(self._config), default_fillfactor=80, section="adk" + ) return f""" CREATE TABLE IF NOT EXISTS {self._user_state_table} ( app_name VARCHAR(128) NOT NULL, @@ -843,6 +850,9 @@ async def _memory_table_ddl(self) -> str: ) if self._vector_index_type == "scann": + if self._scann_quantizer.upper() not in {"SQ8", "FP32", "AH"}: + msg = f"Invalid SCaNN quantizer '{self._scann_quantizer}'. Must be one of: SQ8, FP32, AH" + raise ImproperConfigurationError(msg) indexes.append( f"CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_scann ON {self._memory_table} USING scann (embedding) WITH (num_leaves = {self._scann_num_leaves}, quantizer = '{self._scann_quantizer}');" ) @@ -900,30 +910,3 @@ def _adk_config(config: Any) -> AsyncpgADKConfig: if not isinstance(adk_config, dict): return {} return cast("AsyncpgADKConfig", adk_config) - - -def _postgres_table_options(adk_config: AsyncpgADKConfig, *, include_autovacuum: bool = False) -> str: - options = [_postgres_fillfactor_option(adk_config)] - if include_autovacuum: - options.extend(_postgres_autovacuum_options(adk_config)) - return f" WITH ({', '.join(options)})" - - -def _postgres_fillfactor_option(adk_config: AsyncpgADKConfig) -> str: - value = adk_config.get("fillfactor", 80) - if not isinstance(value, int) or isinstance(value, bool) or value not in range(10, 101): - msg = "extension_config['adk']['fillfactor'] must be an integer from 10 to 100" - raise ValueError(msg) - return f"fillfactor = {value}" - - -def _postgres_autovacuum_options(adk_config: AsyncpgADKConfig) -> "list[str]": - options: list[str] = [] - for key in ("autovacuum_vacuum_scale_factor", "autovacuum_analyze_scale_factor"): - value = adk_config.get(key) - if value is not None: - if not isinstance(value, (int, float)) or isinstance(value, bool) or not 0 <= float(value) <= 1: - msg = f"extension_config['adk']['{key}'] must be a number from 0 to 1" - raise ValueError(msg) - options.append(f"{key} = {float(value):g}") - return options diff --git a/sqlspec/adapters/asyncpg/config.py b/sqlspec/adapters/asyncpg/config.py index 394954974..0a1799027 100644 --- a/sqlspec/adapters/asyncpg/config.py +++ b/sqlspec/adapters/asyncpg/config.py @@ -123,8 +123,6 @@ class AsyncpgConnectionConfig(TypedDict): """Default per-command execution timeout in seconds.""" statement_cache_size: NotRequired[int] """Maximum number of server-side prepared statements cached per connection (``0`` disables caching).""" - pgbouncer: NotRequired[bool] - """Disable prepared statement caching for transaction-pooling proxies such as PgBouncer.""" max_cached_statement_lifetime: NotRequired[int] """Maximum age in seconds before cached prepared statements are evicted.""" max_cacheable_statement_size: NotRequired[int] @@ -201,6 +199,8 @@ class AsyncpgDriverFeatures(TypedDict): """AlloyDB network path (``PUBLIC``, ``PRIVATE``, or ``PSC``; defaults to ``PRIVATE``).""" enable_alloydb_columnar: NotRequired[bool] """Enable ``google_columnar_engine.enable_columnar_scan=on`` at session startup or per transaction.""" + enable_pgcopy_binary: NotRequired[bool] + """Enable binary COPY format for Arrow table ingestion.""" enable_events: NotRequired[bool] """Enable SQLSpec event channel pub/sub support (defaults to ``True`` when ``events`` is configured).""" events_backend: NotRequired[Literal["notify", "notify_queue", "poll_queue"]] @@ -209,10 +209,6 @@ class AsyncpgDriverFeatures(TypedDict): """Async hook invoked with each raw ``asyncpg`` connection after internal codec setup.""" pooler_mode: NotRequired[Literal["session", "transaction", "statement", "auto"] | None] """External pooler compatibility mode; sets ``statement_cache_size=0`` in ``transaction`` or ``statement`` mode.""" - pgbouncer: NotRequired[bool] - """Alias for ``pooler_mode='transaction'`` disabling server-side prepared statements.""" - pgbouncer_mode: NotRequired[bool] - """Alias for ``pooler_mode='transaction'`` disabling server-side prepared statements.""" direct_dsn: NotRequired["str | Mapping[str, Any] | None"] """Direct PostgreSQL DSN or connection mapping bypassing the pooler for migrations and ``LISTEN/NOTIFY``.""" type_codecs: NotRequired["list[dict[str, Any]]"] @@ -292,13 +288,13 @@ class AsyncpgConfig(AsyncDatabaseConfig[AsyncpgConnection, "Pool[Record]", Async """Configuration for AsyncPG database connections using TypedDict.""" driver_type: "ClassVar[type[AsyncpgDriver]]" = AsyncpgDriver - connection_type: "ClassVar[type[AsyncpgConnection]]" = AsyncpgConnection # type: ignore[assignment] + connection_type: "ClassVar[type[Any]]" = cast("type[Any]", AsyncpgConnection) supports_transactional_ddl: "ClassVar[bool]" = True supports_migration_schemas: "ClassVar[bool]" = True - supports_native_arrow_export: "ClassVar[bool]" = True + supports_native_arrow_export: "ClassVar[bool]" = False supports_native_arrow_import: "ClassVar[bool]" = True - supports_native_parquet_export: "ClassVar[bool]" = True - supports_native_parquet_import: "ClassVar[bool]" = True + supports_native_parquet_export: "ClassVar[bool]" = False + supports_native_parquet_import: "ClassVar[bool]" = False supports_native_row_streaming: "ClassVar[bool]" = True type_coercion_capabilities: "ClassVar[TypeCoercionCapabilities]" = TypeCoercionCapabilities( datetime_binding="native", timestamp_precision="microsecond", json_columns_decoded=True, uuid_binding="native" @@ -349,16 +345,9 @@ def __init__( dict(raw_server_settings) if isinstance(raw_server_settings, Mapping) else None ) normalized_conn_cfg = build_connection_config(raw_conn_cfg) - if raw_conn_cfg.get("pgbouncer"): - features_dict["pgbouncer"] = True - features_dict["pgbouncer_mode"] = True - if features_dict.get("pooler_mode") in {None, "session", "auto"}: - features_dict["pooler_mode"] = "transaction" - is_transaction_pooler = features_dict.get("pooler_mode") in {"transaction", "statement"} or bool( - features_dict.get("pgbouncer") or features_dict.get("pgbouncer_mode") - ) + is_transaction_pooler = features_dict.get("pooler_mode") in {"transaction", "statement"} if is_transaction_pooler: - normalized_conn_cfg.setdefault("statement_cache_size", 0) + normalized_conn_cfg["statement_cache_size"] = 0 if bool(features_dict.get("enable_alloydb_columnar", False)) and not is_transaction_pooler: merged_settings = dict(self._explicit_server_settings) if self._explicit_server_settings is not None else {} merged_settings.setdefault("google_columnar_engine.enable_columnar_scan", "on") @@ -412,7 +401,7 @@ async def create_connection(self) -> "AsyncpgConnection": for key in _POOL_ONLY_CONFIG_KEYS: config.pop(key, None) - if self.driver_features.get("pgbouncer"): + if self.driver_features.get("pooler_mode") in {"transaction", "statement"}: config["statement_cache_size"] = 0 if self.driver_features.get("enable_cloud_sql", False): @@ -450,12 +439,14 @@ def provide_session( return AsyncpgSessionContext( acquire_connection=factory.acquire_connection, release_connection=factory.release_connection, - statement_config=statement_config - or (lambda: resolve_runtime_statement_config(None, self.statement_config, default_statement_config)), + statement_config=statement_config or self._default_runtime_statement_config, driver_features=self.driver_features, prepare_driver=self._prepare_driver, ) + def _default_runtime_statement_config(self) -> "StatementConfig": + return resolve_runtime_statement_config(None, self.statement_config, default_statement_config) + def get_signature_namespace(self) -> "dict[str, Any]": """Get the signature namespace for AsyncPG types. @@ -510,8 +501,6 @@ def derive_direct_config(self) -> Self: conn_cfg.update(build_connection_config(direct_dsn)) features: dict[str, Any] = dict(self.driver_features) features["pooler_mode"] = "session" - features["pgbouncer"] = False - features["pgbouncer_mode"] = False features["direct_dsn"] = None if self._user_connection_hook is not None: features["on_connection_create"] = self._user_connection_hook @@ -644,7 +633,7 @@ async def _create_pool(self) -> "Pool[Record]": config["password"] = auth_token_provider config.setdefault("max_inactive_connection_lifetime", _AUTH_TOKEN_DEFAULT_INACTIVE_LIFETIME) - if self.connection_config.get("pgbouncer") or self.driver_features.get("pgbouncer"): + if self.driver_features.get("pooler_mode") in {"transaction", "statement"}: config["statement_cache_size"] = 0 if self.driver_features.get("enable_cloud_sql", False): diff --git a/sqlspec/adapters/asyncpg/core.py b/sqlspec/adapters/asyncpg/core.py index 384f59bb3..0aa4d2689 100644 --- a/sqlspec/adapters/asyncpg/core.py +++ b/sqlspec/adapters/asyncpg/core.py @@ -1,22 +1,19 @@ """AsyncPG adapter compiled helpers.""" +import asyncio import contextlib import datetime as dt import re +import sys from collections.abc import Sized -from typing import TYPE_CHECKING, Any, Final, NamedTuple +from typing import TYPE_CHECKING, Any, NamedTuple from sqlspec.adapters.asyncpg._typing import asyncpg_module as asyncpg from sqlspec.core import DriverParameterProfile, ParameterStyle, StatementConfig, build_statement_config_from_profile -from sqlspec.core.config_runtime import ( - build_postgres_extension_probe_names, - is_postgres_extension_active, - resolve_postgres_extension_state, - resolve_runtime_statement_config, -) +from sqlspec.core.config_runtime import resolve_runtime_statement_config from sqlspec.exceptions import ( CheckViolationError, - ConnectionTimeoutError, + DatabaseConnectionError, DeadlockError, ForeignKeyViolationError, ImproperConfigurationError, @@ -25,6 +22,7 @@ OperationalError, OperationCancelledError, PermissionDeniedError, + QueryTimeoutError, SerializationConflictError, SQLParsingError, SQLSpecError, @@ -50,7 +48,7 @@ "NormalizedStackOperation", "apply_driver_features", "build_connection_config", - "build_postgres_extension_probe_names", + "build_postgres_storage_params_sql", "build_profile", "build_statement_config", "collect_rows", @@ -58,13 +56,10 @@ "create_mapped_exception", "default_statement_config", "driver_profile", - "invoke_prepared_statement", - "is_postgres_extension_active", "parse_status", "register_json_codecs", "register_pgvector_support", "resolve_many_rowcount", - "resolve_postgres_extension_state", "resolve_runtime_statement_config", ) @@ -85,7 +80,6 @@ class NormalizedStackOperation(NamedTuple): parameters: "tuple[Any, ...] | dict[str, Any] | None" -PREPARED_STATEMENT_CACHE_SIZE: Final[int] = 32 _EXCEPTION_MAPPING_DISPATCHER = TypeDispatcher["tuple[str, type[SQLSpecError], str]"]() @@ -132,12 +126,42 @@ def build_connection_config(connection_config: "Mapping[str, Any]") -> "dict[str if found_user: config["user"] = user_val - if config.pop("pgbouncer", False): - config["statement_cache_size"] = 0 - return config +def build_postgres_storage_params_sql( + settings: "dict[str, object] | Mapping[str, object] | None", + *, + include_autovacuum: bool = False, + default_fillfactor: int | None = None, + default_vacuum_scale: float | None = None, + default_analyze_scale: float | None = None, + section: str = "storage", +) -> str: + """Format PostgreSQL table storage parameters (WITH fillfactor = ...) clause.""" + options: list[str] = [] + config_dict = dict(settings) if settings else {} + fillfactor = config_dict.get("fillfactor", default_fillfactor) + if fillfactor is not None: + if not isinstance(fillfactor, int) or isinstance(fillfactor, bool) or fillfactor not in range(10, 101): + msg = f"extension_config['{section}']['fillfactor'] must be an integer from 10 to 100" + raise ValueError(msg) + options.append(f"fillfactor = {fillfactor}") + if include_autovacuum: + scales = ( + ("autovacuum_vacuum_scale_factor", default_vacuum_scale), + ("autovacuum_analyze_scale_factor", default_analyze_scale), + ) + for key, default_val in scales: + value = config_dict.get(key, default_val) + if value is not None: + if not isinstance(value, (int, float)) or isinstance(value, bool) or not 0 <= float(value) <= 1: + msg = f"extension_config['{section}']['{key}'] must be a number from 0 to 1" + raise ValueError(msg) + options.append(f"{key} = {float(value):g}") + return f" WITH ({', '.join(options)})" if options else "" + + def build_profile() -> "DriverParameterProfile": """Create the AsyncPG driver parameter profile.""" @@ -170,39 +194,6 @@ def configure_parameter_serializers( return parameter_config.replace(json_serializer=serializer, json_deserializer=effective_deserializer) -async def invoke_prepared_statement( - prepared: Any, parameters: "tuple[Any, ...] | dict[str, Any] | list[Any] | None", *, fetch: bool, **kwargs: Any -) -> Any: - """Invoke an AsyncPG prepared statement with optional parameters. - - Args: - prepared: AsyncPG prepared statement object. - parameters: Prepared parameters payload. - fetch: Whether to fetch rows. - **kwargs: Native execution options, including timeout. - - Returns: - Query result or status message. - """ - if parameters is None: - if fetch: - return await prepared.fetch(**kwargs) - await prepared.fetch(**kwargs) - return prepared.get_statusmsg() - - if isinstance(parameters, dict): - values = tuple(parameters.values()) - if fetch: - return await prepared.fetch(*values, **kwargs) - await prepared.fetch(*values, **kwargs) - return prepared.get_statusmsg() - - if fetch: - return await prepared.fetch(*parameters, **kwargs) - await prepared.fetch(*parameters, **kwargs) - return prepared.get_statusmsg() - - def build_statement_config( *, json_serializer: "Callable[[Any], str] | None" = None, json_deserializer: "Callable[[str], Any] | None" = None ) -> "StatementConfig": @@ -226,20 +217,56 @@ def build_statement_config( return base_config.replace(parameter_config=parameter_config) +class _JsonEncoder: + __slots__ = ("_encoder",) + + def __init__(self, encoder: "Callable[[Any], str]") -> None: + self._encoder = encoder + + def __call__(self, value: Any) -> bytes: + return _encode_json_payload(value, self._encoder) + + +class _JsonDecoder: + __slots__ = ("_decoder",) + + def __init__(self, decoder: "Callable[[str], Any]") -> None: + self._decoder = decoder + + def __call__(self, value: Any) -> Any: + return _decode_json_payload(value, self._decoder) + + +class _JsonbEncoder: + __slots__ = ("_encoder",) + + def __init__(self, encoder: "Callable[[Any], str]") -> None: + self._encoder = encoder + + def __call__(self, value: Any) -> bytes: + return _encode_jsonb_payload(value, self._encoder) + + +class _JsonbDecoder: + __slots__ = ("_decoder",) + + def __init__(self, decoder: "Callable[[str], Any]") -> None: + self._decoder = decoder + + def __call__(self, value: Any) -> Any: + return _decode_jsonb_payload(value, self._decoder) + + async def register_json_codecs(connection: Any, encoder: Any, decoder: Any) -> None: """Register JSON type codecs on asyncpg connection.""" try: await connection.set_type_codec( - "json", - encoder=lambda value: _encode_json_payload(value, encoder), - decoder=lambda value: _decode_json_payload(value, decoder), - schema="pg_catalog", - format="binary", + "json", encoder=_JsonEncoder(encoder), decoder=_JsonDecoder(decoder), schema="pg_catalog", format="binary" ) await connection.set_type_codec( "jsonb", - encoder=lambda value: _encode_jsonb_payload(value, encoder), - decoder=lambda value: _decode_jsonb_payload(value, decoder), + encoder=_JsonbEncoder(encoder), + decoder=_JsonbDecoder(decoder), schema="pg_catalog", format="binary", ) @@ -281,12 +308,9 @@ def apply_driver_features( processed_features.setdefault("enable_alloydb", False) processed_features.setdefault("enable_events", True) processed_features.setdefault("events_backend", "notify") - pgbouncer_enabled = bool(processed_features.get("pgbouncer") or processed_features.get("pgbouncer_mode")) - processed_features["pgbouncer"] = pgbouncer_enabled - processed_features["pgbouncer_mode"] = pgbouncer_enabled raw_pooler_mode = processed_features.get("pooler_mode") if raw_pooler_mode is None: - processed_features["pooler_mode"] = "transaction" if pgbouncer_enabled else "session" + processed_features["pooler_mode"] = "session" else: normalized_mode = str(raw_pooler_mode).lower() if normalized_mode not in {"session", "transaction", "statement", "auto"}: @@ -294,12 +318,7 @@ def apply_driver_features( f"Invalid pooler_mode: {raw_pooler_mode!r}. Expected 'session', 'transaction', 'statement', or 'auto'." ) raise ImproperConfigurationError(msg) - processed_features["pooler_mode"] = ( - "transaction" if normalized_mode == "auto" and pgbouncer_enabled else normalized_mode - ) - if processed_features["pooler_mode"] in {"transaction", "statement"}: - processed_features["pgbouncer"] = True - processed_features["pgbouncer_mode"] = True + processed_features["pooler_mode"] = "session" if normalized_mode == "auto" else normalized_mode processed_features.setdefault("direct_dsn", None) processed_features.setdefault("enable_alloydb_columnar", False) @@ -381,10 +400,10 @@ def resolve_many_rowcount(parameter_sets: Any, *, fallback_count: "int | None" = asyncpg.exceptions.InvalidAuthorizationSpecificationError, ("28000", PermissionDeniedError, "authorization error") ) _EXCEPTION_MAPPING_DISPATCHER.register( - asyncpg.exceptions.ConnectionDoesNotExistError, ("08003", ConnectionTimeoutError, "connection does not exist") + asyncpg.exceptions.ConnectionDoesNotExistError, ("08003", DatabaseConnectionError, "connection does not exist") ) _EXCEPTION_MAPPING_DISPATCHER.register( - asyncpg.exceptions.CannotConnectNowError, ("57P03", ConnectionTimeoutError, "cannot connect now") + asyncpg.exceptions.CannotConnectNowError, ("57P03", DatabaseConnectionError, "cannot connect now") ) _EXCEPTION_MAPPING_DISPATCHER.register( asyncpg.exceptions.PostgresSyntaxError, ("42601", SQLParsingError, "SQL syntax error") @@ -410,7 +429,14 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpec Returns: A SQLSpec exception that wraps the original error """ - del logger + _ = logger + if isinstance(error, (asyncio.TimeoutError, TimeoutError)): + return QueryTimeoutError(str(error)) + if isinstance(error, OSError) or ( + asyncpg is not None and isinstance(error, getattr(asyncpg, "InterfaceError", ())) + ): + return DatabaseConnectionError(str(error)) + mapped_error = _EXCEPTION_MAPPING_DISPATCHER.get(error) if mapped_error is not None: error_code, error_class, description = mapped_error @@ -453,7 +479,7 @@ def collect_rows(records: "list[Any] | None") -> "tuple[list[Any], list[str]]": class AsyncpgStreamSource: """Compiled async chunk source streaming dict rows from an asyncpg cursor in a stream-owned transaction.""" - __slots__ = ("_chunk_size", "_cursor", "_driver", "_parameters", "_sql", "_timeout_args", "_transaction") + __slots__ = ("_chunk_size", "_cursor", "_driver", "_parameters", "_sql", "_timeout", "_transaction") def __init__( self, @@ -461,9 +487,14 @@ def __init__( sql: str, parameters: "tuple[Any, ...]", chunk_size: int, - timeout_args: "dict[str, Any] | None" = None, + timeout: "float | dict[str, Any] | None" = None, ) -> None: - self._timeout_args = timeout_args or {} + if isinstance(timeout, dict): + self._timeout: float | None = ( + float(timeout["timeout"]) if "timeout" in timeout and timeout["timeout"] is not None else None + ) + else: + self._timeout = float(timeout) if timeout is not None else None self._driver = driver self._sql = sql self._parameters = parameters @@ -473,47 +504,71 @@ def __init__( async def start(self) -> None: handler = self._driver.handle_database_exceptions() - await self._driver._run_with_exception_handler(handler, self._start) + if sys.exc_info()[1] is None: + async with handler: + await self._start() + else: + await self._driver._run_with_exception_handler(handler, self._start) self._driver._check_pending_exception(handler) async def fetch_chunk(self) -> "list[dict[str, Any]]": handler = self._driver.handle_database_exceptions() - records = await self._driver._run_with_exception_handler( - handler, self._cursor.fetch, self._chunk_size, **self._timeout_args - ) + if sys.exc_info()[1] is None: + async with handler: + records = ( + await self._cursor.fetch(self._chunk_size, timeout=self._timeout) + if self._timeout is not None + else await self._cursor.fetch(self._chunk_size) + ) + elif self._timeout is not None: + records = await self._driver._run_with_exception_handler( + handler, self._cursor.fetch, self._chunk_size, timeout=self._timeout + ) + else: + records = await self._driver._run_with_exception_handler(handler, self._cursor.fetch, self._chunk_size) self._driver._check_pending_exception(handler) assert records is not None return [dict(record) for record in records] async def close(self, error: bool = False) -> None: + cursor = self._cursor self._cursor = None transaction = self._transaction self._transaction = None - if transaction is not None: - try: + try: + if cursor is not None and getattr(cursor, "_portal_name", None): + # asyncpg has no public cursor close; CockroachDB requires closing its portal before transaction control. + await cursor._close_portal(self._timeout) + if transaction is not None: if error: await transaction.rollback() else: await transaction.commit() - except Exception: - with contextlib.suppress(Exception): - await transaction.rollback() + except BaseException: + with contextlib.suppress(Exception): + await self._driver.connection.close() + raise async def _start(self) -> None: - in_tx_fn = getattr(self._driver, "_connection_in_transaction", None) - already_in_tx = bool(in_tx_fn()) if callable(in_tx_fn) else False + already_in_tx = self._driver._connection_in_transaction() if not already_in_tx: transaction = self._driver.connection.transaction() await transaction.start() self._transaction = transaction try: - self._cursor = await self._driver.connection.cursor(self._sql, *self._parameters, **self._timeout_args) + if self._timeout is not None: + self._cursor = await self._driver.connection.cursor(self._sql, *self._parameters, timeout=self._timeout) + else: + self._cursor = await self._driver.connection.cursor(self._sql, *self._parameters) except BaseException: transaction = self._transaction self._transaction = None if transaction is not None: - with contextlib.suppress(Exception): + try: await transaction.rollback() + except BaseException: + with contextlib.suppress(Exception): + await self._driver.connection.close() raise @@ -554,10 +609,7 @@ def _custom_type_coercions() -> "dict[type, Callable[[Any], Any]]": def _is_pre_serialized_json(value: str) -> bool: stripped = value.strip() - if not stripped: - return False - first = stripped[0] - if first not in '{["-0123456789tfn': + if not stripped or stripped[0] not in '{["-0123456789tfn': return False try: from_json(stripped) diff --git a/sqlspec/adapters/asyncpg/data_dictionary.py b/sqlspec/adapters/asyncpg/data_dictionary.py index a2978b941..a212bece4 100644 --- a/sqlspec/adapters/asyncpg/data_dictionary.py +++ b/sqlspec/adapters/asyncpg/data_dictionary.py @@ -26,7 +26,12 @@ system_metadata_gated_result, unsupported_system_metadata_capability, ) -from sqlspec.data_dictionary.dialects.postgres import resolve_postgres_json_type +from sqlspec.data_dictionary.dialects.postgres import ( + POSTGRES_METADATA_DOMAINS, + POSTGRES_SUPPORTED_DOMAINS, + POSTGRES_SYSTEM_METADATA_QUERIES, + resolve_postgres_json_type, +) from sqlspec.driver import AsyncDataDictionaryBase if TYPE_CHECKING: @@ -36,37 +41,6 @@ __all__ = ("AsyncpgDataDictionary",) -_POSTGRES_METADATA_DOMAINS: tuple[str, ...] = ( - "schemas", - "objects", - "tables", - "columns", - "constraints", - "indexes", - "views", - "materialized_views", - "sequences", - "types", - "routines", - "triggers", - "comments", - "privileges", - "dependencies", - "extensions", - "partitions", - "ddl", - "system", -) -_POSTGRES_SUPPORTED_DOMAINS = frozenset(d for d in _POSTGRES_METADATA_DOMAINS if d != "system") -_POSTGRES_SYSTEM_METADATA_QUERIES = { - "settings": "settings", - "statement_history": "pg_stat_statements", - "pg_stat_statements": "pg_stat_statements", - "table_statistics": "table_stats", - "table_stats": "table_stats", -} - - class AsyncpgDataDictionary(AsyncDataDictionaryBase): """PostgreSQL-specific async data dictionary.""" @@ -83,7 +57,7 @@ async def get_system_metadata_capabilities( ) -> tuple[SystemMetadataCapability, ...]: """Get PostgreSQL opt-in system metadata capability disclosures.""" _ = driver - requested_domains = tuple(_POSTGRES_SYSTEM_METADATA_QUERIES) if domains is None else tuple(domains) + requested_domains = tuple(POSTGRES_SYSTEM_METADATA_QUERIES) if domains is None else tuple(domains) return tuple(_postgres_system_metadata_capability(domain) for domain in requested_domains) async def get_schemas(self, driver: "AsyncpgDriver") -> MetadataResult: @@ -157,7 +131,7 @@ async def get_system_metadata( if gate_result.capability.support != MetadataSupport.SUPPORTED: return gate_result - query_name = _POSTGRES_SYSTEM_METADATA_QUERIES.get(metadata_request.domain) + query_name = POSTGRES_SYSTEM_METADATA_QUERIES.get(metadata_request.domain) if query_name is None: return gate_result parameters: dict[str, object] = {} @@ -284,10 +258,11 @@ async def get_indexes( ) table_name = self.resolve_identifier(table) - self._log_table_describe(driver, schema_name=schema_name, table_name=table_name, operation="indexes") + table_schema_name = None if schema is None else schema_name + self._log_table_describe(driver, schema_name=table_schema_name, table_name=table_name, operation="indexes") return await driver.select( _postgres_domain_sql("indexes", "by_schema"), - schema_name=schema_name, + schema_name=table_schema_name, table_name=table_name, schema_type=IndexMetadata, ) @@ -352,7 +327,7 @@ def _postgres_metadata_capability(domain: str) -> MetadataCapability: risks=(MetadataRisk.EXPENSIVE, MetadataRisk.PRIVILEGED), warnings=("System metadata is opt-in and disabled by default.",), ) - if domain in _POSTGRES_SUPPORTED_DOMAINS: + if domain in POSTGRES_SUPPORTED_DOMAINS: return MetadataCapability( domain=domain, support=MetadataSupport.SUPPORTED, @@ -363,7 +338,7 @@ def _postgres_metadata_capability(domain: str) -> MetadataCapability: def _postgres_metadata_profile(adapter: str, domains: Sequence[str] | None) -> MetadataCapabilityProfile: - requested_domains = _POSTGRES_METADATA_DOMAINS if domains is None else tuple(domains) + requested_domains = POSTGRES_METADATA_DOMAINS if domains is None else tuple(domains) return MetadataCapabilityProfile( "postgres", adapter=adapter, @@ -422,7 +397,7 @@ def _ddl_result_from_rows( def _postgres_system_metadata_capability(domain: str) -> SystemMetadataCapability: - if domain not in _POSTGRES_SYSTEM_METADATA_QUERIES: + if domain not in POSTGRES_SYSTEM_METADATA_QUERIES: return unsupported_system_metadata_capability(domain) return SystemMetadataCapability( domain, diff --git a/sqlspec/adapters/asyncpg/driver.py b/sqlspec/adapters/asyncpg/driver.py index 1ab2c718f..2a50b6f7a 100644 --- a/sqlspec/adapters/asyncpg/driver.py +++ b/sqlspec/adapters/asyncpg/driver.py @@ -2,7 +2,7 @@ import re from collections.abc import Mapping -from contextlib import suppress +from contextlib import asynccontextmanager, nullcontext, suppress from io import BytesIO from itertools import chain from typing import TYPE_CHECKING, Any, Literal, cast @@ -11,6 +11,7 @@ from sqlglot.errors import ParseError from sqlspec.adapters.asyncpg._typing import AsyncpgCursor, AsyncpgPostgresError, AsyncpgSessionContext +from sqlspec.adapters.asyncpg._typing import asyncpg_module as asyncpg from sqlspec.adapters.asyncpg.core import ( AsyncpgStreamSource, NormalizedStackOperation, @@ -24,6 +25,7 @@ from sqlspec.adapters.asyncpg.data_dictionary import AsyncpgDataDictionary from sqlspec.core import ( SQL, + DMLResult, StackResult, StatementStack, get_cache_config, @@ -37,18 +39,20 @@ BaseAsyncExceptionHandler, StackExecutionObserver, describe_stack_statement, + validate_savepoint_name, ) from sqlspec.exceptions import ImproperConfigurationError, SQLSpecError, StackExecutionError from sqlspec.utils.logging import get_logger from sqlspec.utils.text import normalize_identifier, quote_identifier from sqlspec.utils.type_guards import has_sqlstate +from sqlspec.utils.uuids import uuid4 if TYPE_CHECKING: - from collections.abc import Sequence + from collections.abc import AsyncIterator, Sequence from sqlspec.adapters.asyncpg._typing import AsyncpgConnection from sqlspec.core import ArrowResult, SQLResult, StatementConfig - from sqlspec.driver import ExecutionResult + from sqlspec.driver import CachedQuery, ExecutionResult from sqlspec.storage import StorageBridgeJob, StorageDestination, StorageFormat, StorageTelemetry @@ -74,7 +78,11 @@ class AsyncpgExceptionHandler(BaseAsyncExceptionHandler): def _handle_exception(self, exc_type: "type[BaseException] | None", exc_val: "BaseException") -> bool: _ = exc_type - if isinstance(exc_val, AsyncpgPostgresError) or has_sqlstate(exc_val): + if ( + isinstance(exc_val, (AsyncpgPostgresError, TimeoutError, OSError)) + or has_sqlstate(exc_val) + or (asyncpg is not None and isinstance(exc_val, getattr(asyncpg, "InterfaceError", ()))) + ): self.pending_exception = create_mapped_exception(exc_val) return True return False @@ -132,8 +140,13 @@ async def dispatch_execute(self, cursor: "AsyncpgConnection", statement: "SQL") sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) params: tuple[Any, ...] = cast("tuple[Any, ...]", prepared_parameters) if prepared_parameters else () + timeout = self._timeout_value(statement.statement_config) if statement.returns_rows(): - records = await cursor.fetch(sql, *params, **self._timeout_arguments(statement.statement_config)) + records = ( + await cursor.fetch(sql, *params, timeout=timeout) + if timeout is not None + else await cursor.fetch(sql, *params) + ) data, column_names = collect_rows(records) return self.create_execution_result( @@ -145,8 +158,11 @@ async def dispatch_execute(self, cursor: "AsyncpgConnection", statement: "SQL") row_format="record", ) - result = await cursor.execute(sql, *params, **self._timeout_arguments(statement.statement_config)) - + result = ( + await cursor.execute(sql, *params, timeout=timeout) + if timeout is not None + else await cursor.execute(sql, *params) + ) affected_rows = parse_status(result) return self.create_execution_result(cursor, rowcount_override=affected_rows) @@ -162,16 +178,53 @@ async def dispatch_execute_many(self, cursor: "AsyncpgConnection", statement: "S ExecutionResult with batch execution details """ sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) + timeout = self._timeout_value(statement.statement_config) if prepared_parameters: parameter_sets = cast("list[Sequence[object]]", prepared_parameters) - await cursor.executemany(sql, parameter_sets, **self._timeout_arguments(statement.statement_config)) + if timeout is not None: + await cursor.executemany(sql, parameter_sets, timeout=timeout) + else: + await cursor.executemany(sql, parameter_sets) affected_rows = resolve_many_rowcount(parameter_sets) else: affected_rows = 0 return self.create_execution_result(cursor, rowcount_override=affected_rows, is_many_result=True) + async def _execute_cache_hit_with_cursor( + self, sql: str, params: "tuple[Any, ...] | list[Any] | dict[str, Any]", cached: "CachedQuery" + ) -> "SQLResult": + wire_sql = self._format_wire_sql(cached.compiled_sql) + timeout = self._timeout_value(self.statement_config) + args = ( + params if isinstance(params, (tuple, list)) else tuple(params.values()) if isinstance(params, dict) else () + ) + if cached.operation_profile.returns_rows: + records = ( + await self.connection.fetch(wire_sql, *args, timeout=timeout) + if timeout is not None + else await self.connection.fetch(wire_sql, *args) + ) + data, column_names = collect_rows(records) + return self.build_cached_select_result( + sql, + params, + cached, + selected_data=data, + column_names=column_names, + data_row_count=len(data), + row_format="record", + compiled_sql=wire_sql, + ) + status = ( + await self.connection.execute(wire_sql, *args, timeout=timeout) + if timeout is not None + else await self.connection.execute(wire_sql, *args) + ) + affected_rows = parse_status(status) + return DMLResult(cached.operation_type, affected_rows) + async def dispatch_execute_script(self, cursor: "AsyncpgConnection", statement: "SQL") -> "ExecutionResult": """Execute SQL script with statement splitting and parameter handling. @@ -182,20 +235,7 @@ async def dispatch_execute_script(self, cursor: "AsyncpgConnection", statement: Returns: ExecutionResult with script execution details """ - sql, _ = self._compiled_sql(statement, self.statement_config) - statements = self.split_script_statements(sql, statement.statement_config, strip_trailing_semicolon=True) - - successful_count = 0 - last_result = None - - for stmt in statements: - result = await cursor.execute(stmt, **self._timeout_arguments(statement.statement_config)) - last_result = result - successful_count += 1 - - return self.create_execution_result( - last_result, statement_count=len(statements), successful_statements=successful_count, is_script_result=True - ) + return await super().dispatch_execute_script(cursor, statement) async def dispatch_special_handling(self, cursor: "AsyncpgConnection", statement: "SQL") -> "SQLResult | None": """Handle PostgreSQL COPY operations and other special cases. @@ -233,10 +273,15 @@ async def begin(self) -> None: transaction = self.connection.transaction() try: await transaction.start() + except Exception as e: + self._release_failed_transaction_claim(transaction) + raise create_mapped_exception(e) from e + try: if self._should_set_local_columnar(): await self.connection.execute("SET LOCAL google_columnar_engine.enable_columnar_scan = 'on'") - except AsyncpgPostgresError as e: - self._release_failed_transaction_claim(transaction) + except Exception as e: + with suppress(Exception): + await transaction.rollback() raise create_mapped_exception(e) from e self._transaction = transaction @@ -249,7 +294,7 @@ async def commit(self) -> None: await transaction.commit() elif self._connection_in_transaction(): await self.connection.execute("COMMIT") - except AsyncpgPostgresError as e: + except Exception as e: raise create_mapped_exception(e) from e async def rollback(self) -> None: @@ -261,7 +306,7 @@ async def rollback(self) -> None: await transaction.rollback() elif self._connection_in_transaction(): await self.connection.execute("ROLLBACK") - except AsyncpgPostgresError as e: + except Exception as e: raise create_mapped_exception(e) from e async def set_migration_session_schema(self, schema: str) -> None: @@ -299,9 +344,8 @@ def dispatch_select_stream(self, statement: "SQL", chunk_size: int) -> "AsyncRow return None sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) params: tuple[Any, ...] = cast("tuple[Any, ...]", prepared_parameters) if prepared_parameters else () - return AsyncRowStream( - AsyncpgStreamSource(self, sql, params, chunk_size, self._timeout_arguments(statement.statement_config)) - ) + timeout = self._timeout_value(statement.statement_config) + return AsyncRowStream(AsyncpgStreamSource(self, sql, params, chunk_size, timeout)) def handle_database_exceptions(self) -> "AsyncpgExceptionHandler": """Handle database exceptions with PostgreSQL error codes.""" @@ -336,7 +380,6 @@ async def select_to_storage( ``STDIN``/``STDOUT`` stream form. """ - self._require_capability("arrow_export_enabled") arrow_result = await self.select_to_arrow(statement, *parameters, statement_config=statement_config, **kwargs) async_pipeline = self._storage_pipeline() telemetry_payload = await self._write_storage_result( @@ -385,29 +428,33 @@ async def load_from_arrow( if overwrite: try: await self.connection.execute(f"TRUNCATE TABLE {quoted_target}") - except AsyncpgPostgresError as exc: + except Exception as exc: + mapped = create_mapped_exception(exc) msg = f"Failed to truncate table '{table}': {exc}" - raise SQLSpecError(msg) from exc + raise type(mapped)(msg) from exc columns, batch_iter = self._coerce_arrow_batches(source) first_batch = next((b for b in batch_iter if b.num_rows > 0), None) metrics: dict[str, int] = {"rows_processed": 0, "bytes_processed": 0} if first_batch is not None: all_batches = chain((first_batch,), (b for b in batch_iter if b.num_rows > 0)) use_binary = bool(self.driver_features.get("enable_pgcopy_binary", False)) - if use_binary: - async_chunks = self._iter_pgcopy_binary_chunks_async(all_batches, columns=columns, metrics=metrics) - await self.connection.copy_to_table( - table_name, source=async_chunks, columns=columns, schema_name=schema_name, format="binary" - ) - else: - for batch in all_batches: - metrics["rows_processed"] += int(batch.num_rows) - metrics["bytes_processed"] += int(batch.nbytes) - records = self._arrow_batch_to_rows(batch, columns) - if records: - await self.connection.copy_records_to_table( - table_name, records=records, columns=columns, schema_name=schema_name - ) + try: + if use_binary: + async_chunks = self._iter_pgcopy_binary_chunks_async(all_batches, columns=columns, metrics=metrics) + await self.connection.copy_to_table( + table_name, source=async_chunks, columns=columns, schema_name=schema_name, format="binary" + ) + else: + for batch in all_batches: + metrics["rows_processed"] += int(batch.num_rows) + metrics["bytes_processed"] += int(batch.nbytes) + records = self._arrow_batch_to_rows(batch, columns) + if records: + await self.connection.copy_records_to_table( + table_name, records=records, columns=columns, schema_name=schema_name + ) + except Exception as exc: + raise create_mapped_exception(exc) from exc telemetry_payload = self._ingest_telemetry_from_batches(metrics["rows_processed"], metrics["bytes_processed"]) telemetry_payload["destination"] = table self._attach_partition_telemetry(telemetry_payload, partitioner) @@ -448,6 +495,21 @@ async def load_from_records( self._require_capability("arrow_import_enabled") materialized = list(records) if not materialized: + if overwrite: + table_name, schema_name, quoted_target = self._copy_target(table) + try: + await self.connection.execute(f"TRUNCATE TABLE {quoted_target}") + except Exception as exc: + mapped = create_mapped_exception(exc) + msg = f"Failed to truncate table '{table}': {exc}" + raise type(mapped)(msg) from exc + empty_telemetry: StorageTelemetry = { + "bytes_processed": 0, + "destination": table, + "format": "records", + "rows_processed": 0, + } + return self._storage_job(empty_telemetry) msg = "load_from_records requires at least one record." raise ImproperConfigurationError(msg) @@ -484,12 +546,16 @@ async def load_from_records( if overwrite: try: await self.connection.execute(f"TRUNCATE TABLE {quoted_target}") - except AsyncpgPostgresError as exc: + except Exception as exc: + mapped = create_mapped_exception(exc) msg = f"Failed to truncate table '{table}': {exc}" - raise SQLSpecError(msg) from exc - await self.connection.copy_records_to_table( - table_name, records=copy_rows, columns=resolved_columns, schema_name=schema_name - ) + raise type(mapped)(msg) from exc + try: + await self.connection.copy_records_to_table( + table_name, records=copy_rows, columns=resolved_columns, schema_name=schema_name + ) + except Exception as exc: + raise create_mapped_exception(exc) from exc telemetry_payload: StorageTelemetry = { "bytes_processed": 0, "destination": table, @@ -498,29 +564,18 @@ async def load_from_records( } return self._storage_job(telemetry_payload) - async def load_from_storage( - self, - table: str, - source: "StorageDestination", - *, - file_format: "StorageFormat", - partitioner: "dict[str, object] | None" = None, - overwrite: bool = False, - **kwargs: Any, - ) -> "StorageBridgeJob": - """Read an artifact from storage and ingest it via COPY.""" - - return await super().load_from_storage( - table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs - ) - - def _timeout_arguments(self, config: "StatementConfig") -> "dict[str, Any]": + def _timeout_value(self, config: "StatementConfig") -> float | None: for options in (config.execution_args, self.statement_config.execution_args): if options: for key in ("timeout", "command_timeout"): if key in options: - return {"timeout": options[key]} - return {} + val = options[key] + return float(val) if val is not None else None + return None + + def _timeout_arguments(self, config: "StatementConfig") -> "dict[str, Any]": + timeout = self._timeout_value(config) + return {"timeout": timeout} if timeout is not None else {} def _release_failed_transaction_claim(self, transaction: Any) -> None: """Clear the connection's transaction claim left by a failed start.""" @@ -570,16 +625,19 @@ async def _execute_stack_native( ) -> "tuple[StackResult, ...]": results: list[StackResult] = [] + caller_in_tx = self._connection_in_transaction() transaction_cm = None - if not continue_on_error and not self._connection_in_transaction(): + if not continue_on_error and not caller_in_tx: transaction_cm = self.connection.transaction() with StackExecutionObserver(self, stack, continue_on_error, native_pipeline=False) as observer: if transaction_cm is not None: async with transaction_cm: - await self._run_stack_operations(stack, continue_on_error, observer, results) + await self._run_stack_operations( + stack, continue_on_error, observer, results, caller_in_tx=caller_in_tx + ) else: - await self._run_stack_operations(stack, continue_on_error, observer, results) + await self._run_stack_operations(stack, continue_on_error, observer, results, caller_in_tx=caller_in_tx) return tuple(results) @@ -589,6 +647,8 @@ async def _run_stack_operations( continue_on_error: bool, observer: "StackExecutionObserver", results: "list[StackResult]", + *, + caller_in_tx: bool = False, ) -> None: """Run operations for statement stack execution. @@ -596,27 +656,32 @@ async def _run_stack_operations( """ for index, operation in enumerate(stack.operations): try: - normalized: NormalizedStackOperation | None = None - if operation.method == "execute": - kwargs = dict(operation.keyword_arguments) if operation.keyword_arguments else {} - statement_config = kwargs.pop("statement_config", None) - config = statement_config or self.statement_config - - sql_statement = self.prepare_statement( - operation.statement, operation.arguments, statement_config=config, kwargs=kwargs - ) - if not sql_statement.is_script and not sql_statement.is_many: - sql_text, prepared_parameters = self._compiled_sql(sql_statement, config) - prepared_parameters = cast("tuple[Any, ...] | dict[str, Any] | None", prepared_parameters) - normalized = NormalizedStackOperation( - operation=operation, statement=sql_statement, sql=sql_text, parameters=prepared_parameters + scope = self._stack_savepoint() if continue_on_error and caller_in_tx else nullcontext() + async with scope: + normalized: NormalizedStackOperation | None = None + if operation.method == "execute": + kwargs = dict(operation.keyword_arguments) if operation.keyword_arguments else {} + statement_config = kwargs.pop("statement_config", None) + config = statement_config or self.statement_config + + sql_statement = self.prepare_statement( + operation.statement, operation.arguments, statement_config=config, kwargs=kwargs ) - - if normalized is not None: - stack_result = await self._execute_stack_operation_prepared(normalized) - else: - result = await self._execute_stack_operation(operation) - stack_result = StackResult(result=result) + if not sql_statement.is_script and not sql_statement.is_many: + sql_text, prepared_parameters = self._compiled_sql(sql_statement, config) + prepared_parameters = cast("tuple[Any, ...] | dict[str, Any] | None", prepared_parameters) + normalized = NormalizedStackOperation( + operation=operation, + statement=sql_statement, + sql=sql_text, + parameters=prepared_parameters, + ) + + if normalized is not None: + stack_result = await self._execute_stack_operation_prepared(normalized) + else: + result = await self._execute_stack_operation(operation) + stack_result = StackResult(result=result) except Exception as exc: stack_error = StackExecutionError( index, @@ -626,16 +691,29 @@ async def _run_stack_operations( mode="continue-on-error" if continue_on_error else "fail-fast", ) if continue_on_error: - await self._rollback_failed_stack() + if not caller_in_tx: + await self._rollback_failed_stack() observer.record_operation_error(stack_error) results.append(StackResult.from_error(stack_error)) continue raise stack_error from exc results.append(stack_result) - if continue_on_error: + if continue_on_error and not caller_in_tx: await self._commit_stack_success() + @asynccontextmanager + async def _stack_savepoint(self) -> "AsyncIterator[None]": + name = validate_savepoint_name(f"sqlspec_stack_{uuid4().hex}") + await self.connection.execute(f"SAVEPOINT {name}") + try: + yield + except BaseException: + await self.connection.execute(f"ROLLBACK TO SAVEPOINT {name}") + raise + finally: + await self.connection.execute(f"RELEASE SAVEPOINT {name}") + async def _execute_stack_operation_prepared(self, normalized: "NormalizedStackOperation") -> StackResult: result = await self._execute_cached_statement(normalized.statement) if not self._is_transaction_pooler_mode(): @@ -644,9 +722,7 @@ async def _execute_stack_operation_prepared(self, normalized: "NormalizedStackOp def _is_transaction_pooler_mode(self) -> bool: pooler_mode = self.driver_features.get("pooler_mode", "session") - return pooler_mode in {"transaction", "statement"} or bool( - self.driver_features.get("pgbouncer") or self.driver_features.get("pgbouncer_mode") - ) + return pooler_mode in {"transaction", "statement"} def _should_set_local_columnar(self) -> bool: return bool(self.driver_features.get("enable_alloydb_columnar", False)) and self._is_transaction_pooler_mode() diff --git a/sqlspec/adapters/asyncpg/events/_hub.py b/sqlspec/adapters/asyncpg/events/_hub.py index 5cafccc63..0215f62de 100644 --- a/sqlspec/adapters/asyncpg/events/_hub.py +++ b/sqlspec/adapters/asyncpg/events/_hub.py @@ -7,11 +7,10 @@ notifications into one queue per consumer task. """ -# pyright: reportPrivateUsage=false - import asyncio import contextlib import logging +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, cast from weakref import WeakKeyDictionary @@ -161,7 +160,7 @@ def is_subscribed(self, channel: str) -> bool: def _get_consumer_queue(self, channel: str) -> "asyncio.Queue[str] | None": task = asyncio.current_task() - if task is None: # pragma: no cover + if task is None: msg = "AsyncpgListenerHub.dequeue requires an active asyncio task" raise RuntimeError(msg) queues = self._queues.get(channel) @@ -178,16 +177,9 @@ def _resolve_listener_config(self) -> "AsyncpgConfig": return self._listener_config raw_features = getattr(self._config, "driver_features", None) driver_features = raw_features if isinstance(raw_features, dict) else {} - pooler_mode = str( - driver_features.get("pooler_mode") - or ( - "transaction" - if driver_features.get("pgbouncer") or driver_features.get("pgbouncer_mode") - else "session" - ) - ).lower() + pooler_mode = str(driver_features.get("pooler_mode") or "session").lower() raw_direct_dsn = driver_features.get("direct_dsn") - direct_dsn = raw_direct_dsn if isinstance(raw_direct_dsn, str) and raw_direct_dsn else None + direct_dsn = raw_direct_dsn if isinstance(raw_direct_dsn, (str, Mapping)) and raw_direct_dsn else None if pooler_mode in {"transaction", "statement"} and not direct_dsn: logger.warning( "PostgreSQL LISTEN/NOTIFY cannot run under pooler_mode=%r without configuring driver_features['direct_dsn']", diff --git a/sqlspec/adapters/asyncpg/events/backend.py b/sqlspec/adapters/asyncpg/events/backend.py index 44f8eef83..9b3b04661 100644 --- a/sqlspec/adapters/asyncpg/events/backend.py +++ b/sqlspec/adapters/asyncpg/events/backend.py @@ -1,4 +1,3 @@ -# pyright: reportPrivateUsage=false """Native and hybrid PostgreSQL backends for EventChannel.""" import logging @@ -105,14 +104,15 @@ async def dequeue(self, channel: str, poll_interval: float) -> EventMessage | No return await self._dequeue_reconciled(channel) self._queue._reset_empty_poll_delay(channel) self._runtime.increment_metric("events.listener.wakeup") - batch_size = _extract_batch_size(payload) - is_new, should_drain = self._marker_state.register(channel, _extract_marker_id(payload), batch_size) + parsed_payload = _parse_notify_payload(payload) + batch_size = _extract_batch_size(parsed_payload) + is_new, should_drain = self._marker_state.register(channel, _extract_marker_id(parsed_payload), batch_size) if is_new and batch_size > 1: self._runtime.increment_metric("events.marker.coalesced", batch_size - 1) if not should_drain: self._runtime.increment_metric("events.marker.duplicate") continue - event_id = _extract_event_id(payload) + event_id = _extract_event_id(parsed_payload) if event_id is not None: event = await self._queue.dequeue_by_event_id(event_id) if event is not None: @@ -353,40 +353,50 @@ def reset(self) -> None: def _validate_listener_pool_capacity(config: "AsyncpgConfig") -> None: + raw_features = config.driver_features + driver_features = raw_features if isinstance(raw_features, dict) else {} + if driver_features.get("direct_dsn"): + return max_size = config.connection_config.get("max_size") if max_size is not None and max_size < _MIN_LISTENER_POOL_SIZE: msg = f"{type(config).__name__} native event listeners require pool max_size >= {_MIN_LISTENER_POOL_SIZE}" raise ImproperConfigurationError(msg) -def _extract_event_id(payload: "str | None") -> "str | None": - if not payload: +def _parse_notify_payload(payload: "dict[str, Any] | str | None") -> "dict[str, Any] | None": + if isinstance(payload, dict): + return payload + if not payload or not isinstance(payload, str): + return None + try: + raw = from_json(payload) + return raw if isinstance(raw, dict) else None + except (TypeError, ValueError): + return None + + +def _extract_event_id(payload: "dict[str, Any] | str | None") -> "str | None": + parsed = payload if isinstance(payload, dict) else _parse_notify_payload(payload) + if parsed is None: return None - raw = from_json(payload) - if isinstance(raw, dict): - event_id = raw.get("event_id") - return event_id if isinstance(event_id, str) else None - return None + event_id = parsed.get("event_id") + return event_id if isinstance(event_id, str) else None -def _extract_batch_size(payload: "str | None") -> int: - if not payload: +def _extract_batch_size(payload: "dict[str, Any] | str | None") -> int: + parsed = payload if isinstance(payload, dict) else _parse_notify_payload(payload) + if parsed is None: return 0 - raw = from_json(payload) - if isinstance(raw, dict): - batch_size = raw.get("batch_size") - return batch_size if isinstance(batch_size, int) else 0 - return 0 + batch_size = parsed.get("batch_size") + return batch_size if isinstance(batch_size, int) else 0 -def _extract_marker_id(payload: "str | None") -> "str | None": - if not payload: +def _extract_marker_id(payload: "dict[str, Any] | str | None") -> "str | None": + parsed = payload if isinstance(payload, dict) else _parse_notify_payload(payload) + if parsed is None: return None - raw = from_json(payload) - if isinstance(raw, dict): - marker_id = raw.get("marker_id", raw.get("event_id")) - return marker_id if isinstance(marker_id, str) else None - return None + marker_id = parsed.get("marker_id", parsed.get("event_id")) + return marker_id if isinstance(marker_id, str) else None def _record_dequeue_result(runtime: Any, event: "EventMessage | None") -> None: diff --git a/sqlspec/adapters/asyncpg/events/store.py b/sqlspec/adapters/asyncpg/events/store.py index 9c55518df..ecd9a7edc 100644 --- a/sqlspec/adapters/asyncpg/events/store.py +++ b/sqlspec/adapters/asyncpg/events/store.py @@ -3,6 +3,7 @@ from typing import NotRequired from sqlspec.adapters.asyncpg.config import AsyncpgConfig +from sqlspec.adapters.asyncpg.core import build_postgres_storage_params_sql from sqlspec.config import EventsConfig from sqlspec.extensions.events import BaseEventQueueStore @@ -53,27 +54,4 @@ def _column_types(self) -> "tuple[str, str, str]": def _table_clause(self) -> str: """Return explicitly configured PostgreSQL table tuning options.""" - return _postgres_table_options(self.settings) - - -def _postgres_table_options(settings: "dict[str, object]") -> str: - options = _postgres_option_values(settings) - return f" WITH ({', '.join(options)})" if options else "" - - -def _postgres_option_values(settings: "dict[str, object]") -> "list[str]": - options: list[str] = [] - fillfactor = settings.get("fillfactor") - if fillfactor is not None: - if not isinstance(fillfactor, int) or isinstance(fillfactor, bool) or fillfactor not in range(10, 101): - msg = "extension_config['events']['fillfactor'] must be an integer from 10 to 100" - raise ValueError(msg) - options.append(f"fillfactor = {fillfactor}") - for key in ("autovacuum_vacuum_scale_factor", "autovacuum_analyze_scale_factor"): - value = settings.get(key) - if value is not None: - if not isinstance(value, (int, float)) or isinstance(value, bool) or not 0 <= float(value) <= 1: - msg = f"extension_config['events']['{key}'] must be a number from 0 to 1" - raise ValueError(msg) - options.append(f"{key} = {float(value):g}") - return options + return build_postgres_storage_params_sql(self.settings, include_autovacuum=True, section="events") diff --git a/sqlspec/adapters/asyncpg/litestar/store.py b/sqlspec/adapters/asyncpg/litestar/store.py index cb8ff5863..30ee2cd74 100644 --- a/sqlspec/adapters/asyncpg/litestar/store.py +++ b/sqlspec/adapters/asyncpg/litestar/store.py @@ -1,10 +1,10 @@ """AsyncPG session store for Litestar integration.""" from datetime import UTC, datetime, timedelta -from typing import TYPE_CHECKING, Any, NotRequired, cast +from typing import TYPE_CHECKING, NotRequired from sqlspec.adapters.asyncpg._typing import asyncpg_module as asyncpg -from sqlspec.adapters.asyncpg.core import parse_status +from sqlspec.adapters.asyncpg.core import build_postgres_storage_params_sql, parse_status from sqlspec.config import LitestarConfig from sqlspec.extensions.litestar.store import BaseSQLSpecStore @@ -240,7 +240,14 @@ def _table_ddl(self) -> str: Returns: SQL statement to create the sessions table with proper indexes. """ - fillfactor, vacuum_scale, analyze_scale = _postgres_litestar_tuning(self._config) + tuning_clause = build_postgres_storage_params_sql( + self._config.extension_config.get("litestar", {}), + default_fillfactor=80, + default_vacuum_scale=0.05, + default_analyze_scale=0.02, + include_autovacuum=True, + section="litestar", + ) return f""" CREATE TABLE IF NOT EXISTS {self._table_name} ( session_id TEXT PRIMARY KEY, @@ -248,11 +255,7 @@ def _table_ddl(self) -> str: expires_at TIMESTAMPTZ, created_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP, updated_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP - ) WITH ( - fillfactor = {fillfactor}, - autovacuum_vacuum_scale_factor = {vacuum_scale:g}, - autovacuum_analyze_scale_factor = {analyze_scale:g} - ); + ){tuning_clause}; CREATE INDEX IF NOT EXISTS idx_{self._table_name}_expires_at ON {self._table_name}(expires_at) WHERE expires_at IS NOT NULL; @@ -265,19 +268,3 @@ def _drop_table_sql(self) -> "list[str]": List of SQL statements to drop indexes and table. """ return [f"DROP INDEX IF EXISTS idx_{self._table_name}_expires_at", f"DROP TABLE IF EXISTS {self._table_name}"] - - -def _postgres_litestar_tuning(config: Any) -> "tuple[int, float, float]": - settings = cast("dict[str, Any]", config.extension_config.get("litestar", {})) - fillfactor = settings.get("fillfactor", 80) - if not isinstance(fillfactor, int) or isinstance(fillfactor, bool) or fillfactor not in range(10, 101): - msg = "extension_config['litestar']['fillfactor'] must be an integer from 10 to 100" - raise ValueError(msg) - scales: list[float] = [] - for key, default in (("autovacuum_vacuum_scale_factor", 0.05), ("autovacuum_analyze_scale_factor", 0.02)): - value = settings.get(key, default) - if not isinstance(value, (int, float)) or isinstance(value, bool) or not 0 <= float(value) <= 1: - msg = f"extension_config['litestar']['{key}'] must be a number from 0 to 1" - raise ValueError(msg) - scales.append(float(value)) - return fillfactor, scales[0], scales[1] diff --git a/sqlspec/adapters/bigquery/_typing.py b/sqlspec/adapters/bigquery/_typing.py index c4c1ddbc5..c1029d3fe 100644 --- a/sqlspec/adapters/bigquery/_typing.py +++ b/sqlspec/adapters/bigquery/_typing.py @@ -12,6 +12,7 @@ from google.api_core import exceptions as bigquery_exceptions from google.api_core.client_info import ClientInfo as BigQueryClientInfo from google.api_core.client_options import ClientOptions as BigQueryClientOptions +from google.api_core.exceptions import GoogleAPIError from google.api_core.retry import Retry as BigQueryRetry from google.auth.credentials import Credentials as BigQueryCredentials from google.cloud.bigquery import ArrayQueryParameter, Client, ScalarQueryParameter, StructQueryParameter @@ -27,19 +28,17 @@ BigQueryConnection: TypeAlias = Client BigQueryParam: TypeAlias = ArrayQueryParameter | ScalarQueryParameter | StructQueryParameter -BigQueryStorageWriteModule: Any = import_optional("google.cloud.bigquery_storage_v1") -BigQueryStorageWriteTypes: Any = import_optional("google.cloud.bigquery_storage_v1.types") if TYPE_CHECKING: from collections.abc import Callable from types import TracebackType import google.cloud.bigquery_storage as bigquery_storage_read_module + import google.cloud.bigquery_storage_v1 as BigQueryStorageWriteModule # noqa: N812 + import google.cloud.bigquery_storage_v1.types as BigQueryStorageWriteTypes # noqa: N812 from sqlspec.adapters.bigquery.driver import BigQueryDriver from sqlspec.core import StatementConfig -else: - bigquery_storage_read_module = import_optional("google.cloud.bigquery_storage") __all__ = ( "BIGQUERY_DEFAULT_RETRY", @@ -58,6 +57,7 @@ "BigQuerySessionContext", "BigQueryStorageWriteModule", "BigQueryStorageWriteTypes", + "GoogleAPIError", "GoogleCloudError", "bigquery_exceptions", "bigquery_module", @@ -67,6 +67,17 @@ _DRIVER_TYPE: "type[BigQueryDriver] | None" = None +def __getattr__(name: str) -> Any: + if name == "BigQueryStorageWriteModule": + return import_optional("google.cloud.bigquery_storage_v1") + if name == "BigQueryStorageWriteTypes": + return import_optional("google.cloud.bigquery_storage_v1.types") + if name == "bigquery_storage_read_module": + return import_optional("google.cloud.bigquery_storage") + msg = f"module {__name__!r} has no attribute {name!r}" + raise AttributeError(msg) + + class BigQueryCursor: """BigQuery cursor with resource management.""" diff --git a/sqlspec/adapters/bigquery/adk/store.py b/sqlspec/adapters/bigquery/adk/store.py index 254e53acb..876fec8b8 100644 --- a/sqlspec/adapters/bigquery/adk/store.py +++ b/sqlspec/adapters/bigquery/adk/store.py @@ -15,10 +15,10 @@ from typing_extensions import TypedDict -from sqlspec.adapters.bigquery._typing import GoogleCloudError +from sqlspec.adapters.bigquery._typing import GoogleAPIError, GoogleCloudError from sqlspec.adapters.bigquery._typing import bigquery_module as bigquery from sqlspec.adapters.bigquery.config import BigQueryConfig -from sqlspec.adapters.bigquery.core import build_dml_rowcount, create_mapped_exception +from sqlspec.adapters.bigquery.core import build_dml_rowcount, create_mapped_exception, resolve_dataset_id from sqlspec.config import ADKConfig from sqlspec.extensions.adk import BaseSyncADKStore, StoredEvent, StoredSession, normalize_session_list_options from sqlspec.extensions.adk._config_utils import _adk_config_from_extension @@ -77,6 +77,8 @@ class BigQueryADKStore(BaseSyncADKStore[BigQueryConfig]): "_require_partition_filter", ) + _max_table_name_parts = 3 + def __init__(self, config: BigQueryConfig) -> None: """Initialize the BigQuery ADK store. @@ -238,7 +240,7 @@ def _run_query(self, sql: str, parameters: "Iterable[Any] | None" = None) -> "li try: job = client.query(sql, job_config=job_config) rows = [dict(row) for row in job.result()] - except GoogleCloudError as exc: + except (GoogleCloudError, GoogleAPIError) as exc: raise create_mapped_exception(exc) from exc self._last_dml_affected_rows = build_dml_rowcount(job, 0) return rows @@ -672,23 +674,7 @@ def _decode_json(value: Any, *, deserializer: "Callable[[str], Any]" = from_json def _resolve_config_dataset_id(config: BigQueryConfig) -> str | None: - dataset_id = config.connection_config.get("dataset_id") - if isinstance(dataset_id, str) and dataset_id: - return dataset_id - default_job_config = config.connection_config.get("default_query_job_config") or config.driver_features.get( - "default_query_job_config" - ) - if default_job_config is None: - return None - default_dataset = getattr(default_job_config, "default_dataset", None) - if default_dataset is None: - return None - if isinstance(default_dataset, str): - return default_dataset or None - resolved_id = getattr(default_dataset, "dataset_id", None) - if resolved_id: - return str(resolved_id) - return str(default_dataset) + return resolve_dataset_id(config) def _session_record_from_row( diff --git a/sqlspec/adapters/bigquery/config.py b/sqlspec/adapters/bigquery/config.py index 0cb60d6f4..7a1d6338c 100644 --- a/sqlspec/adapters/bigquery/config.py +++ b/sqlspec/adapters/bigquery/config.py @@ -103,10 +103,14 @@ class BigQueryDriverFeatures(TypedDict): connection_instance: NotRequired["BigQueryConnection"] """Pre-created ``google.cloud.bigquery.Client`` reused without transferring lifecycle ownership.""" + default_query_job_config: NotRequired[QueryJobConfig] + """Default ``QueryJobConfig`` for a directly constructed driver; config-managed drivers use the client's default.""" on_connection_create: NotRequired["Callable[[Any], None]"] """Callback invoked with the ``bigquery.Client`` when instantiated by the config.""" json_serializer: NotRequired["Callable[[Any], str]"] """Serializer converting dict/list parameters to JSON strings (defaults to ``sqlspec.utils.serializers.to_json``).""" + json_deserializer: NotRequired["Callable[[str], Any]"] + """Deserializer converting JSON strings to native dict/list objects.""" enable_uuid_conversion: NotRequired[bool] """Automatically coerce ``UUID`` strings and parameters (defaults to ``True``).""" enable_events: NotRequired[bool] @@ -117,7 +121,7 @@ class BigQueryDriverFeatures(TypedDict): """Total seconds to retry transient query failures (defaults to ``60.0``; ``<= 0`` disables retries).""" job_result_timeout: NotRequired[float] """Maximum seconds to poll ``QueryJob.result()`` before raising a timeout error.""" - use_query_and_wait: NotRequired[bool] + enable_query_and_wait: NotRequired[bool] """Execute single-statement queries via low-latency ``Client.query_and_wait()`` outside transactions (defaults to ``True``).""" request_timeout: NotRequired[float] """Per-request HTTP transport timeout in seconds for job submission calls (defaults to ``job_result_timeout`` or ``120.0``).""" @@ -232,6 +236,10 @@ def __init__( self.connection_config = build_connection_config(connection_config) + if "json_deserializer" in kwargs: + driver_features = dict(driver_features) if driver_features else {} + driver_features["json_deserializer"] = kwargs.pop("json_deserializer") + statement_config = statement_config or default_statement_config statement_config, driver_features = apply_driver_features(statement_config, driver_features) user_connection_hook = cast( @@ -435,7 +443,8 @@ def _provide_storage_write_client(self, connection: "BigQueryConnection") -> Any client_options = self.connection_config.get("client_options") if client_options is not None and client_options is not Empty: client_kwargs["client_options"] = client_options - self._storage_write_client = BigQueryStorageWriteModule.BigQueryWriteClient(**client_kwargs) + write_client_cls = cast("Any", BigQueryStorageWriteModule).BigQueryWriteClient + self._storage_write_client = write_client_cls(**client_kwargs) return self._storage_write_client def _resolve_default_dataset(self) -> "str | None": diff --git a/sqlspec/adapters/bigquery/core.py b/sqlspec/adapters/bigquery/core.py index 2f632a681..f6f62ce86 100644 --- a/sqlspec/adapters/bigquery/core.py +++ b/sqlspec/adapters/bigquery/core.py @@ -7,7 +7,7 @@ import ipaddress from decimal import Decimal from enum import Enum -from typing import TYPE_CHECKING, Any, Protocol, cast, get_args +from typing import TYPE_CHECKING, Any, Final, Protocol, cast, get_args from urllib.parse import urlparse import sqlglot @@ -37,6 +37,7 @@ OperationalError, OperationCancelledError, PermissionDeniedError, + QueryTimeoutError, SQLParsingError, SQLSpecError, StorageCapabilityError, @@ -44,8 +45,9 @@ UniqueViolationError, _classify_timeout_or_cancellation, ) +from sqlspec.storage._utils import import_pyarrow, import_pyarrow_parquet +from sqlspec.typing import PYARROW_INSTALLED from sqlspec.utils.logging import get_logger -from sqlspec.utils.module_loader import ensure_pyarrow, import_optional from sqlspec.utils.serializers import from_json, to_json from sqlspec.utils.type_converters import build_uuid_coercions from sqlspec.utils.type_guards import has_errors @@ -60,9 +62,6 @@ BigQueryLoadFormat = Literal["jsonl", "json", "parquet", "arrow-ipc", "csv", "avro", "orc"] -_pa: Any = import_optional("pyarrow") -_pq: Any = import_optional("pyarrow.parquet") - __all__ = ( "BigQueryStreamSource", "apply_driver_features", @@ -86,6 +85,7 @@ "is_simple_insert", "normalize_script_rowcount", "resolve_column_names", + "resolve_dataset_id", "resolve_export_format", "run_query_and_wait", "run_query_job", @@ -101,6 +101,7 @@ HTTP_SERVER_ERROR = 500 COLUMN_CACHE_MAX_SIZE = 256 _CONNECTION_NAME_PARTS = 3 +none_type: Final = type(None) DEFAULT_REQUEST_TIMEOUT = 120.0 @@ -181,21 +182,23 @@ def try_bulk_insert( if _has_synthetic_positional_keys(remapped_parameters): return None + if not PYARROW_INSTALLED: + logger.debug("pyarrow not available, falling back to literal inlining") + return None + try: - ensure_pyarrow() - arrow_table = _pa.Table.from_pylist(remapped_parameters) + pa = import_pyarrow() + pq = import_pyarrow_parquet() + arrow_table = pa.Table.from_pylist(remapped_parameters) buffer = io.BytesIO() - _pq.write_table(arrow_table, buffer) + pq.write_table(arrow_table, buffer) buffer.seek(0) job_config = build_load_job_config("parquet", overwrite=False) job = connection.load_table_from_file(buffer, table_name, job_config=job_config, timeout=result_timeout) job.result(timeout=result_timeout) return len(remapped_parameters) - except ImportError: - logger.debug("pyarrow not available, falling back to literal inlining") - return None except Exception as exc: logger.debug("Bulk insert failed, falling back to literal inlining: %s", exc) return None @@ -690,7 +693,7 @@ def build_profile() -> "DriverParameterProfile": Decimal: _identity, dict: _identity, list: _identity, - type(None): _return_none, + none_type: _return_none, **build_uuid_coercions(), }, default_ast_transformer=build_null_pruning_transform(dialect="bigquery"), @@ -751,7 +754,7 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpec Returns: A SQLSpec exception that wraps the original error """ - del logger + _ = logger try: status_code = error.code except AttributeError: @@ -767,6 +770,9 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpec if isinstance(error, api_exceptions.Cancelled): return _create_bigquery_error(error, status_code, OperationCancelledError, "query cancelled") + if isinstance(error, api_exceptions.RetryError): + return _create_bigquery_error(error, status_code, QueryTimeoutError, "query retry timeout expired") + if error_class := _classify_timeout_or_cancellation(error_msg): return _create_bigquery_error(error, status_code, error_class, "query terminated") @@ -1378,3 +1384,55 @@ def _create_bigquery_error( driver_profile = build_profile() default_statement_config = build_statement_config() + + +def _positive_int_or_none(value: object, key: str) -> "int | None": + """Validate that an extension config setting is a positive integer or None.""" + if value is None: + return None + if not isinstance(value, int) or isinstance(value, bool) or value <= 0: + msg = f"'{key}' must be a positive integer" + raise ValueError(msg) + return value + + +def _bigquery_partition_clauses( + column: str, *, enabled: bool, expiration_days: "int | None", require_filter: bool +) -> "tuple[str, str]": + """Build BigQuery PARTITION BY and OPTIONS clauses for table DDL.""" + if not enabled: + return "", "" + options: list[str] = [] + if require_filter: + options.append("require_partition_filter = TRUE") + if expiration_days is not None: + options.append(f"partition_expiration_days = {expiration_days}") + options_clause = f" OPTIONS({', '.join(options)})" if options else "" + return f" PARTITION BY DATE({column})", options_clause + + +def resolve_dataset_id(config: Any) -> "str | None": + """Resolve the dataset identifier from connection configuration or default job config.""" + connection_config = getattr(config, "connection_config", None) + if isinstance(connection_config, dict): + dataset_id = connection_config.get("dataset_id") + if isinstance(dataset_id, str) and dataset_id: + return dataset_id + default_job_config = getattr(config, "_default_query_job_config", None) + if default_job_config is None and isinstance(connection_config, dict): + default_job_config = connection_config.get("default_query_job_config") + if default_job_config is None: + driver_features = getattr(config, "driver_features", None) + if isinstance(driver_features, dict): + default_job_config = driver_features.get("default_query_job_config") + if default_job_config is None: + return None + default_dataset = getattr(default_job_config, "default_dataset", None) + if default_dataset is None: + return None + if isinstance(default_dataset, str): + return default_dataset or None + resolved_id = getattr(default_dataset, "dataset_id", None) + if resolved_id: + return str(resolved_id) + return str(default_dataset) diff --git a/sqlspec/adapters/bigquery/driver.py b/sqlspec/adapters/bigquery/driver.py index f89bdfadb..c37b5d942 100644 --- a/sqlspec/adapters/bigquery/driver.py +++ b/sqlspec/adapters/bigquery/driver.py @@ -5,6 +5,7 @@ """ import io +import sys from collections.abc import Mapping from contextlib import suppress from itertools import chain @@ -22,6 +23,7 @@ BigQuerySessionContext, BigQueryStorageWriteModule, BigQueryStorageWriteTypes, + GoogleAPIError, GoogleCloudError, ) from sqlspec.adapters.bigquery.core import ( @@ -60,8 +62,9 @@ ) from sqlspec.driver import BaseSyncExceptionHandler, ExecutionResult, SyncDriverAdapterBase, SyncRowStream from sqlspec.exceptions import ImproperConfigurationError, StorageOperationFailedError +from sqlspec.storage._utils import import_pyarrow, import_pyarrow_parquet from sqlspec.utils.logging import get_logger -from sqlspec.utils.module_loader import ensure_pyarrow, import_optional +from sqlspec.utils.module_loader import ensure_pyarrow from sqlspec.utils.serializers import to_json from sqlspec.utils.text import split_qualified_identifier @@ -83,8 +86,6 @@ logger = get_logger(__name__) _DATASET_TABLE_PARTS = 2 _PROJECT_DATASET_TABLE_PARTS = 3 -_pa: Any = import_optional("pyarrow") -_pq: Any = import_optional("pyarrow.parquet") class BigQueryExceptionHandler(BaseSyncExceptionHandler): @@ -99,7 +100,7 @@ class BigQueryExceptionHandler(BaseSyncExceptionHandler): def _handle_exception(self, exc_type: "type[BaseException] | None", exc_val: "BaseException") -> bool: if exc_type is None: return False - if issubclass(exc_type, GoogleCloudError): + if issubclass(exc_type, (GoogleCloudError, GoogleAPIError)): self.pending_exception = create_mapped_exception(exc_val) return True return False @@ -116,6 +117,7 @@ class BigQueryDriver(SyncDriverAdapterBase): "_column_name_cache", "_data_dictionary", "_default_query_job_config", + "_enable_query_and_wait", "_enable_storage_write_api", "_in_transaction", "_job_result_kwargs_defaults", @@ -126,7 +128,6 @@ class BigQueryDriver(SyncDriverAdapterBase): "_literal_inliner", "_request_timeout", "_session_id", - "_use_query_and_wait", ) dialect = "bigquery" @@ -164,7 +165,7 @@ def __init__( self._job_retry: Retry | None = build_retry(self._job_retry_deadline) if self._job_retry_deadline > 0 else None self._job_result_timeout: float | object = features.get("job_result_timeout", POLLING_DEFAULT_VALUE) self._request_timeout = self._resolve_request_timeout(features) - self._use_query_and_wait = bool(features.get("use_query_and_wait", True)) + self._enable_query_and_wait = bool(features.get("enable_query_and_wait", True)) self._enable_storage_write_api = bool(features.get("enable_storage_write_api", False)) self._in_transaction: bool = False self._session_id: str | None = None @@ -318,7 +319,8 @@ def begin(self) -> None: if not hasattr(self.connection, "query"): return exc_handler = self.handle_database_exceptions() - with exc_handler: + + def _begin_op() -> None: job_config = ( BigQueryQueryJobConfig(create_session=True) if self._session_id is None @@ -344,8 +346,13 @@ def begin(self) -> None: session_id = getattr(session_info, "session_id", None) if isinstance(session_id, str) and session_id: self._session_id = session_id - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + + if sys.exc_info()[1] is None: + with exc_handler: + _begin_op() + else: + self._run_with_exception_handler(exc_handler, _begin_op) + self._check_pending_exception(exc_handler) def commit(self) -> None: """Commit the active BigQuery session transaction.""" @@ -354,13 +361,19 @@ def commit(self) -> None: if not was_in_tx or not hasattr(self.connection, "query"): return exc_handler = self.handle_database_exceptions() - with exc_handler: + + def _commit_op() -> None: job = self._run_query_job(self.connection, "COMMIT TRANSACTION", None) result_fn = getattr(job, "result", None) if callable(result_fn): result_fn(job_retry=self._job_retry, timeout=self._job_result_timeout) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + + if sys.exc_info()[1] is None: + with exc_handler: + _commit_op() + else: + self._run_with_exception_handler(exc_handler, _commit_op) + self._check_pending_exception(exc_handler) def rollback(self) -> None: """Rollback the active BigQuery session transaction.""" @@ -369,13 +382,19 @@ def rollback(self) -> None: if not was_in_tx or not hasattr(self.connection, "query"): return exc_handler = self.handle_database_exceptions() - with exc_handler: + + def _rollback_op() -> None: job = self._run_query_job(self.connection, "ROLLBACK TRANSACTION", None) result_fn = getattr(job, "result", None) if callable(result_fn): result_fn(job_retry=self._job_retry, timeout=self._job_result_timeout) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + + if sys.exc_info()[1] is None: + with exc_handler: + _rollback_op() + else: + self._run_with_exception_handler(exc_handler, _rollback_op) + self._check_pending_exception(exc_handler) def dry_run( self, @@ -394,7 +413,8 @@ def dry_run( job_config.connection_properties = [BigQueryConnectionProperty("session_id", self._session_id)] exc_handler = self.handle_database_exceptions() total_bytes = 0 - with exc_handler: + + def _dry_run_op() -> int: job = run_query_job( self.connection, sql, @@ -406,9 +426,16 @@ def dry_run( timeout=self._job_request_timeout(), job_retry=self._job_retry, ) - total_bytes = int(getattr(job, "total_bytes_processed", 0) or 0) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + return int(getattr(job, "total_bytes_processed", 0) or 0) + + if sys.exc_info()[1] is None: + with exc_handler: + total_bytes = _dry_run_op() + else: + res = self._run_with_exception_handler(exc_handler, _dry_run_op) + if res is not None: + total_bytes = res + self._check_pending_exception(exc_handler) return total_bytes def create_savepoint(self, name: str) -> None: @@ -523,7 +550,8 @@ def select_to_arrow( exc_handler = self.handle_database_exceptions() streaming_result: ArrowResult | None = None - with exc_handler: + + def _streaming_op() -> Any: query_job = self._run_query_job(self.connection, sql, driver_params) row_iterator = query_job.result( page_size=batch_size, job_retry=self._job_retry, timeout=self._job_result_timeout @@ -535,7 +563,7 @@ def select_to_arrow( else (query_job.schema or getattr(row_iterator, "schema", None)) ) arrow_reader = _bigquery_arrow_reader_from_iterable(arrow_batches, schema=fallback_schema) - streaming_result = build_arrow_result_from_reader( + return build_arrow_result_from_reader( prepared_statement, arrow_reader, return_format=return_format, @@ -543,8 +571,12 @@ def select_to_arrow( arrow_schema=arrow_schema, ) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + if sys.exc_info()[1] is None: + with exc_handler: + streaming_result = _streaming_op() + else: + streaming_result = self._run_with_exception_handler(exc_handler, _streaming_op) + self._check_pending_exception(exc_handler) if streaming_result is None: msg = "Unreachable" raise RuntimeError(msg) @@ -596,8 +628,7 @@ def select_to_arrow( arrow_schema=arrow_schema, ) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + self._check_pending_exception(exc_handler) if arrow_result is None: msg = "Unreachable" @@ -661,7 +692,7 @@ def select_to_storage( raise if handler.pending_exception is not None: self.observability.end_storage_span(span, error=handler.pending_exception) - raise handler.pending_exception from None + self._check_pending_exception(handler) native_telemetry = self.observability.annotate_storage_telemetry(native_telemetry) self.observability.end_storage_span(span, telemetry=native_telemetry) self._attach_partition_telemetry(native_telemetry, partitioner) @@ -732,8 +763,7 @@ def load_from_records( rows, qualified_table, job_config=job_config, timeout=self._job_request_timeout() ) job.result(timeout=self._job_request_timeout()) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + self._check_pending_exception(exc_handler) telemetry_payload = build_load_job_telemetry(job, qualified_table, format_label="jsonl") return self._storage_job(telemetry_payload) @@ -795,20 +825,23 @@ def load_from_arrow( ensure_pyarrow() qualified_table = self._qualify_table_name(table) - if self._enable_storage_write_api and not overwrite and not uses_local_bigquery_endpoint(self.connection): - try: - telemetry_payload = self._load_arrow_via_storage_write_api(qualified_table, arrow_table) - except ImportError as exc: - logger.warning("Storage Write API unavailable, falling back to Parquet load job: %s", exc) - else: - if telemetry: - telemetry_payload.setdefault("extra", {}) - telemetry_payload["extra"]["arrow_rows"] = telemetry.get("rows_processed") - self._attach_partition_telemetry(telemetry_payload, partitioner) - return self._storage_job(telemetry_payload) - + if ( + self._enable_storage_write_api + and not overwrite + and not uses_local_bigquery_endpoint(self.connection) + and BigQueryStorageWriteModule is not None + and BigQueryStorageWriteTypes is not None + ): + telemetry_payload = self._load_arrow_via_storage_write_api(qualified_table, arrow_table) + if telemetry: + telemetry_payload.setdefault("extra", {}) + telemetry_payload["extra"]["arrow_rows"] = telemetry.get("rows_processed") + self._attach_partition_telemetry(telemetry_payload, partitioner) + return self._storage_job(telemetry_payload) + + pq = import_pyarrow_parquet() buffer = io.BytesIO() - _pq.write_table(arrow_table, buffer) + pq.write_table(arrow_table, buffer) buffer.seek(0) job_config = build_load_job_config("parquet", overwrite) exc_handler = self.handle_database_exceptions() @@ -818,8 +851,7 @@ def load_from_arrow( buffer, qualified_table, job_config=job_config, timeout=self._job_request_timeout() ) job.result(timeout=self._job_request_timeout()) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + self._check_pending_exception(exc_handler) telemetry_payload = build_load_job_telemetry(job, qualified_table, format_label="parquet") if telemetry: telemetry_payload.setdefault("extra", {}) @@ -878,8 +910,7 @@ def load_from_storage( buffer, qualified_table, job_config=job_config, timeout=self._job_request_timeout() ) job.result(timeout=self._job_request_timeout()) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + self._check_pending_exception(exc_handler) telemetry_payload = build_load_job_telemetry(job, qualified_table, format_label=file_format) self._attach_partition_telemetry(telemetry_payload, partitioner) return self._storage_job(telemetry_payload, source_telemetry) @@ -921,7 +952,7 @@ def _qualify_table_name(self, table: str) -> str: def _can_use_query_and_wait(self, statement: "SQL", sql: str) -> bool: """Use the interactive fast path only for compatible query jobs.""" - if not self._use_query_and_wait or not hasattr(self.connection, "query_and_wait"): + if not self._enable_query_and_wait or not hasattr(self.connection, "query_and_wait"): return False if self._in_transaction or self._session_id is not None: return False @@ -1012,7 +1043,8 @@ def _load_arrow_via_storage_write_api( client = provider(self.connection) else: credentials = getattr(self.connection, "_credentials", None) - client = BigQueryStorageWriteModule.BigQueryWriteClient(credentials=credentials) + write_client_cls = cast("Any", BigQueryStorageWriteModule).BigQueryWriteClient + client = write_client_cls(credentials=credentials) parent = f"projects/{project}/datasets/{dataset}/tables/{table_name}" write_stream_type_enum = types.WriteStream.Type.COMMITTED if is_committed else types.WriteStream.Type.PENDING write_stream = client.create_write_stream( @@ -1086,20 +1118,42 @@ def _bigquery_arrow_reader_from_iterable( batches: "Iterable[ArrowRecordBatch]", schema: Any = None ) -> "ArrowRecordBatchReader": ensure_pyarrow() + pa = import_pyarrow() iterator = iter(batches) try: first_batch = next(iterator) except StopIteration: - if isinstance(schema, _pa.Schema): + if isinstance(schema, pa.Schema): empty_schema = schema elif isinstance(schema, (list, tuple)): - empty_schema = _pa.schema([_pa.field(getattr(field, "name", str(field)), _pa.null()) for field in schema]) + type_map = { + "INTEGER": pa.int64(), + "INT64": pa.int64(), + "FLOAT": pa.float64(), + "FLOAT64": pa.float64(), + "BOOLEAN": pa.bool_(), + "BOOL": pa.bool_(), + "STRING": pa.string(), + "BYTES": pa.binary(), + "DATE": pa.date32(), + "DATETIME": pa.timestamp("us"), + "TIMESTAMP": pa.timestamp("us", tz="UTC"), + "NUMERIC": pa.decimal128(38, 9), + "BIGNUMERIC": pa.decimal256(76, 38), + } + fields = [ + pa.field( + getattr(field, "name", str(field)), + type_map.get(str(getattr(field, "field_type", "")).upper(), pa.null()), + ) + for field in schema + ] + empty_schema = pa.schema(fields) else: - empty_schema = _pa.schema([]) - return cast("ArrowRecordBatchReader", _pa.RecordBatchReader.from_batches(empty_schema, ())) + empty_schema = pa.schema([]) + return cast("ArrowRecordBatchReader", pa.RecordBatchReader.from_batches(empty_schema, ())) return cast( - "ArrowRecordBatchReader", - _pa.RecordBatchReader.from_batches(first_batch.schema, chain((first_batch,), iterator)), + "ArrowRecordBatchReader", pa.RecordBatchReader.from_batches(first_batch.schema, chain((first_batch,), iterator)) ) diff --git a/sqlspec/adapters/bigquery/events/store.py b/sqlspec/adapters/bigquery/events/store.py index dc9fd9b3b..bfb0ff8c9 100644 --- a/sqlspec/adapters/bigquery/events/store.py +++ b/sqlspec/adapters/bigquery/events/store.py @@ -15,6 +15,7 @@ from typing import NotRequired from sqlspec.adapters.bigquery.config import BigQueryConfig +from sqlspec.adapters.bigquery.core import _bigquery_partition_clauses, _positive_int_or_none, resolve_dataset_id from sqlspec.config import EventsConfig from sqlspec.extensions.events import BaseEventQueueStore @@ -88,7 +89,7 @@ def _column_types(self) -> "tuple[str, str, str]": def _string_type(self, length: int) -> str: """Return BigQuery STRING type (length is ignored).""" - del length + _ = length return "STRING" def _integer_type(self) -> str: @@ -136,43 +137,4 @@ def _index_ddl(self) -> str | None: def _resolve_events_dataset_id(config: BigQueryConfig) -> "str | None": - dataset_id = config.connection_config.get("dataset_id") - if isinstance(dataset_id, str) and dataset_id: - return dataset_id - default_job_config = config.connection_config.get("default_query_job_config") or config.driver_features.get( - "default_query_job_config" - ) - if default_job_config is None: - return None - default_dataset = getattr(default_job_config, "default_dataset", None) - if default_dataset is None: - return None - if isinstance(default_dataset, str): - return default_dataset or None - resolved_id = getattr(default_dataset, "dataset_id", None) - if resolved_id: - return str(resolved_id) - return str(default_dataset) - - -def _positive_int_or_none(value: object, key: str) -> "int | None": - if value is None: - return None - if not isinstance(value, int) or isinstance(value, bool) or value <= 0: - msg = f"extension_config['events']['{key}'] must be a positive integer" - raise ValueError(msg) - return value - - -def _bigquery_partition_clauses( - column: str, *, enabled: bool, expiration_days: "int | None", require_filter: bool -) -> "tuple[str, str]": - if not enabled: - return "", "" - options: list[str] = [] - if require_filter: - options.append("require_partition_filter = TRUE") - if expiration_days is not None: - options.append(f"partition_expiration_days = {expiration_days}") - options_clause = f" OPTIONS({', '.join(options)})" if options else "" - return f" PARTITION BY DATE({column})", options_clause + return resolve_dataset_id(config) diff --git a/sqlspec/adapters/bigquery/litestar/store.py b/sqlspec/adapters/bigquery/litestar/store.py index c26aca41f..f6ddac2e6 100644 --- a/sqlspec/adapters/bigquery/litestar/store.py +++ b/sqlspec/adapters/bigquery/litestar/store.py @@ -3,6 +3,7 @@ from datetime import UTC, datetime, timedelta from typing import TYPE_CHECKING, Any, NotRequired, cast +from sqlspec.adapters.bigquery.core import _bigquery_partition_clauses, _positive_int_or_none, resolve_dataset_id from sqlspec.config import LitestarConfig from sqlspec.extensions.litestar.store import BaseSQLSpecStore from sqlspec.utils.sync_tools import async_ @@ -345,43 +346,4 @@ def _delete_expired(self) -> int: def _resolve_litestar_dataset_id(config: "BigQueryConfig") -> "str | None": - dataset_id = config.connection_config.get("dataset_id") - if isinstance(dataset_id, str) and dataset_id: - return dataset_id - default_job_config = config.connection_config.get("default_query_job_config") or config.driver_features.get( - "default_query_job_config" - ) - if default_job_config is None: - return None - default_dataset = getattr(default_job_config, "default_dataset", None) - if default_dataset is None: - return None - if isinstance(default_dataset, str): - return default_dataset or None - resolved_id = getattr(default_dataset, "dataset_id", None) - if resolved_id: - return str(resolved_id) - return str(default_dataset) - - -def _positive_int_or_none(value: object, key: str) -> "int | None": - if value is None: - return None - if not isinstance(value, int) or isinstance(value, bool) or value <= 0: - msg = f"extension_config['litestar']['{key}'] must be a positive integer" - raise ValueError(msg) - return value - - -def _bigquery_partition_clauses( - column: str, *, enabled: bool, expiration_days: "int | None", require_filter: bool -) -> "tuple[str, str]": - if not enabled: - return "", "" - options: list[str] = [] - if require_filter: - options.append("require_partition_filter = TRUE") - if expiration_days is not None: - options.append(f"partition_expiration_days = {expiration_days}") - options_clause = f"\n OPTIONS({', '.join(options)})" if options else "" - return f"PARTITION BY DATE({column})\n ", options_clause + return resolve_dataset_id(config) diff --git a/sqlspec/adapters/cockroach_asyncpg/adk/store.py b/sqlspec/adapters/cockroach_asyncpg/adk/store.py index f8e3a0876..3be8d81da 100644 --- a/sqlspec/adapters/cockroach_asyncpg/adk/store.py +++ b/sqlspec/adapters/cockroach_asyncpg/adk/store.py @@ -72,6 +72,8 @@ class CockroachAsyncpgADKStore(BaseAsyncADKStore["CockroachAsyncpgConfig"]): __slots__ = () + _max_table_name_parts = 2 + def __init__(self, config: "CockroachAsyncpgConfig") -> None: super().__init__(config) @@ -587,6 +589,8 @@ class CockroachAsyncpgADKMemoryStore(BaseAsyncADKMemoryStore["CockroachAsyncpgCo __slots__ = () + _max_table_name_parts = 2 + def __init__(self, config: "CockroachAsyncpgConfig") -> None: super().__init__(config) @@ -609,66 +613,81 @@ async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: " if not entries: return 0 + ids = [entry["id"] for entry in entries] + session_ids = [entry["session_id"] for entry in entries] + app_names = [entry["app_name"] for entry in entries] + user_ids = [entry["user_id"] for entry in entries] + scopes = [entry.get("scope", "user") for entry in entries] + event_ids = [entry["event_id"] for entry in entries] + authors = [entry["author"] for entry in entries] + timestamps = [entry["timestamp"] for entry in entries] + content_jsons = [entry["content_json"] for entry in entries] + content_texts = [entry["content_text"] for entry in entries] + metadata_jsons = [entry["metadata_json"] for entry in entries] + inserted_ats = [entry["inserted_at"] for entry in entries] + if self._owner_id_column_name: sql = f""" INSERT INTO {self._memory_table} (id, session_id, app_name, user_id, scope, 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, $13) + SELECT id, session_id, app_name, user_id, scope, event_id, author, + $13, timestamp, content_json, content_text, metadata_json, inserted_at + FROM unnest( + $1::varchar[], $2::varchar[], $3::varchar[], $4::varchar[], + $5::varchar[], $6::varchar[], $7::varchar[], + $8::timestamptz[], $9::jsonb[], $10::text[], $11::jsonb[], $12::timestamptz[] + ) AS batch(id, session_id, app_name, user_id, scope, event_id, author, + timestamp, content_json, content_text, metadata_json, inserted_at) ON CONFLICT (event_id) DO NOTHING """ + params: tuple[Any, ...] = ( + ids, + session_ids, + app_names, + user_ids, + scopes, + event_ids, + authors, + timestamps, + content_jsons, + content_texts, + metadata_jsons, + inserted_ats, + owner_id, + ) else: sql = f""" INSERT INTO {self._memory_table} (id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json, inserted_at) - VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12) + SELECT * FROM unnest( + $1::varchar[], $2::varchar[], $3::varchar[], $4::varchar[], + $5::varchar[], $6::varchar[], $7::varchar[], + $8::timestamptz[], $9::jsonb[], $10::text[], $11::jsonb[], $12::timestamptz[] + ) ON CONFLICT (event_id) DO NOTHING """ + params = ( + ids, + session_ids, + app_names, + user_ids, + scopes, + event_ids, + authors, + timestamps, + content_jsons, + content_texts, + metadata_jsons, + inserted_ats, + ) - inserted_count = 0 async with self._config.provide_connection() as conn, conn.transaction(): - for entry in entries: - if self._owner_id_column_name: - result = await conn.execute( - sql, - entry["id"], - entry["session_id"], - entry["app_name"], - entry["user_id"], - entry.get("scope", "user"), - entry["event_id"], - entry["author"], - owner_id, - entry["timestamp"], - entry["content_json"], - entry["content_text"], - entry["metadata_json"], - entry["inserted_at"], - ) - else: - result = await conn.execute( - sql, - entry["id"], - entry["session_id"], - entry["app_name"], - entry["user_id"], - entry.get("scope", "user"), - entry["event_id"], - entry["author"], - entry["timestamp"], - entry["content_json"], - entry["content_text"], - entry["metadata_json"], - entry["inserted_at"], - ) - - parts = result.split() if result else [] - if parts and parts[-1].isdigit() and int(parts[-1]) > 0: - inserted_count += 1 - - return inserted_count + result = await conn.execute(sql, *params) + parts = result.split() if result else [] + return int(parts[-1]) if parts and parts[-1].isdigit() else 0 async def search_entries( self, @@ -691,20 +710,18 @@ async def search_entries( if scope_filter == "all": where_scope = "app_name = $1 AND ((scope = 'user' AND user_id = $2) OR scope = 'app')" scope_params: tuple[Any, ...] = (app_name, user_id) - p_lim = "$4" - p_pat = "$3" elif scope_filter == "user": where_scope = "app_name = $1 AND scope = 'user' AND user_id = $2" scope_params = (app_name, user_id) - p_lim = "$4" - p_pat = "$3" else: where_scope = "app_name = $1 AND scope = 'app'" scope_params = (app_name,) - p_lim = "$3" - p_pat = "$2" + + scope_count = len(scope_params) if self._use_fts: + p_pat = f"${scope_count + 1}" + p_lim = f"${scope_count + 2}" sql = f""" SELECT * FROM {self._memory_table} WHERE {where_scope} @@ -712,17 +729,17 @@ async def search_entries( ORDER BY timestamp DESC LIMIT {p_lim} """ - search_param = query + params = (*scope_params, query, effective_limit) else: + p_pat = f"${scope_count + 1}" + p_lim = f"${scope_count + 2}" sql = f""" SELECT * FROM {self._memory_table} WHERE {where_scope} AND content_text ILIKE {p_pat} ORDER BY timestamp DESC LIMIT {p_lim} """ - search_param = f"%{query}%" - - params = (*scope_params, search_param, effective_limit) + params = (*scope_params, f"%{query}%", effective_limit) try: async with self._config.provide_connection() as conn: @@ -857,9 +874,12 @@ def _cockroach_hash_shard_clause(adk_config: CockroachAsyncpgADKConfig) -> str: if not adk_config.get("enable_hash_sharded_indexes", False): return "" bucket_count = adk_config.get("hash_shard_bucket_count") - if isinstance(bucket_count, int) and not isinstance(bucket_count, bool) and bucket_count > 0: - return f" USING HASH WITH (bucket_count = {bucket_count})" - return " USING HASH" + if bucket_count is None: + return " USING HASH" + if not isinstance(bucket_count, int) or isinstance(bucket_count, bool) or bucket_count <= 0: + msg = "extension_config['adk']['hash_shard_bucket_count'] must be a positive integer" + raise ValueError(msg) + return f" USING HASH WITH (bucket_count = {bucket_count})" def _cockroach_storing_clause(adk_config: CockroachAsyncpgADKConfig, columns: tuple[str, ...]) -> str: diff --git a/sqlspec/adapters/cockroach_asyncpg/config.py b/sqlspec/adapters/cockroach_asyncpg/config.py index e5bad024f..1680b951b 100644 --- a/sqlspec/adapters/cockroach_asyncpg/config.py +++ b/sqlspec/adapters/cockroach_asyncpg/config.py @@ -95,8 +95,6 @@ class CockroachAsyncpgConnectionConfig(TypedDict): """Application identifier forwarded via ``server_settings``.""" cluster: NotRequired[str] """CockroachDB Serverless / Cloud routing cluster identifier.""" - pgbouncer: NotRequired[bool] - """Disable prepared statement caching when connecting through a transaction pooler.""" default_transaction_use_follower_reads: NotRequired[bool] """Session setting enabling bounded-staleness follower reads on read-only transactions.""" results_buffer_size: NotRequired[int] @@ -111,6 +109,10 @@ class CockroachAsyncpgConnectionConfig(TypedDict): """Startup CockroachDB session variables passed at connection time.""" target_session_attrs: NotRequired[CockroachAsyncpgTargetSessionAttrs] """Required node session attribute when connecting across multiple hosts.""" + statement_timeout: NotRequired[float | int | None] + """Statement execution timeout in seconds forwarded via server_settings.""" + idle_in_transaction_session_timeout: NotRequired[float | int | None] + """Idle in transaction timeout in seconds forwarded via server_settings.""" krbsrvname: NotRequired[str] """Kerberos service principal name for GSSAPI authentication.""" gsslib: NotRequired[CockroachAsyncpgGSSLib] @@ -192,24 +194,9 @@ class CockroachAsyncpgDriverFeatures(TypedDict): class _CockroachAsyncpgSessionFactory(AsyncPoolSessionFactory): - """Uses pool.acquire() context manager pattern instead of direct acquire/release.""" + """Factory creating AsyncPG database connections for CockroachDB sessions.""" - __slots__ = ("_ctx",) - - def __init__(self, config: "CockroachAsyncpgConfig") -> None: - super().__init__(config) - self._ctx: Any | None = None - - async def acquire_connection(self) -> "CockroachAsyncpgConnection": - pool = await self._config.provide_pool() - ctx = pool.acquire() - self._ctx = ctx - return cast("CockroachAsyncpgConnection", await ctx.__aenter__()) - - async def release_connection(self, _conn: "CockroachAsyncpgConnection", **kwargs: Any) -> None: - if self._ctx is not None: - await self._ctx.__aexit__(kwargs.get("exc_type"), kwargs.get("exc_val"), kwargs.get("exc_tb")) - self._ctx = None + __slots__ = () class CockroachAsyncpgConnectionContext(AsyncPoolConnectionContext): @@ -224,7 +211,7 @@ class CockroachAsyncpgConfig( """Configuration for CockroachDB using AsyncPG.""" driver_type: "ClassVar[type[CockroachAsyncpgDriver]]" = CockroachAsyncpgDriver - connection_type: "ClassVar[type[CockroachAsyncpgConnection]]" = CockroachAsyncpgConnection # type: ignore[assignment] + connection_type: "ClassVar[type[Any]]" = cast("type[Any]", CockroachAsyncpgConnection) supports_transactional_ddl: "ClassVar[bool]" = False supports_migration_schemas: "ClassVar[bool]" = True supports_native_arrow_export: "ClassVar[bool]" = True diff --git a/sqlspec/adapters/cockroach_asyncpg/core.py b/sqlspec/adapters/cockroach_asyncpg/core.py index d6352d7e8..4073f748a 100644 --- a/sqlspec/adapters/cockroach_asyncpg/core.py +++ b/sqlspec/adapters/cockroach_asyncpg/core.py @@ -234,7 +234,7 @@ def build_native_export( query: str, parameters: "list[Any]", uri: str, file_format: str, options: "dict[str, Any]" ) -> "tuple[str, list[Any]]": """Wrap compiled query SQL without embedding destination or CSV values.""" - format_sql = {"csv": "CSV", "parquet": "PARQUET"}[file_format] + format_sql = "PARQUET" if file_format == "parquet" else "CSV" values = [*parameters, uri] command = "EXPORT INTO " + format_sql + " " + ("$" + str(len(values))) if file_format == "csv" and "nullas" in options: @@ -250,7 +250,7 @@ def build_native_import(table: str, uri: str, file_format: str, options: "dict[s msg = "Table name must not be empty" raise ValueError(msg) target = ".".join(quote_identifier(part) for part in parts) - format_sql = {"csv": "CSV", "parquet": "PARQUET"}[file_format] + format_sql = "PARQUET" if file_format == "parquet" else "CSV" values: list[Any] = [uri] command = "IMPORT INTO " + target + " " + format_sql + " DATA (" + "$1" + ")" clauses = [] diff --git a/sqlspec/adapters/cockroach_asyncpg/data_dictionary.py b/sqlspec/adapters/cockroach_asyncpg/data_dictionary.py index 157546b46..888fef1bd 100644 --- a/sqlspec/adapters/cockroach_asyncpg/data_dictionary.py +++ b/sqlspec/adapters/cockroach_asyncpg/data_dictionary.py @@ -5,6 +5,7 @@ from mypy_extensions import mypyc_attr +from sqlspec.adapters.asyncpg.data_dictionary import AsyncpgDataDictionary from sqlspec.data_dictionary import ( ColumnMetadata, DDLResult, @@ -28,66 +29,48 @@ system_metadata_gated_result, unsupported_system_metadata_capability, ) -from sqlspec.data_dictionary.dialects.cockroachdb import resolve_cockroachdb_json_type -from sqlspec.driver import AsyncDataDictionaryBase +from sqlspec.data_dictionary.dialects.cockroachdb import ( + COCKROACH_METADATA_DOMAINS, + COCKROACH_SUPPORTED_DOMAINS, + resolve_cockroachdb_json_type, +) if TYPE_CHECKING: - from sqlspec.adapters.cockroach_asyncpg.driver import CockroachAsyncpgDriver + from sqlspec.adapters.asyncpg.driver import AsyncpgDriver from sqlspec.core import SQL __all__ = ("CockroachAsyncpgDataDictionary",) -_COCKROACH_METADATA_DOMAINS: tuple[str, ...] = ( - "schemas", - "objects", - "tables", - "columns", - "constraints", - "indexes", - "views", - "sequences", - "comments", - "privileges", - "dependencies", - "ddl", - "crdb_internal", - "system", -) -_COCKROACH_SUPPORTED_DOMAINS = frozenset(d for d in _COCKROACH_METADATA_DOMAINS if d not in {"crdb_internal", "system"}) - - @mypyc_attr(allow_interpreted_subclasses=True, native_class=False) -class CockroachAsyncpgDataDictionary(AsyncDataDictionaryBase): +class CockroachAsyncpgDataDictionary(AsyncpgDataDictionary): """CockroachDB async data dictionary (AsyncPG).""" dialect: ClassVar[str] = "cockroachdb" async def get_metadata_capabilities( - self, driver: "CockroachAsyncpgDriver", domains: Sequence[str] | None = None + self, driver: "AsyncpgDriver", domains: Sequence[str] | None = None ) -> MetadataCapabilityProfile: """Get CockroachDB data-dictionary capability profile.""" return _cockroach_metadata_profile(type(self).__name__, domains) async def get_system_metadata_capabilities( - self, driver: "CockroachAsyncpgDriver", domains: Sequence[str] | None = None + self, driver: "AsyncpgDriver", domains: Sequence[str] | None = None ) -> tuple[SystemMetadataCapability, ...]: """Get CockroachDB opt-in system metadata capability disclosures.""" _ = driver requested_domains = ("system",) if domains is None else tuple(domains) return tuple(unsupported_system_metadata_capability(domain) for domain in requested_domains) - async def get_schemas(self, driver: "CockroachAsyncpgDriver") -> MetadataResult: + async def get_schemas(self, driver: "AsyncpgDriver") -> MetadataResult: """Get schema metadata.""" return await self._select_domain(driver, "schemas", "list") - async def get_objects(self, driver: "CockroachAsyncpgDriver", schema: str | None = None) -> MetadataResult: + async def get_objects(self, driver: "AsyncpgDriver", schema: str | None = None) -> MetadataResult: """Get database object metadata.""" return await self._select_domain(driver, "objects", "by_schema", schema_name=self.resolve_schema(schema)) - async def get_table_details( - self, driver: "CockroachAsyncpgDriver", table: str, schema: str | None = None - ) -> MetadataResult: + async def get_table_details(self, driver: "AsyncpgDriver", table: str, schema: str | None = None) -> MetadataResult: """Get rich table metadata.""" return await self._select_domain( driver, @@ -98,7 +81,7 @@ async def get_table_details( ) async def get_constraints( - self, driver: "CockroachAsyncpgDriver", table: str | None = None, schema: str | None = None + self, driver: "AsyncpgDriver", table: str | None = None, schema: str | None = None ) -> MetadataResult: """Get constraint metadata.""" table_name = self.resolve_identifier(table) if table is not None else None @@ -106,16 +89,16 @@ async def get_constraints( driver, "constraints", "by_schema", schema_name=self.resolve_schema(schema), table_name=table_name ) - async def get_views(self, driver: "CockroachAsyncpgDriver", schema: str | None = None) -> MetadataResult: + async def get_views(self, driver: "AsyncpgDriver", schema: str | None = None) -> MetadataResult: """Get view metadata.""" return await self._select_domain(driver, "views", "by_schema", schema_name=self.resolve_schema(schema)) - async def get_routines(self, driver: "CockroachAsyncpgDriver", schema: str | None = None) -> MetadataResult: + async def get_routines(self, driver: "AsyncpgDriver", schema: str | None = None) -> MetadataResult: """Get routine metadata.""" return _metadata_result("routines", MetadataCapability.unsupported("routines")) async def get_privileges( - self, driver: "CockroachAsyncpgDriver", object_name: str | None = None, schema: str | None = None + self, driver: "AsyncpgDriver", object_name: str | None = None, schema: str | None = None ) -> MetadataResult: """Get privilege metadata.""" resolved_object = self.resolve_identifier(object_name) if object_name is not None else None @@ -124,7 +107,7 @@ async def get_privileges( ) async def get_dependencies( - self, driver: "CockroachAsyncpgDriver", object_name: str | None = None, schema: str | None = None + self, driver: "AsyncpgDriver", object_name: str | None = None, schema: str | None = None ) -> MetadataResult: """Get stable dependency metadata.""" resolved_object = self.resolve_identifier(object_name) if object_name is not None else None @@ -132,12 +115,12 @@ async def get_dependencies( driver, "dependencies", "by_schema", schema_name=self.resolve_schema(schema), object_name=resolved_object ) - async def get_sequences(self, driver: "CockroachAsyncpgDriver", schema: str | None = None) -> MetadataResult: + async def get_sequences(self, driver: "AsyncpgDriver", schema: str | None = None) -> MetadataResult: """Get sequence metadata.""" return await self._select_domain(driver, "sequences", "by_schema", schema_name=self.resolve_schema(schema)) async def get_system_metadata( - self, driver: "CockroachAsyncpgDriver", request: SystemMetadataRequest | str | None = None, **kwargs: Any + self, driver: "AsyncpgDriver", request: SystemMetadataRequest | str | None = None, **kwargs: Any ) -> SystemMetadataResult: """Get opt-in CockroachDB system metadata.""" _ = driver @@ -145,7 +128,7 @@ async def get_system_metadata( capability = unsupported_system_metadata_capability(metadata_request.domain) return system_metadata_gated_result(metadata_request, capability) - async def get_version(self, driver: "CockroachAsyncpgDriver") -> "VersionInfo | None": + async def get_version(self, driver: "AsyncpgDriver") -> "VersionInfo | None": driver_id = id(getattr(driver, "connection", None) or driver) if driver_id in self._version_fetch_attempted: return self._version_cache.get(driver_id) @@ -166,17 +149,17 @@ async def get_version(self, driver: "CockroachAsyncpgDriver") -> "VersionInfo | self.cache_version(driver_id, version_info) return version_info - async def get_feature_flag(self, driver: "CockroachAsyncpgDriver", feature: str) -> bool: + async def get_feature_flag(self, driver: "AsyncpgDriver", feature: str) -> bool: version_info = await self.get_version(driver) return self.resolve_feature_flag(feature, version_info) - async def get_optimal_type(self, driver: "CockroachAsyncpgDriver", type_category: str) -> str: + async def get_optimal_type(self, driver: "AsyncpgDriver", type_category: str) -> str: config = self.get_dialect_config() if type_category == "json": return resolve_cockroachdb_json_type(await self.get_version(driver)) return config.get_optimal_type(type_category) - async def get_tables(self, driver: "CockroachAsyncpgDriver", schema: "str | None" = None) -> "list[TableMetadata]": + async def get_tables(self, driver: "AsyncpgDriver", schema: "str | None" = None) -> "list[TableMetadata]": schema_name = self.resolve_schema(schema) self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="tables") return await driver.select( @@ -187,7 +170,7 @@ async def get_tables(self, driver: "CockroachAsyncpgDriver", schema: "str | None ) async def get_columns( - self, driver: "CockroachAsyncpgDriver", table: "str | None" = None, schema: "str | None" = None + self, driver: "AsyncpgDriver", table: "str | None" = None, schema: "str | None" = None ) -> "list[ColumnMetadata]": schema_name = self.resolve_schema(schema) if table is None: @@ -203,14 +186,14 @@ async def get_columns( table_schema_name = None if schema is None else schema_name self._log_table_describe(driver, schema_name=table_schema_name, table_name=table_name, operation="columns") return await driver.select( - _cockroach_domain_sql("columns", "by_schema"), + _cockroach_domain_sql("columns", "by_table"), schema_name=table_schema_name, table_name=table_name, schema_type=ColumnMetadata, ) async def get_indexes( - self, driver: "CockroachAsyncpgDriver", table: "str | None" = None, schema: "str | None" = None + self, driver: "AsyncpgDriver", table: "str | None" = None, schema: "str | None" = None ) -> "list[IndexMetadata]": schema_name = self.resolve_schema(schema) if table is None: @@ -223,16 +206,17 @@ async def get_indexes( ) table_name = self.resolve_identifier(table) - self._log_table_describe(driver, schema_name=schema_name, table_name=table_name, operation="indexes") + table_schema_name = None if schema is None else schema_name + self._log_table_describe(driver, schema_name=table_schema_name, table_name=table_name, operation="indexes") return await driver.select( - _cockroach_domain_sql("indexes", "by_schema"), - schema_name=schema_name, + _cockroach_domain_sql("indexes", "by_table"), + schema_name=table_schema_name, table_name=table_name, schema_type=IndexMetadata, ) async def get_foreign_keys( - self, driver: "CockroachAsyncpgDriver", table: "str | None" = None, schema: "str | None" = None + self, driver: "AsyncpgDriver", table: "str | None" = None, schema: "str | None" = None ) -> "list[ForeignKeyMetadata]": """Get foreign key metadata, including ON DELETE / ON UPDATE actions.""" schema_name = self.resolve_schema(schema) @@ -249,7 +233,7 @@ async def get_foreign_keys( async def _get_native_object_ddl( self, - driver: "CockroachAsyncpgDriver", + driver: "AsyncpgDriver", object_name: str, *, schema: str | None, @@ -276,7 +260,7 @@ async def _get_native_object_ddl( ) async def _select_domain( - self, driver: "CockroachAsyncpgDriver", domain: str, query_name: str, **parameters: Any + self, driver: "AsyncpgDriver", domain: str, query_name: str, **parameters: Any ) -> MetadataResult: query = get_data_dictionary_loader().get_domain_query(type(self).dialect, domain, query_name) if not query.is_supported or query.sql is None: @@ -314,7 +298,7 @@ def _cockroach_metadata_capability(domain: str) -> MetadataCapability: risks=(MetadataRisk.EXPENSIVE, MetadataRisk.PRIVILEGED), warnings=("System metadata is opt-in and disabled by default.",), ) - if domain in _COCKROACH_SUPPORTED_DOMAINS: + if domain in COCKROACH_SUPPORTED_DOMAINS: return MetadataCapability( domain=domain, support=MetadataSupport.SUPPORTED, @@ -325,7 +309,7 @@ def _cockroach_metadata_capability(domain: str) -> MetadataCapability: def _cockroach_metadata_profile(adapter: str, domains: Sequence[str] | None) -> MetadataCapabilityProfile: - requested_domains = _COCKROACH_METADATA_DOMAINS if domains is None else tuple(domains) + requested_domains = COCKROACH_METADATA_DOMAINS if domains is None else tuple(domains) return MetadataCapabilityProfile( "cockroachdb", adapter=adapter, diff --git a/sqlspec/adapters/cockroach_asyncpg/driver.py b/sqlspec/adapters/cockroach_asyncpg/driver.py index 51c9bc890..72003e29f 100644 --- a/sqlspec/adapters/cockroach_asyncpg/driver.py +++ b/sqlspec/adapters/cockroach_asyncpg/driver.py @@ -1,6 +1,7 @@ """CockroachDB AsyncPG driver implementation.""" import asyncio +import sys from contextlib import suppress from typing import TYPE_CHECKING, Any, TypeVar, cast @@ -20,23 +21,67 @@ ) from sqlspec.adapters.cockroach_asyncpg.data_dictionary import CockroachAsyncpgDataDictionary from sqlspec.core import SQL, register_driver_profile +from sqlspec.driver import AsyncRowStream 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.asyncpg.core import AsyncpgStreamSource from sqlspec.adapters.cockroach_asyncpg._typing import CockroachAsyncpgConnection - from sqlspec.core import StatementConfig - from sqlspec.driver import ExecutionResult + from sqlspec.core import SQLResult, StatementConfig + from sqlspec.driver import CachedQuery, ExecutionResult from sqlspec.storage import StorageBridgeJob, StorageDestination, StorageFormat, StorageTelemetry -__all__ = ("CockroachAsyncpgDriver", "CockroachAsyncpgExceptionHandler", "CockroachAsyncpgSessionContext") +__all__ = ( + "CockroachAsyncpgDriver", + "CockroachAsyncpgExceptionHandler", + "CockroachAsyncpgSessionContext", + "CockroachAsyncpgStreamSource", +) logger = get_logger("sqlspec.adapters.cockroach_asyncpg") _T = TypeVar("_T") +class CockroachAsyncpgStreamSource: + """Async chunk source for CockroachDB AsyncPG with follower read support.""" + + __slots__ = ("_auto_started", "_driver", "_statement", "_underlying") + + def __init__(self, driver: "CockroachAsyncpgDriver", statement: "SQL", underlying: "AsyncpgStreamSource") -> None: + self._driver = driver + self._statement = statement + self._underlying = underlying + self._auto_started = False + + async def start(self) -> None: + if self._driver._wants_follower_read(self._statement): + self._auto_started = await self._driver._begin_follower_read_transaction() + try: + await self._underlying.start() + except BaseException: + if self._auto_started and self._driver._connection_in_transaction(): + with suppress(Exception): + await self._driver.rollback() + raise + + async def fetch_chunk(self) -> "list[dict[str, Any]]": + return await self._underlying.fetch_chunk() + + async def close(self, error: bool = False) -> None: + try: + await self._underlying.close(error=error) + finally: + if self._auto_started and self._driver._connection_in_transaction(): + if error: + with suppress(Exception): + await self._driver.rollback() + else: + await self._driver.commit() + + class CockroachAsyncpgExceptionHandler(AsyncpgExceptionHandler): """Async context manager for CockroachDB AsyncPG exceptions.""" @@ -69,7 +114,7 @@ def __init__( self._data_dictionary = None @property - def data_dictionary(self) -> "CockroachAsyncpgDataDictionary": # type: ignore[override] + def data_dictionary(self) -> "CockroachAsyncpgDataDictionary": if self._data_dictionary is None: object.__setattr__(self, "_data_dictionary", CockroachAsyncpgDataDictionary()) return cast("CockroachAsyncpgDataDictionary", self._data_dictionary) @@ -164,6 +209,10 @@ async def load_from_storage( table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs ) + def _wants_follower_read(self, statement: "SQL") -> bool: + """Return True when a statement may run under a follower-read transaction.""" + return statement.returns_rows() and not statement.is_modifying_operation() + async def begin(self) -> None: """Begin a transaction and apply follower-read staleness to it. @@ -176,7 +225,12 @@ async def begin(self) -> None: await super().begin() return await super().begin() - await self._apply_follower_reads() + try: + await self._apply_follower_reads() + except Exception as exc: + with suppress(Exception): + await self.rollback() + raise create_mapped_exception(exc) from exc async def run_transaction_with_retry(self, operation: "Callable[[], Awaitable[_T]]") -> _T: """Execute a full CockroachDB transaction callback with serialization retries. @@ -211,7 +265,7 @@ async def run_transaction_with_retry(self, operation: "Callable[[], Awaitable[_T attempt += 1 async def dispatch_execute(self, cursor: Any, statement: SQL) -> "ExecutionResult": - auto_started = await self._begin_follower_read_transaction() if statement.returns_rows() else False + auto_started = await self._begin_follower_read_transaction() if self._wants_follower_read(statement) else False try: result = await super().dispatch_execute(cursor, statement) except BaseException: @@ -223,6 +277,33 @@ async def dispatch_execute(self, cursor: Any, statement: SQL) -> "ExecutionResul await self.commit() return result + async def _execute_cache_hit( + self, sql: str, params: "tuple[Any, ...] | list[Any] | dict[str, Any]", cached: "CachedQuery" + ) -> "SQLResult": + auto_started = ( + await self._begin_follower_read_transaction() + if (cached.operation_profile.returns_rows and not cached.operation_profile.modifies_rows) + else False + ) + try: + result = await super()._execute_cache_hit(sql, params, cached) + except BaseException: + if auto_started and self._connection_in_transaction(): + with suppress(Exception): + await self.rollback() + raise + if auto_started and self._connection_in_transaction(): + await self.commit() + return result + + def dispatch_select_stream(self, statement: "SQL", chunk_size: int) -> "AsyncRowStream[dict[str, Any]] | None": + """Return a native asyncpg row stream wrapped for follower reads.""" + stream = super().dispatch_select_stream(statement, chunk_size) + if stream is None: + return None + source = CockroachAsyncpgStreamSource(self, statement, cast("AsyncpgStreamSource", stream._source)) + return AsyncRowStream(source) + def handle_database_exceptions(self) -> "CockroachAsyncpgExceptionHandler": return CockroachAsyncpgExceptionHandler() @@ -231,11 +312,18 @@ def _native_storage_ready(self) -> bool: async def _execute_native_storage(self, command: str, parameters: "list[Any]") -> "list[dict[str, Any]]": handler = self.handle_database_exceptions() - rows = [] - async with handler: - rows = [dict(row) for row in await self.connection.fetch(command, *parameters)] - if handler.pending_exception is not None: - raise handler.pending_exception + rows: list[dict[str, Any]] = [] + + async def _fetch() -> list[dict[str, Any]]: + return [dict(row) for row in await self.connection.fetch(command, *parameters)] + + if sys.exc_info()[1] is None: + async with handler: + rows = await _fetch() + else: + res = await self._run_with_exception_handler(handler, _fetch) + rows = res if res is not None else [] + self._check_pending_exception(handler) return rows async def _apply_follower_reads(self) -> None: diff --git a/sqlspec/adapters/cockroach_psycopg/_typing.py b/sqlspec/adapters/cockroach_psycopg/_typing.py index f9db5c984..933503c58 100644 --- a/sqlspec/adapters/cockroach_psycopg/_typing.py +++ b/sqlspec/adapters/cockroach_psycopg/_typing.py @@ -32,10 +32,10 @@ from sqlspec.adapters.cockroach_psycopg.driver import CockroachPsycopgAsyncDriver, CockroachPsycopgSyncDriver from sqlspec.core import StatementConfig - CockroachSyncConnection: TypeAlias = CrdbConnection[PsycopgDictRow] - CockroachAsyncConnection: TypeAlias = AsyncCrdbConnection[PsycopgDictRow] - CockroachSyncCursor: TypeAlias = Cursor[PsycopgDictRow] - CockroachAsyncCursor: TypeAlias = AsyncCursor[PsycopgDictRow] + CockroachSyncConnection: TypeAlias = CrdbConnection[Any] + CockroachAsyncConnection: TypeAlias = AsyncCrdbConnection[Any] + CockroachSyncCursor: TypeAlias = Cursor[Any] + CockroachAsyncCursor: TypeAlias = AsyncCursor[Any] else: CockroachSyncConnection = cockroach_psycopg_crdb.CrdbConnection CockroachAsyncConnection = cockroach_psycopg_crdb.AsyncCrdbConnection diff --git a/sqlspec/adapters/cockroach_psycopg/adk/store.py b/sqlspec/adapters/cockroach_psycopg/adk/store.py index 12eac5f36..4dd797397 100644 --- a/sqlspec/adapters/cockroach_psycopg/adk/store.py +++ b/sqlspec/adapters/cockroach_psycopg/adk/store.py @@ -210,6 +210,8 @@ class CockroachPsycopgAsyncADKStore(BaseAsyncADKStore["CockroachPsycopgAsyncConf __slots__ = () + _max_table_name_parts = 2 + def __init__(self, config: "CockroachPsycopgAsyncConfig") -> None: super().__init__(config) @@ -700,6 +702,8 @@ class CockroachPsycopgSyncADKStore(BaseSyncADKStore["CockroachPsycopgSyncConfig" __slots__ = () + _max_table_name_parts = 2 + def __init__(self, config: "CockroachPsycopgSyncConfig") -> None: super().__init__(config) @@ -1211,6 +1215,8 @@ class CockroachPsycopgAsyncADKMemoryStore(BaseAsyncADKMemoryStore["CockroachPsyc __slots__ = () + _max_table_name_parts = 2 + def __init__(self, config: "CockroachPsycopgAsyncConfig") -> None: super().__init__(config) @@ -1234,7 +1240,21 @@ async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: " if not entries: return 0 - inserted_count = 0 + ids = [entry["id"] for entry in entries] + session_ids = [entry["session_id"] for entry in entries] + app_names = [entry["app_name"] for entry in entries] + user_ids = [entry["user_id"] for entry in entries] + scopes = [entry.get("scope", "user") for entry in entries] + event_ids = [entry["event_id"] for entry in entries] + authors = [entry["author"] for entry in entries] + timestamps = [entry["timestamp"] for entry in entries] + content_jsons = [Jsonb(entry["content_json"]) for entry in entries] + content_texts = [entry["content_text"] for entry in entries] + metadata_jsons = [ + Jsonb(entry["metadata_json"]) if entry["metadata_json"] is not None else None for entry in entries + ] + inserted_ats = [entry["inserted_at"] for entry in entries] + table_ident = pg_sql.Identifier(*self._memory_table.split(".")) if self._owner_id_column_name: query = pg_sql.SQL(""" @@ -1242,30 +1262,63 @@ async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: " id, session_id, app_name, user_id, scope, event_id, author, {owner_id_col}, timestamp, content_json, content_text, metadata_json, inserted_at - ) VALUES ( - %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s ) + SELECT id, session_id, app_name, user_id, scope, event_id, author, + %s, timestamp, content_json, content_text, metadata_json, inserted_at + FROM unnest( + %s::varchar[], %s::varchar[], %s::varchar[], %s::varchar[], + %s::varchar[], %s::varchar[], %s::varchar[], + %s::timestamptz[], %s::jsonb[], %s::text[], %s::jsonb[], %s::timestamptz[] + ) AS batch(id, session_id, app_name, user_id, scope, event_id, author, + timestamp, content_json, content_text, metadata_json, inserted_at) ON CONFLICT (event_id) DO NOTHING """).format(table=table_ident, owner_id_col=pg_sql.Identifier(self._owner_id_column_name)) + params: tuple[Any, ...] = ( + owner_id, + ids, + session_ids, + app_names, + user_ids, + scopes, + event_ids, + authors, + timestamps, + content_jsons, + content_texts, + metadata_jsons, + inserted_ats, + ) else: query = pg_sql.SQL(""" INSERT INTO {table} ( id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json, inserted_at - ) VALUES ( - %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s + ) + SELECT * FROM unnest( + %s::varchar[], %s::varchar[], %s::varchar[], %s::varchar[], + %s::varchar[], %s::varchar[], %s::varchar[], + %s::timestamptz[], %s::jsonb[], %s::text[], %s::jsonb[], %s::timestamptz[] ) ON CONFLICT (event_id) DO NOTHING """).format(table=table_ident) + params = ( + ids, + session_ids, + app_names, + user_ids, + scopes, + event_ids, + authors, + timestamps, + content_jsons, + content_texts, + metadata_jsons, + inserted_ats, + ) async with self._config.provide_connection() as conn, conn.cursor() as cur: - for entry in entries: - if self._owner_id_column_name: - await cur.execute(query, _build_insert_params_with_owner(entry, owner_id)) - else: - await cur.execute(query, _build_insert_params(entry)) - if cur.rowcount and cur.rowcount > 0: - inserted_count += cur.rowcount + await cur.execute(query, params) + inserted_count = cur.rowcount if cur.rowcount and cur.rowcount > 0 else 0 await conn.commit() return inserted_count @@ -1298,6 +1351,7 @@ async def search_entries( LIMIT %s """ search_param = query + params = (*scope_params, search_param, effective_limit) else: sql = f""" SELECT * FROM {self._memory_table} @@ -1306,8 +1360,7 @@ async def search_entries( LIMIT %s """ search_param = f"%{query}%" - - params = (*scope_params, search_param, effective_limit) + params = (*scope_params, search_param, effective_limit) try: async with self._config.provide_connection() as conn, conn.cursor() as cur: @@ -1395,6 +1448,8 @@ class CockroachPsycopgSyncADKMemoryStore(BaseSyncADKMemoryStore["CockroachPsycop __slots__ = () + _max_table_name_parts = 2 + def __init__(self, config: "CockroachPsycopgSyncConfig") -> None: super().__init__(config) @@ -1420,7 +1475,21 @@ def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object if not entries: return 0 - inserted_count = 0 + ids = [entry["id"] for entry in entries] + session_ids = [entry["session_id"] for entry in entries] + app_names = [entry["app_name"] for entry in entries] + user_ids = [entry["user_id"] for entry in entries] + scopes = [entry.get("scope", "user") for entry in entries] + event_ids = [entry["event_id"] for entry in entries] + authors = [entry["author"] for entry in entries] + timestamps = [entry["timestamp"] for entry in entries] + content_jsons = [Jsonb(entry["content_json"]) for entry in entries] + content_texts = [entry["content_text"] for entry in entries] + metadata_jsons = [ + Jsonb(entry["metadata_json"]) if entry["metadata_json"] is not None else None for entry in entries + ] + inserted_ats = [entry["inserted_at"] for entry in entries] + table_ident = pg_sql.Identifier(*self._memory_table.split(".")) if self._owner_id_column_name: query = pg_sql.SQL(""" @@ -1428,31 +1497,65 @@ def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object id, session_id, app_name, user_id, scope, event_id, author, {owner_id_col}, timestamp, content_json, content_text, metadata_json, inserted_at - ) VALUES ( - %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s ) + SELECT id, session_id, app_name, user_id, scope, event_id, author, + %s, timestamp, content_json, content_text, metadata_json, inserted_at + FROM unnest( + %s::varchar[], %s::varchar[], %s::varchar[], %s::varchar[], + %s::varchar[], %s::varchar[], %s::varchar[], + %s::timestamptz[], %s::jsonb[], %s::text[], %s::jsonb[], %s::timestamptz[] + ) AS batch(id, session_id, app_name, user_id, scope, event_id, author, + timestamp, content_json, content_text, metadata_json, inserted_at) ON CONFLICT (event_id) DO NOTHING """).format(table=table_ident, owner_id_col=pg_sql.Identifier(self._owner_id_column_name)) + params: tuple[Any, ...] = ( + owner_id, + ids, + session_ids, + app_names, + user_ids, + scopes, + event_ids, + authors, + timestamps, + content_jsons, + content_texts, + metadata_jsons, + inserted_ats, + ) else: query = pg_sql.SQL(""" INSERT INTO {table} ( id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json, inserted_at - ) VALUES ( - %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s + ) + SELECT * FROM unnest( + %s::varchar[], %s::varchar[], %s::varchar[], %s::varchar[], + %s::varchar[], %s::varchar[], %s::varchar[], + %s::timestamptz[], %s::jsonb[], %s::text[], %s::jsonb[], %s::timestamptz[] ) ON CONFLICT (event_id) DO NOTHING """).format(table=table_ident) + params = ( + ids, + session_ids, + app_names, + user_ids, + scopes, + event_ids, + authors, + timestamps, + content_jsons, + content_texts, + metadata_jsons, + inserted_ats, + ) with self._config.provide_connection() as conn, conn.cursor() as cur: - for entry in entries: - if self._owner_id_column_name: - cur.execute(query, _build_insert_params_with_owner(entry, owner_id)) - else: - cur.execute(query, _build_insert_params(entry)) - if cur.rowcount and cur.rowcount > 0: - inserted_count += cur.rowcount + cur.execute(query, params) + inserted_count = cur.rowcount if cur.rowcount and cur.rowcount > 0 else 0 conn.commit() + return inserted_count def search_entries( @@ -1484,6 +1587,7 @@ def search_entries( LIMIT %s """ search_param = query + params = (*scope_params, search_param, effective_limit) else: sql = f""" SELECT * FROM {self._memory_table} @@ -1492,8 +1596,7 @@ def search_entries( LIMIT %s """ search_param = f"%{query}%" - - params = (*scope_params, search_param, effective_limit) + params = (*scope_params, search_param, effective_limit) try: with self._config.provide_connection() as conn, conn.cursor() as cur: @@ -1647,9 +1750,12 @@ def _cockroach_hash_shard_clause(adk_config: CockroachPsycopgADKConfig) -> str: if not adk_config.get("enable_hash_sharded_indexes", False): return "" bucket_count = adk_config.get("hash_shard_bucket_count") - if isinstance(bucket_count, int) and not isinstance(bucket_count, bool) and bucket_count > 0: - return f" USING HASH WITH (bucket_count = {bucket_count})" - return " USING HASH" + if bucket_count is None: + return " USING HASH" + if not isinstance(bucket_count, int) or isinstance(bucket_count, bool) or bucket_count <= 0: + msg = "extension_config['adk']['hash_shard_bucket_count'] must be a positive integer" + raise ValueError(msg) + return f" USING HASH WITH (bucket_count = {bucket_count})" def _cockroach_storing_clause(adk_config: CockroachPsycopgADKConfig, columns: tuple[str, ...]) -> str: diff --git a/sqlspec/adapters/cockroach_psycopg/config.py b/sqlspec/adapters/cockroach_psycopg/config.py index 1bacb0375..5161946b7 100644 --- a/sqlspec/adapters/cockroach_psycopg/config.py +++ b/sqlspec/adapters/cockroach_psycopg/config.py @@ -1,5 +1,6 @@ """CockroachDB configuration using psycopg.""" +import inspect import re from contextlib import suppress from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, NotRequired, TypedDict, cast @@ -279,6 +280,7 @@ def __init__( self._user_connection_hook: Callable[[CockroachSyncConnection], None] | None = features_dict.pop( "on_connection_create", None ) + self._user_configure: Callable[[CockroachSyncConnection], Any] | None = connection_config.get("configure") super().__init__( connection_config=connection_config, @@ -302,14 +304,13 @@ def create_connection(self) -> "CockroachSyncConnection": """ conninfo, connection_kwargs = _standalone_connection_kwargs(dict(self.connection_config)) connection = psycopg_crdb.CrdbConnection.connect(conninfo, **connection_kwargs) - configure = self.connection_config.get("configure", self._configure_connection) try: - configure(cast("CockroachSyncConnection", connection)) + self._configure_connection(connection) except BaseException: with suppress(Exception): connection.close() raise - return cast("CockroachSyncConnection", connection) + return connection def provide_session( self, @@ -362,6 +363,7 @@ def get_event_runtime_hints(self) -> "EventRuntimeHints": def _create_pool(self) -> "ConnectionPool": all_config = dict(self.connection_config) + all_config.pop("configure", None) pool_parameters = { "min_size": all_config.pop("min_size", 4), @@ -380,7 +382,7 @@ def _create_pool(self) -> "ConnectionPool": "close_returns": all_config.pop("close_returns", False), } - pool_parameters["configure"] = all_config.pop("configure", self._configure_connection) + pool_parameters["configure"] = self._configure_connection pool_parameters = {k: v for k, v in pool_parameters.items() if v is not None} conninfo = all_config.pop("conninfo", None) @@ -412,6 +414,12 @@ def _configure_connection(self, conn: "CockroachSyncConnection") -> None: if self._user_connection_hook is not None: self._user_connection_hook(conn) + if self._user_configure is not None: + result = self._user_configure(conn) + if inspect.isawaitable(result): + msg = "Async configure callback cannot be used with CockroachPsycopgSyncConfig" + raise ImproperConfigurationError(msg) + def _close_pool(self) -> None: if not self.connection_instance: return @@ -515,6 +523,7 @@ def __init__( self._user_connection_hook: Callable[[CockroachAsyncConnection], Awaitable[None]] | None = features_dict.pop( "on_connection_create", None ) + self._user_configure: Callable[[CockroachAsyncConnection], Any] | None = connection_config.get("configure") super().__init__( connection_config=connection_config, @@ -538,14 +547,13 @@ async def create_connection(self) -> "CockroachAsyncConnection": """ conninfo, connection_kwargs = _standalone_connection_kwargs(dict(self.connection_config)) connection = await psycopg_crdb.AsyncCrdbConnection.connect(conninfo, **connection_kwargs) - configure = self.connection_config.get("configure", self._configure_async_connection) try: - await configure(cast("CockroachAsyncConnection", connection)) + await self._configure_async_connection(connection) except BaseException: with suppress(Exception): await connection.close() raise - return cast("CockroachAsyncConnection", connection) + return connection def provide_session( self, @@ -598,6 +606,7 @@ def get_event_runtime_hints(self) -> "EventRuntimeHints": async def _create_pool(self) -> "AsyncConnectionPool": all_config = dict(self.connection_config) + all_config.pop("configure", None) pool_parameters = { "min_size": all_config.pop("min_size", 4), @@ -617,7 +626,7 @@ async def _create_pool(self) -> "AsyncConnectionPool": open_pool = all_config.pop("open", True) pool_parameters["open"] = False if open_pool is True else open_pool - pool_parameters["configure"] = all_config.pop("configure", self._configure_async_connection) + pool_parameters["configure"] = self._configure_async_connection pool_parameters = {k: v for k, v in pool_parameters.items() if v is not None} conninfo = all_config.pop("conninfo", None) @@ -655,6 +664,11 @@ async def _configure_async_connection(self, conn: "CockroachAsyncConnection") -> if self._user_connection_hook is not None: await self._user_connection_hook(conn) + if self._user_configure is not None: + result = self._user_configure(conn) + if inspect.isawaitable(result): + await result + async def _close_pool(self) -> None: if not self.connection_instance: return diff --git a/sqlspec/adapters/cockroach_psycopg/data_dictionary.py b/sqlspec/adapters/cockroach_psycopg/data_dictionary.py index e4b5e39a7..e6007b5be 100644 --- a/sqlspec/adapters/cockroach_psycopg/data_dictionary.py +++ b/sqlspec/adapters/cockroach_psycopg/data_dictionary.py @@ -28,7 +28,11 @@ system_metadata_gated_result, unsupported_system_metadata_capability, ) -from sqlspec.data_dictionary.dialects.cockroachdb import resolve_cockroachdb_json_type +from sqlspec.data_dictionary.dialects.cockroachdb import ( + COCKROACH_METADATA_DOMAINS, + COCKROACH_SUPPORTED_DOMAINS, + resolve_cockroachdb_json_type, +) from sqlspec.driver import AsyncDataDictionaryBase, SyncDataDictionaryBase if TYPE_CHECKING: @@ -38,25 +42,6 @@ __all__ = ("CockroachPsycopgAsyncDataDictionary", "CockroachPsycopgSyncDataDictionary") -_COCKROACH_METADATA_DOMAINS: tuple[str, ...] = ( - "schemas", - "objects", - "tables", - "columns", - "constraints", - "indexes", - "views", - "sequences", - "comments", - "privileges", - "dependencies", - "ddl", - "crdb_internal", - "system", -) -_COCKROACH_SUPPORTED_DOMAINS = frozenset(d for d in _COCKROACH_METADATA_DOMAINS if d not in {"crdb_internal", "system"}) - - @mypyc_attr(allow_interpreted_subclasses=True, native_class=False) class CockroachPsycopgSyncDataDictionary(SyncDataDictionaryBase): """CockroachDB sync data dictionary.""" @@ -211,7 +196,7 @@ def get_columns( table_schema_name = None if schema is None else schema_name self._log_table_describe(driver, schema_name=table_schema_name, table_name=table_name, operation="columns") return driver.select( - _cockroach_domain_sql("columns", "by_schema"), + _cockroach_domain_sql("columns", "by_table"), schema_name=table_schema_name, table_name=table_name, schema_type=ColumnMetadata, @@ -232,10 +217,11 @@ def get_indexes( ) table_name = self.resolve_identifier(table) - self._log_table_describe(driver, schema_name=schema_name, table_name=table_name, operation="indexes") + table_schema_name = None if schema is None else schema_name + self._log_table_describe(driver, schema_name=table_schema_name, table_name=table_name, operation="indexes") return driver.select( - _cockroach_domain_sql("indexes", "by_schema"), - schema_name=schema_name, + _cockroach_domain_sql("indexes", "by_table"), + schema_name=table_schema_name, table_name=table_name, schema_type=IndexMetadata, ) @@ -450,7 +436,7 @@ async def get_columns( table_schema_name = None if schema is None else schema_name self._log_table_describe(driver, schema_name=table_schema_name, table_name=table_name, operation="columns") return await driver.select( - _cockroach_domain_sql("columns", "by_schema"), + _cockroach_domain_sql("columns", "by_table"), schema_name=table_schema_name, table_name=table_name, schema_type=ColumnMetadata, @@ -471,10 +457,11 @@ async def get_indexes( ) table_name = self.resolve_identifier(table) - self._log_table_describe(driver, schema_name=schema_name, table_name=table_name, operation="indexes") + table_schema_name = None if schema is None else schema_name + self._log_table_describe(driver, schema_name=table_schema_name, table_name=table_name, operation="indexes") return await driver.select( - _cockroach_domain_sql("indexes", "by_schema"), - schema_name=schema_name, + _cockroach_domain_sql("indexes", "by_table"), + schema_name=table_schema_name, table_name=table_name, schema_type=IndexMetadata, ) @@ -562,7 +549,7 @@ def _cockroach_metadata_capability(domain: str) -> MetadataCapability: risks=(MetadataRisk.EXPENSIVE, MetadataRisk.PRIVILEGED), warnings=("System metadata is opt-in and disabled by default.",), ) - if domain in _COCKROACH_SUPPORTED_DOMAINS: + if domain in COCKROACH_SUPPORTED_DOMAINS: return MetadataCapability( domain=domain, support=MetadataSupport.SUPPORTED, @@ -573,7 +560,7 @@ def _cockroach_metadata_capability(domain: str) -> MetadataCapability: def _cockroach_metadata_profile(adapter: str, domains: Sequence[str] | None) -> MetadataCapabilityProfile: - requested_domains = _COCKROACH_METADATA_DOMAINS if domains is None else tuple(domains) + requested_domains = COCKROACH_METADATA_DOMAINS if domains is None else tuple(domains) return MetadataCapabilityProfile( "cockroachdb", adapter=adapter, diff --git a/sqlspec/adapters/cockroach_psycopg/driver.py b/sqlspec/adapters/cockroach_psycopg/driver.py index cbc4d9afd..654f82e99 100644 --- a/sqlspec/adapters/cockroach_psycopg/driver.py +++ b/sqlspec/adapters/cockroach_psycopg/driver.py @@ -1,6 +1,7 @@ """CockroachDB psycopg driver implementation.""" import asyncio +import sys import time from contextlib import suppress from typing import TYPE_CHECKING, Any, TypeVar, cast @@ -29,6 +30,7 @@ CockroachPsycopgAsyncDataDictionary, CockroachPsycopgSyncDataDictionary, ) +from sqlspec.adapters.psycopg.core import create_mapped_exception from sqlspec.adapters.psycopg.driver import ( PsycopgAsyncDriver, PsycopgAsyncExceptionHandler, @@ -36,12 +38,14 @@ PsycopgSyncExceptionHandler, ) from sqlspec.core import SQL, StatementConfig, get_cache_config, register_driver_profile +from sqlspec.driver import AsyncRowStream, SyncRowStream from sqlspec.utils.logging import get_logger if TYPE_CHECKING: from collections.abc import Awaitable, Callable from sqlspec.adapters.cockroach_psycopg._typing import CockroachAsyncCursor, CockroachSyncCursor + from sqlspec.adapters.psycopg.core import PsycopgAsyncStreamSource, PsycopgSyncStreamSource from sqlspec.core import SQLResult from sqlspec.driver import CachedQuery, ExecutionResult from sqlspec.storage import StorageBridgeJob, StorageDestination, StorageFormat, StorageTelemetry @@ -50,15 +54,92 @@ "CockroachPsycopgAsyncDriver", "CockroachPsycopgAsyncExceptionHandler", "CockroachPsycopgAsyncSessionContext", + "CockroachPsycopgAsyncStreamSource", "CockroachPsycopgSyncDriver", "CockroachPsycopgSyncExceptionHandler", "CockroachPsycopgSyncSessionContext", + "CockroachPsycopgSyncStreamSource", ) logger = get_logger("sqlspec.adapters.cockroach_psycopg") _T = TypeVar("_T") +class CockroachPsycopgSyncStreamSource: + """Sync chunk source for CockroachDB Psycopg with follower read support.""" + + __slots__ = ("_auto_started", "_driver", "_underlying") + + def __init__( + self, driver: "CockroachPsycopgSyncDriver", underlying: "PsycopgSyncStreamSource", auto_started: bool + ) -> None: + self._driver = driver + self._underlying = underlying + self._auto_started = auto_started + + def start(self) -> None: + try: + self._underlying.start() + except BaseException: + if self._auto_started and self._driver._connection_in_transaction(): + with suppress(Exception): + self._driver.rollback() + raise + + def fetch_chunk(self) -> "list[dict[str, Any]]": + return self._underlying.fetch_chunk() + + def close(self, error: bool = False) -> None: + try: + self._underlying.close(error=error) + finally: + if self._auto_started and self._driver._connection_in_transaction(): + if error: + with suppress(Exception): + self._driver.rollback() + else: + self._driver.commit() + + +class CockroachPsycopgAsyncStreamSource: + """Async chunk source for CockroachDB Psycopg with follower read support.""" + + __slots__ = ("_auto_started", "_driver", "_statement", "_underlying") + + def __init__( + self, driver: "CockroachPsycopgAsyncDriver", statement: "SQL", underlying: "PsycopgAsyncStreamSource" + ) -> None: + self._driver = driver + self._statement = statement + self._underlying = underlying + self._auto_started = False + + async def start(self) -> None: + if self._driver._wants_follower_read(self._statement): + self._auto_started = await self._driver._begin_follower_read_transaction() + try: + await self._underlying.start() + except BaseException: + if self._auto_started and self._driver._connection_in_transaction(): + with suppress(Exception): + await self._driver.rollback() + raise + + async def fetch_chunk(self) -> "list[dict[str, Any]]": + return await self._underlying.fetch_chunk() + + async def close(self, error: bool = False) -> None: + try: + await self._underlying.close(error=error) + finally: + if self._auto_started and self._driver._connection_in_transaction(): + if error: + with suppress(Exception): + await self._driver.rollback() + else: + await self._driver.commit() + + class CockroachPsycopgSyncExceptionHandler(PsycopgSyncExceptionHandler): """Context manager for handling CockroachDB psycopg exceptions.""" @@ -74,7 +155,7 @@ class CockroachPsycopgAsyncExceptionHandler(PsycopgAsyncExceptionHandler): class CockroachPsycopgSyncDriver(PsycopgSyncDriver): """CockroachDB sync driver using psycopg.crdb.""" - __slots__ = ("_enable_retry", "_follower_staleness", "_retry_config") + __slots__ = ("_enable_retry", "_follower_reads_enabled", "_follower_staleness", "_retry_config") dialect = "postgres" def __init__( @@ -92,7 +173,9 @@ def __init__( 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")) + self._follower_reads_enabled = bool(self.driver_features.get("enable_follower_reads", False)) + raw_staleness = cast("str | None", self.driver_features.get("default_staleness")) + self._follower_staleness = validate_follower_read_staleness(raw_staleness) if raw_staleness else None self._data_dictionary = None @property @@ -191,6 +274,10 @@ def load_from_storage( table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs ) + def _wants_follower_read(self, statement: "SQL") -> bool: + """Return True when a statement may run under a follower-read transaction.""" + return statement.returns_rows() and not statement.is_modifying_operation() + def begin(self) -> None: """Begin a transaction and apply follower-read staleness to it. @@ -203,7 +290,13 @@ def begin(self) -> None: super().begin() return super().begin() - self._apply_follower_reads() + try: + self._apply_follower_reads() + except Exception as exc: + with suppress(Exception): + self.rollback() + self._transaction_active = False + raise create_mapped_exception(exc) from exc def run_transaction_with_retry(self, operation: "Callable[[], _T]") -> _T: """Execute a full CockroachDB transaction callback with serialization retries. @@ -238,7 +331,7 @@ def run_transaction_with_retry(self, operation: "Callable[[], _T]") -> _T: attempt += 1 def dispatch_execute(self, cursor: "CockroachSyncCursor", statement: SQL) -> "ExecutionResult": - auto_started = self._begin_follower_read_transaction() if statement.returns_rows() else False + auto_started = self._begin_follower_read_transaction() if self._wants_follower_read(statement) else False try: result = super().dispatch_execute(cursor, statement) except BaseException: @@ -250,13 +343,40 @@ def dispatch_execute(self, cursor: "CockroachSyncCursor", statement: SQL) -> "Ex self.commit() return result + def dispatch_select_stream(self, statement: "SQL", chunk_size: int) -> "SyncRowStream[dict[str, Any]] | None": + """Return a native psycopg row stream wrapped for follower reads.""" + auto_started = self._begin_follower_read_transaction() if self._wants_follower_read(statement) else False + try: + stream = super().dispatch_select_stream(statement, chunk_size) + if stream is None: + if auto_started and self._connection_in_transaction(): + self.commit() + return None + source = CockroachPsycopgSyncStreamSource( + self, cast("PsycopgSyncStreamSource", stream._source), auto_started + ) + return SyncRowStream(source) + except BaseException: + if auto_started and self._connection_in_transaction(): + with suppress(Exception): + self.rollback() + raise + def handle_database_exceptions(self) -> "CockroachPsycopgSyncExceptionHandler": return CockroachPsycopgSyncExceptionHandler() def _execute_cache_hit( self, sql: str, params: "tuple[Any, ...] | list[Any] | dict[str, Any]", cached: "CachedQuery" ) -> "SQLResult": - auto_started = self._begin_follower_read_transaction() if cached.operation_profile.returns_rows else False + auto_started = ( + self._begin_follower_read_transaction() + if ( + self._follower_reads_enabled + and cached.operation_profile.returns_rows + and not cached.operation_profile.modifies_rows + ) + else False + ) try: result = super()._execute_cache_hit(sql, params, cached) except BaseException: @@ -277,22 +397,27 @@ def _native_storage_ready(self) -> bool: def _execute_native_storage(self, command: str, parameters: "list[Any]") -> "list[dict[str, Any]]": handler = self.handle_database_exceptions() - rows = [] - with self.with_cursor(self.connection) as cursor, handler: - cursor.row_factory = dict_row - cursor.execute(command.encode("utf-8"), parameters) - rows = cursor.fetchall() - if handler.pending_exception is not None: - raise handler.pending_exception + rows: list[dict[str, Any]] = [] + + def _fetch() -> list[dict[str, Any]]: + with self.with_cursor(self.connection) as cursor: + cursor.row_factory = dict_row + cursor.execute(command.encode("utf-8"), parameters) + return cursor.fetchall() + + if sys.exc_info()[1] is None: + with handler: + rows = _fetch() + else: + res = self._run_with_exception_handler(handler, _fetch) + rows = res if res is not None else [] + self._check_pending_exception(handler) return rows def _apply_follower_reads(self) -> None: - if not self.driver_features.get("enable_follower_reads", False): - return - if not self._follower_staleness: + if not self._follower_reads_enabled or not self._follower_staleness: return - staleness = validate_follower_read_staleness(self._follower_staleness) - self.connection.execute(cast("Any", f"SET TRANSACTION AS OF SYSTEM TIME {staleness}")).close() + self.connection.execute(cast("Any", f"SET TRANSACTION AS OF SYSTEM TIME {self._follower_staleness}")).close() def _begin_follower_read_transaction(self) -> bool: """Open the transaction a follower read needs so the staleness clause can lead it. @@ -301,9 +426,7 @@ def _begin_follower_read_transaction(self) -> bool: clause with nowhere to go, so a read opens one here when the caller has not already done so. """ - if not self.driver_features.get("enable_follower_reads", False): - return False - if not self._follower_staleness: + if not self._follower_reads_enabled or not self._follower_staleness: return False if self._connection_in_transaction(): return False @@ -314,7 +437,7 @@ def _begin_follower_read_transaction(self) -> bool: class CockroachPsycopgAsyncDriver(PsycopgAsyncDriver): """CockroachDB async driver using psycopg.crdb.""" - __slots__ = ("_enable_retry", "_follower_staleness", "_retry_config") + __slots__ = ("_enable_retry", "_follower_reads_enabled", "_follower_staleness", "_retry_config") dialect = "postgres" def __init__( @@ -332,7 +455,9 @@ def __init__( 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")) + self._follower_reads_enabled = bool(self.driver_features.get("enable_follower_reads", False)) + raw_staleness = cast("str | None", self.driver_features.get("default_staleness")) + self._follower_staleness = validate_follower_read_staleness(raw_staleness) if raw_staleness else None self._data_dictionary = None @property @@ -431,6 +556,10 @@ async def load_from_storage( table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs ) + def _wants_follower_read(self, statement: "SQL") -> bool: + """Return True when a statement may run under a follower-read transaction.""" + return statement.returns_rows() and not statement.is_modifying_operation() + async def begin(self) -> None: """Begin a transaction and apply follower-read staleness to it. @@ -443,7 +572,13 @@ async def begin(self) -> None: await super().begin() return await super().begin() - await self._apply_follower_reads() + try: + await self._apply_follower_reads() + except Exception as exc: + with suppress(Exception): + await self.rollback() + self._transaction_active = False + raise create_mapped_exception(exc) from exc async def run_transaction_with_retry(self, operation: "Callable[[], Awaitable[_T]]") -> _T: """Execute a full CockroachDB transaction callback with serialization retries. @@ -478,7 +613,7 @@ async def run_transaction_with_retry(self, operation: "Callable[[], Awaitable[_T attempt += 1 async def dispatch_execute(self, cursor: "CockroachAsyncCursor", statement: SQL) -> "ExecutionResult": - auto_started = await self._begin_follower_read_transaction() if statement.returns_rows() else False + auto_started = await self._begin_follower_read_transaction() if self._wants_follower_read(statement) else False try: result = await super().dispatch_execute(cursor, statement) except BaseException: @@ -490,13 +625,29 @@ async def dispatch_execute(self, cursor: "CockroachAsyncCursor", statement: SQL) await self.commit() return result + def dispatch_select_stream(self, statement: "SQL", chunk_size: int) -> "AsyncRowStream[dict[str, Any]] | None": + """Return a native psycopg row stream wrapped for follower reads.""" + stream = super().dispatch_select_stream(statement, chunk_size) + if stream is None: + return None + source = CockroachPsycopgAsyncStreamSource(self, statement, cast("PsycopgAsyncStreamSource", stream._source)) + return AsyncRowStream(source) + def handle_database_exceptions(self) -> "CockroachPsycopgAsyncExceptionHandler": return CockroachPsycopgAsyncExceptionHandler() async def _execute_cache_hit( self, sql: str, params: "tuple[Any, ...] | list[Any] | dict[str, Any]", cached: "CachedQuery" ) -> "SQLResult": - auto_started = await self._begin_follower_read_transaction() if cached.operation_profile.returns_rows else False + auto_started = ( + await self._begin_follower_read_transaction() + if ( + self._follower_reads_enabled + and cached.operation_profile.returns_rows + and not cached.operation_profile.modifies_rows + ) + else False + ) try: result = await super()._execute_cache_hit(sql, params, cached) except BaseException: @@ -517,22 +668,29 @@ def _native_storage_ready(self) -> bool: async def _execute_native_storage(self, command: str, parameters: "list[Any]") -> "list[dict[str, Any]]": handler = self.handle_database_exceptions() - rows = [] - async with self.with_cursor(self.connection) as cursor, handler: - cursor.row_factory = dict_row - await cursor.execute(command.encode("utf-8"), parameters) - rows = await cursor.fetchall() - if handler.pending_exception is not None: - raise handler.pending_exception + rows: list[dict[str, Any]] = [] + + async def _fetch() -> list[dict[str, Any]]: + async with self.with_cursor(self.connection) as cursor: + cursor.row_factory = dict_row + await cursor.execute(command.encode("utf-8"), parameters) + return await cursor.fetchall() + + if sys.exc_info()[1] is None: + async with handler: + rows = await _fetch() + else: + res = await self._run_with_exception_handler(handler, _fetch) + rows = res if res is not None else [] + self._check_pending_exception(handler) return rows async def _apply_follower_reads(self) -> None: - if not self.driver_features.get("enable_follower_reads", False): + if not self._follower_reads_enabled or not self._follower_staleness: return - if not self._follower_staleness: - return - staleness = validate_follower_read_staleness(self._follower_staleness) - cursor = await self.connection.execute(cast("Any", f"SET TRANSACTION AS OF SYSTEM TIME {staleness}")) + cursor = await self.connection.execute( + cast("Any", f"SET TRANSACTION AS OF SYSTEM TIME {self._follower_staleness}") + ) await cursor.close() async def _begin_follower_read_transaction(self) -> bool: @@ -542,9 +700,7 @@ async def _begin_follower_read_transaction(self) -> bool: clause with nowhere to go, so a read opens one here when the caller has not already done so. """ - if not self.driver_features.get("enable_follower_reads", False): - return False - if not self._follower_staleness: + if not self._follower_reads_enabled or not self._follower_staleness: return False if self._connection_in_transaction(): return False diff --git a/sqlspec/adapters/db2/_typing.py b/sqlspec/adapters/db2/_typing.py index ad5fafdb3..8ee4a4c8a 100644 --- a/sqlspec/adapters/db2/_typing.py +++ b/sqlspec/adapters/db2/_typing.py @@ -4,38 +4,23 @@ import importlib from typing import TYPE_CHECKING, Any, cast -from sqlspec.exceptions import MissingDependencyError -from sqlspec.typing import import_optional_attr -from sqlspec.utils.module_loader import import_optional - - -class _Db2UnavailableError(Exception): - """Fallback Db2 exception base when ibm_db_dbi is not installed.""" +import ibm_db +import ibm_db_dbi +from ibm_db_dbi import AsyncConnection as Db2AsyncConnection +from ibm_db_dbi import AsyncCursor as Db2AsyncRawCursor +from ibm_db_dbi import Connection as Db2SyncConnection +from ibm_db_dbi import Cursor as Db2RawCursor +from ibm_db_dbi import Error as Db2Error +from sqlspec.exceptions import MissingDependencyError if TYPE_CHECKING: from collections.abc import Callable from types import TracebackType - from ibm_db_dbi import AsyncConnection as Db2AsyncConnection - from ibm_db_dbi import AsyncCursor as Db2AsyncRawCursor - from ibm_db_dbi import Connection as Db2SyncConnection - from ibm_db_dbi import Cursor as Db2RawCursor - from ibm_db_dbi import Error as Db2Error - from sqlspec.adapters.db2.driver import Db2AsyncDriver, Db2SyncDriver from sqlspec.core import StatementConfig -if not TYPE_CHECKING: - Db2SyncConnection = import_optional_attr("ibm_db_dbi", "Connection") or Any - Db2RawCursor = import_optional_attr("ibm_db_dbi", "Cursor") or Any - Db2AsyncConnection = import_optional_attr("ibm_db_dbi", "AsyncConnection") or Any - Db2AsyncRawCursor = import_optional_attr("ibm_db_dbi", "AsyncCursor") or Any - Db2Error = import_optional_attr("ibm_db_dbi", "Error") or _Db2UnavailableError - -ibm_db = import_optional("ibm_db") -ibm_db_dbi = import_optional("ibm_db_dbi") - __all__ = ( "Db2AsyncConnection", "Db2AsyncCursor", @@ -318,11 +303,27 @@ async def __aexit__( return None +_SYNC_DRIVER_CLS: "Callable[..., Any] | None" = None +_ASYNC_DRIVER_CLS: "Callable[..., Any] | None" = None + + def _resolve_sync_driver_type() -> "Callable[..., Any]": """Resolve Db2SyncDriver without a function-scoped import.""" - return cast("Callable[..., Any]", importlib.import_module("sqlspec.adapters.db2.driver").Db2SyncDriver) + global _SYNC_DRIVER_CLS + if _SYNC_DRIVER_CLS is None: + _SYNC_DRIVER_CLS = cast( + "Callable[..., Any]", importlib.import_module("sqlspec.adapters.db2.driver").Db2SyncDriver + ) + assert _SYNC_DRIVER_CLS is not None + return _SYNC_DRIVER_CLS def _resolve_async_driver_type() -> "Callable[..., Any]": """Resolve Db2AsyncDriver without a function-scoped import.""" - return cast("Callable[..., Any]", importlib.import_module("sqlspec.adapters.db2.driver").Db2AsyncDriver) + global _ASYNC_DRIVER_CLS + if _ASYNC_DRIVER_CLS is None: + _ASYNC_DRIVER_CLS = cast( + "Callable[..., Any]", importlib.import_module("sqlspec.adapters.db2.driver").Db2AsyncDriver + ) + assert _ASYNC_DRIVER_CLS is not None + return _ASYNC_DRIVER_CLS diff --git a/sqlspec/adapters/db2/adk/store.py b/sqlspec/adapters/db2/adk/store.py index 72dc4fad5..08fab252d 100644 --- a/sqlspec/adapters/db2/adk/store.py +++ b/sqlspec/adapters/db2/adk/store.py @@ -24,6 +24,7 @@ ) from sqlspec.extensions.adk.memory import BaseAsyncADKMemoryStore, BaseSyncADKMemoryStore, StoredMemory from sqlspec.utils.serializers import from_json, to_json +from sqlspec.utils.text import quote_identifier if TYPE_CHECKING: from sqlspec.adapters.db2.driver import Db2AsyncDriver, Db2SyncDriver @@ -39,6 +40,8 @@ class _Db2ADKStoreMixin: __slots__ = () + _max_table_name_parts: ClassVar[int] = 2 + if TYPE_CHECKING: _session_table: str _events_table: str @@ -115,7 +118,11 @@ def create_tables(self) -> None: _create_missing(driver, TABLE_EXISTS_SQL, table, ddl) for index_name, index_table, columns in self._index_specs(): _create_missing( - driver, INDEX_EXISTS_SQL, index_name, _create_index_sql(index_table, index_name, columns) + driver, + INDEX_EXISTS_SQL, + index_name, + _create_index_sql(index_table, index_name, columns), + table=index_table, ) driver.commit() @@ -356,7 +363,11 @@ async def create_tables(self) -> None: await _create_missing_async(driver, TABLE_EXISTS_SQL, table, ddl) for index_name, index_table, columns in self._index_specs(): await _create_missing_async( - driver, INDEX_EXISTS_SQL, index_name, _create_index_sql(index_table, index_name, columns) + driver, + INDEX_EXISTS_SQL, + index_name, + _create_index_sql(index_table, index_name, columns), + table=index_table, ) await driver.commit() @@ -588,6 +599,8 @@ class _Db2ADKMemoryStoreMixin: __slots__ = () + _max_table_name_parts: ClassVar[int] = 2 + if TYPE_CHECKING: _memory_table: str @@ -615,7 +628,11 @@ def create_tables(self) -> None: _create_missing(driver, TABLE_EXISTS_SQL, self._memory_table, self._memory_table_ddl()) for index_name, index_table, columns in self._memory_index_specs(): _create_missing( - driver, INDEX_EXISTS_SQL, index_name, _create_index_sql(index_table, index_name, columns) + driver, + INDEX_EXISTS_SQL, + index_name, + _create_index_sql(index_table, index_name, columns), + table=index_table, ) driver.commit() @@ -625,13 +642,29 @@ def insert_memory_entries(self, entries: list[StoredMemory], owner_id: object | if not entries: return 0 + unique_entries: list[StoredMemory] = [] + seen_batch: set[str] = set() + for entry in entries: + eid = entry["event_id"] + if eid not in seen_batch: + seen_batch.add(eid) + unique_entries.append(entry) + inserted_count = 0 insert_sql = _memory_insert_sql(self._memory_table, self._owner_id_column_name) + has_owner = self._owner_id_column_name is not None + placeholders = ", ".join("?" for _ in unique_entries) + existing_rows = self._execute_fetchall( + f"SELECT id, event_id FROM {self._memory_table} WHERE event_id IN ({placeholders})", + tuple(e["event_id"] for e in unique_entries), + ) + existing_ids = {str(row["event_id"] if "event_id" in row else row["EVENT_ID"]) for row in existing_rows} + with self._config.provide_session() as driver: - for entry in entries: - if driver.select_one_or_none(_memory_exists_sql(self._memory_table), (entry["event_id"],)) is not None: + for entry in unique_entries: + if entry["event_id"] in existing_ids: continue - driver.execute(insert_sql, _memory_insert_row(entry, self._owner_id_column_name is not None, owner_id)) + driver.execute(insert_sql, _memory_insert_row(entry, has_owner, owner_id)) inserted_count += 1 driver.commit() return inserted_count @@ -715,7 +748,11 @@ async def create_tables(self) -> None: await _create_missing_async(driver, TABLE_EXISTS_SQL, self._memory_table, await self._memory_table_ddl()) for index_name, index_table, columns in self._memory_index_specs(): await _create_missing_async( - driver, INDEX_EXISTS_SQL, index_name, _create_index_sql(index_table, index_name, columns) + driver, + INDEX_EXISTS_SQL, + index_name, + _create_index_sql(index_table, index_name, columns), + table=index_table, ) await driver.commit() @@ -725,16 +762,29 @@ async def insert_memory_entries(self, entries: list[StoredMemory], owner_id: obj if not entries: return 0 + unique_entries: list[StoredMemory] = [] + seen_batch: set[str] = set() + for entry in entries: + eid = entry["event_id"] + if eid not in seen_batch: + seen_batch.add(eid) + unique_entries.append(entry) + inserted_count = 0 insert_sql = _memory_insert_sql(self._memory_table, self._owner_id_column_name) + has_owner = self._owner_id_column_name is not None + placeholders = ", ".join("?" for _ in unique_entries) + existing_rows = await self._execute_fetchall( + f"SELECT id, event_id FROM {self._memory_table} WHERE event_id IN ({placeholders})", + tuple(e["event_id"] for e in unique_entries), + ) + existing_ids = {str(row["event_id"] if "event_id" in row else row["EVENT_ID"]) for row in existing_rows} + async with self._config.provide_session() as driver: - for entry in entries: - existing = await driver.select_one_or_none(_memory_exists_sql(self._memory_table), (entry["event_id"],)) - if existing is not None: + for entry in unique_entries: + if entry["event_id"] in existing_ids: continue - await driver.execute( - insert_sql, _memory_insert_row(entry, self._owner_id_column_name is not None, owner_id) - ) + await driver.execute(insert_sql, _memory_insert_row(entry, has_owner, owner_id)) inserted_count += 1 await driver.commit() return inserted_count @@ -893,6 +943,9 @@ def _metadata_table_ddl(table: str) -> str: def _create_index_sql(table: str, index_name: str, columns: str) -> str: + schema, _ = split_db2_table_name(table) + if schema is not None: + return f"CREATE INDEX {quote_identifier(schema)}.{quote_identifier(index_name.upper())} ON {table} ({columns})" return f"CREATE INDEX {index_name} ON {table} ({columns})" @@ -1139,7 +1192,7 @@ def _catalog_names(name: str) -> "tuple[str | None, str]": return split_db2_table_name(name) -def _create_missing(driver: "Db2SyncDriver", probe_sql: str, name: str, ddl: str) -> None: +def _create_missing(driver: "Db2SyncDriver", probe_sql: str, name: str, ddl: str, table: "str | None" = None) -> None: """Run ``ddl`` unless the catalog probe finds the object. Object names are written unquoted in the DDL, so they are probed by their upper-folded @@ -1150,12 +1203,20 @@ def _create_missing(driver: "Db2SyncDriver", probe_sql: str, name: str, ddl: str probe_sql: ``TABLE_EXISTS_SQL`` or ``INDEX_EXISTS_SQL``. name: Unquoted table or index name. ddl: Statement creating the object. + table: Optional owning table name to derive index schema qualification. """ - if driver.select_one_or_none(probe_sql, _catalog_names(name)) is None: + schema, object_name = _catalog_names(name) + if table is not None: + table_schema, _ = split_db2_table_name(table) + if table_schema is not None: + schema = table_schema + if driver.select_one_or_none(probe_sql, (schema, object_name)) is None: driver.execute(ddl) -async def _create_missing_async(driver: "Db2AsyncDriver", probe_sql: str, name: str, ddl: str) -> None: +async def _create_missing_async( + driver: "Db2AsyncDriver", probe_sql: str, name: str, ddl: str, table: "str | None" = None +) -> None: """Run ``ddl`` unless the catalog probe finds the object, on an async driver. Args: @@ -1163,8 +1224,14 @@ async def _create_missing_async(driver: "Db2AsyncDriver", probe_sql: str, name: probe_sql: ``TABLE_EXISTS_SQL`` or ``INDEX_EXISTS_SQL``. name: Unquoted table or index name. ddl: Statement creating the object. + table: Optional owning table name to derive index schema qualification. """ - if await driver.select_one_or_none(probe_sql, _catalog_names(name)) is None: + schema, object_name = _catalog_names(name) + if table is not None: + table_schema, _ = split_db2_table_name(table) + if table_schema is not None: + schema = table_schema + if await driver.select_one_or_none(probe_sql, (schema, object_name)) is None: await driver.execute(ddl) diff --git a/sqlspec/adapters/db2/core.py b/sqlspec/adapters/db2/core.py index d9c4f88b6..5285431ee 100644 --- a/sqlspec/adapters/db2/core.py +++ b/sqlspec/adapters/db2/core.py @@ -124,6 +124,9 @@ "22007": (DataError, "invalid datetime format"), "22012": (DataError, "division by zero"), "42601": (SQLParsingError, "sql syntax error"), + "42703": (SQLParsingError, "undefined column"), + "42710": (SQLParsingError, "duplicate object"), + "42884": (SQLParsingError, "routine not found"), } _SQLCODE_MAP: Final[dict[int, tuple[type[SQLSpecError], str]]] = { @@ -138,9 +141,13 @@ 1042: (DatabaseConnectionError, "unexpected system error occurred"), 30081: (DatabaseConnectionError, "communication failure"), 104: (SQLParsingError, "sql syntax error"), + 206: (SQLParsingError, "undefined column"), + 440: (SQLParsingError, "routine not found"), + 601: (SQLParsingError, "duplicate object"), 302: (DataError, "numeric or string value out of range"), 433: (DataError, "value is too long"), 952: (QueryTimeoutError, "statement execution interrupted or timed out"), + 30082: (PermissionDeniedError, "security processing failed"), } diff --git a/sqlspec/adapters/db2/driver.py b/sqlspec/adapters/db2/driver.py index 8a39a9fdd..398e3d159 100644 --- a/sqlspec/adapters/db2/driver.py +++ b/sqlspec/adapters/db2/driver.py @@ -2,6 +2,7 @@ import contextlib import logging +import sys from typing import TYPE_CHECKING, Any, cast from sqlspec.adapters.db2._typing import ( @@ -96,8 +97,13 @@ def start(self) -> None: try: cursor = cursor_manager.__enter__() handler = self._driver.handle_database_exceptions() - with handler: - cursor.execute(self._sql, normalize_execute_parameters(self._parameters)) + if sys.exc_info()[1] is None: + with handler: + cursor.execute(self._sql, normalize_execute_parameters(self._parameters)) + else: + self._driver._run_with_exception_handler( + handler, cursor.execute, self._sql, normalize_execute_parameters(self._parameters) + ) self._driver._check_pending_exception(handler) except BaseException: with contextlib.suppress(Exception): @@ -112,8 +118,11 @@ def fetch_chunk(self) -> "list[dict[str, Any]]": cursor = cursor_manager.cursor handler = self._driver.handle_database_exceptions() rows: Any = [] - with handler: - rows = cursor.fetchmany(self._chunk_size) + if sys.exc_info()[1] is None: + with handler: + rows = cursor.fetchmany(self._chunk_size) + else: + rows = self._driver._run_with_exception_handler(handler, cursor.fetchmany, self._chunk_size) self._driver._check_pending_exception(handler) if not rows: return [] @@ -305,21 +314,30 @@ def set_migration_session_schema(self, schema: str) -> None: Args: schema: Schema to make current. Unquoted all-lowercase names fold to uppercase. """ - with self.with_cursor(self.connection) as cursor: + handler = self.handle_database_exceptions() + with handler, self.with_cursor(self.connection) as cursor: if self._migration_schema_restore is None: cursor.execute("VALUES CURRENT SCHEMA") row = cursor.fetchone() self._migration_schema_restore = str(row[0]) cursor.execute(build_set_schema_sql(schema)) + self._check_pending_exception(handler) def reset_migration_session_schema(self) -> None: - """Restore the current schema captured by ``set_migration_session_schema``.""" + """Restore the current schema captured by ``set_migration_session_schema``. + + The captured schema is quoted before passing to ``build_set_schema_sql`` + to preserve lowercase or mixed-case catalog schema names from being folded + to uppercase. + """ previous_schema = self._migration_schema_restore if previous_schema is None: return self._migration_schema_restore = None - with self.with_cursor(self.connection) as cursor: + handler = self.handle_database_exceptions() + with handler, self.with_cursor(self.connection) as cursor: cursor.execute(build_set_schema_sql(quote_identifier(previous_schema))) + self._check_pending_exception(handler) def has_schema(self, schema: str) -> bool: """Return whether the schema exists in the catalog. @@ -330,9 +348,13 @@ def has_schema(self, schema: str) -> bool: Returns: True when ``SYSCAT.SCHEMATA`` lists the schema. """ - with self.with_cursor(self.connection) as cursor: + exists = False + handler = self.handle_database_exceptions() + with handler, self.with_cursor(self.connection) as cursor: cursor.execute("SELECT 1 FROM SYSCAT.SCHEMATA WHERE SCHEMANAME = ?", (normalize_identifier(schema, "db2"),)) - return cursor.fetchone() is not None + exists = cursor.fetchone() is not None + self._check_pending_exception(handler) + return exists def collect_rows(self, cursor: Any, fetched: "list[Any]") -> "tuple[list[Any], list[str], int]": column_names = resolve_column_names( @@ -637,20 +659,32 @@ async def set_migration_session_schema(self, schema: str) -> None: schema: Schema to make current. Unquoted all-lowercase names fold to uppercase. """ async with self.with_cursor(self.connection) as cursor: + handler = self.handle_database_exceptions() if self._migration_schema_restore is None: - await cursor.execute("VALUES CURRENT SCHEMA") + await self._run_with_exception_handler(handler, cursor.execute, "VALUES CURRENT SCHEMA") + self._check_pending_exception(handler) row = await cursor.fetchone() self._migration_schema_restore = str(row[0]) - await cursor.execute(build_set_schema_sql(schema)) + await self._run_with_exception_handler(handler, cursor.execute, build_set_schema_sql(schema)) + self._check_pending_exception(handler) async def reset_migration_session_schema(self) -> None: - """Restore the current schema captured by ``set_migration_session_schema``.""" + """Restore the current schema captured by ``set_migration_session_schema``. + + The captured schema is quoted before passing to ``build_set_schema_sql`` + to preserve lowercase or mixed-case catalog schema names from being folded + to uppercase. + """ previous_schema = self._migration_schema_restore if previous_schema is None: return self._migration_schema_restore = None async with self.with_cursor(self.connection) as cursor: - await cursor.execute(build_set_schema_sql(quote_identifier(previous_schema))) + handler = self.handle_database_exceptions() + await self._run_with_exception_handler( + handler, cursor.execute, build_set_schema_sql(quote_identifier(previous_schema)) + ) + self._check_pending_exception(handler) async def has_schema(self, schema: str) -> bool: """Return whether the schema exists in the catalog. @@ -662,10 +696,16 @@ async def has_schema(self, schema: str) -> bool: True when ``SYSCAT.SCHEMATA`` lists the schema. """ async with self.with_cursor(self.connection) as cursor: - await cursor.execute( - "SELECT 1 FROM SYSCAT.SCHEMATA WHERE SCHEMANAME = ?", (normalize_identifier(schema, "db2"),) + handler = self.handle_database_exceptions() + await self._run_with_exception_handler( + handler, + cursor.execute, + "SELECT 1 FROM SYSCAT.SCHEMATA WHERE SCHEMANAME = ?", + (normalize_identifier(schema, "db2"),), ) - return await cursor.fetchone() is not None + self._check_pending_exception(handler) + row = await cursor.fetchone() + return row is not None def collect_rows(self, cursor: Any, fetched: "list[Any]") -> "tuple[list[Any], list[str], int]": column_names = resolve_column_names( diff --git a/sqlspec/adapters/db2/events/__init__.py b/sqlspec/adapters/db2/events/__init__.py index 77190e284..fe1a4dda8 100644 --- a/sqlspec/adapters/db2/events/__init__.py +++ b/sqlspec/adapters/db2/events/__init__.py @@ -1,6 +1,5 @@ """IBM Db2 Events integration.""" -from sqlspec.adapters.db2.events.config import Db2EventsConfig from sqlspec.adapters.db2.events.store import Db2AsyncEventQueueStore, Db2SyncEventQueueStore -__all__ = ("Db2AsyncEventQueueStore", "Db2EventsConfig", "Db2SyncEventQueueStore") +__all__ = ("Db2AsyncEventQueueStore", "Db2SyncEventQueueStore") diff --git a/sqlspec/adapters/db2/events/config.py b/sqlspec/adapters/db2/events/config.py deleted file mode 100644 index 308d930a8..000000000 --- a/sqlspec/adapters/db2/events/config.py +++ /dev/null @@ -1,9 +0,0 @@ -"""IBM Db2 Events configuration.""" - -from sqlspec.config import EventsConfig - -__all__ = ("Db2EventsConfig",) - - -class Db2EventsConfig(EventsConfig): - """Db2 events settings for queue storage.""" diff --git a/sqlspec/adapters/db2/litestar/__init__.py b/sqlspec/adapters/db2/litestar/__init__.py index f9d3bfc0c..b8df1a043 100644 --- a/sqlspec/adapters/db2/litestar/__init__.py +++ b/sqlspec/adapters/db2/litestar/__init__.py @@ -1,6 +1,5 @@ """IBM Db2 Litestar integration.""" -from sqlspec.adapters.db2.litestar.config import Db2LitestarConfig from sqlspec.adapters.db2.litestar.store import Db2AsyncStore, Db2SyncStore -__all__ = ("Db2AsyncStore", "Db2LitestarConfig", "Db2SyncStore") +__all__ = ("Db2AsyncStore", "Db2SyncStore") diff --git a/sqlspec/adapters/db2/litestar/config.py b/sqlspec/adapters/db2/litestar/config.py deleted file mode 100644 index 0b8ad02a6..000000000 --- a/sqlspec/adapters/db2/litestar/config.py +++ /dev/null @@ -1,12 +0,0 @@ -"""IBM Db2 Litestar configuration.""" - -from sqlspec.config import LitestarConfig - -__all__ = ("Db2LitestarConfig",) - - -class Db2LitestarConfig(LitestarConfig): - """Db2-specific Litestar settings. - - Use inside ``extension_config["litestar"]`` with Db2 session store. - """ diff --git a/sqlspec/adapters/db2/litestar/store.py b/sqlspec/adapters/db2/litestar/store.py index 787fbe6c5..b688f7a3e 100644 --- a/sqlspec/adapters/db2/litestar/store.py +++ b/sqlspec/adapters/db2/litestar/store.py @@ -6,6 +6,7 @@ from sqlspec.adapters.db2.core import INDEX_EXISTS_SQL, TABLE_EXISTS_SQL, split_db2_table_name, to_db_timestamp, utc_now from sqlspec.extensions.litestar.store import BaseSQLSpecStore from sqlspec.utils.sync_tools import async_ +from sqlspec.utils.text import quote_identifier if TYPE_CHECKING: from sqlspec.adapters.db2.config import Db2AsyncConfig, Db2SyncConfig @@ -36,7 +37,7 @@ def session_index_name(table: str) -> str: Returns: The upper-folded index name. """ - return f"IX_{session_catalog_names(table)[1]}_EXP" + return f"IX_{session_catalog_names(table)[1]}_EXP".upper() def session_table_ddl(table: str) -> str: @@ -68,7 +69,11 @@ def session_index_ddl(table: str) -> str: Returns: Db2 DDL for the expiry index. """ - return f"CREATE INDEX {session_index_name(table)} ON {table}(expires_at)" + schema, _ = split_db2_table_name(table) + idx_name = session_index_name(table) + if schema is not None: + return f"CREATE INDEX {quote_identifier(schema)}.{quote_identifier(idx_name)} ON {table}(expires_at)" + return f"CREATE INDEX {idx_name} ON {table}(expires_at)" def select_session_sql(table: str) -> str: diff --git a/sqlspec/adapters/db2/pool.py b/sqlspec/adapters/db2/pool.py index 12ba77161..b63c28b5b 100644 --- a/sqlspec/adapters/db2/pool.py +++ b/sqlspec/adapters/db2/pool.py @@ -9,10 +9,11 @@ from contextlib import contextmanager from typing import TYPE_CHECKING, Any +from sqlspec.adapters.db2 import _typing as typing_module +from sqlspec.adapters.db2._typing import ibm_db_dbi from sqlspec.adapters.db2.core import build_dsn_string -from sqlspec.exceptions import ConnectionTimeoutError, DatabaseConnectionError, MissingDependencyError +from sqlspec.exceptions import ConnectionTimeoutError, DatabaseConnectionError from sqlspec.utils.logging import POOL_LOGGER_NAME, get_logger, log_with_context -from sqlspec.utils.module_loader import import_optional from sqlspec.utils.uuids import uuid4 if TYPE_CHECKING: @@ -23,7 +24,43 @@ logger = get_logger(POOL_LOGGER_NAME) _ADAPTER_NAME = "db2" -_IBM_DB_DBI: "Any | None" = None +_IBM_DB_DBI = ibm_db_dbi + + +def _set_current_schema(connection: Any, schema: str) -> None: + """Set the SQL register and SDK catalog schema without CLI keyword quoting.""" + cursor = connection.cursor() + pending_exception: BaseException | None = None + try: + cursor.execute("SET CURRENT SCHEMA = ?", (schema,)) + except BaseException as exc: + pending_exception = exc + try: + cursor.close() + except BaseException as exc: + if pending_exception is None: + pending_exception = exc + if pending_exception is not None: + raise pending_exception + connection.set_current_schema(schema) + + +async def _set_current_schema_async(connection: Any, schema: str) -> None: + """Set both schema states, closing the cursor outside the query's exception handler.""" + cursor = await connection.cursor() + pending_exception: BaseException | None = None + try: + await cursor.execute("SET CURRENT SCHEMA = ?", (schema,)) + except BaseException as exc: + pending_exception = exc + try: + await cursor.close() + except BaseException as exc: + if pending_exception is None: + pending_exception = exc + if pending_exception is not None: + raise pending_exception + await connection.set_current_schema(schema) class Db2SyncConnectionPool: @@ -61,7 +98,9 @@ def __init__( on_connection_create: Callback executed when connection is created. """ self._connection_parameters = connection_parameters - self._dsn = build_dsn_string(connection_parameters) + native_parameters = dict(connection_parameters) + native_parameters.pop("current_schema", None) + self._dsn = build_dsn_string(native_parameters) self._autocommit = bool(connection_parameters.get("autocommit", True)) self._thread_local = threading.local() self._connection_registry: set[Any] = set() @@ -82,6 +121,9 @@ def new_connection(self) -> Any: The connection opens in the pool's autocommit mode. + The schema is applied with a bound statement because CLI keyword parsing cannot preserve + semicolon-bearing schema values. The SDK catalog schema is synchronized before the hook. + The result is owned by the caller: it is not thread-local and is not tracked for pool shutdown. @@ -91,17 +133,25 @@ def new_connection(self) -> Any: Raises: MissingDependencyError: When ibm_db is not installed. """ - ibm_db_dbi = _require_ibm_db_dbi() - autocommit_mode = ibm_db_dbi.SQL_AUTOCOMMIT_ON if self._autocommit else ibm_db_dbi.SQL_AUTOCOMMIT_OFF - connection = ibm_db_dbi.connect(self._dsn, "", "", "", "", {ibm_db_dbi.SQL_ATTR_AUTOCOMMIT: autocommit_mode}) + autocommit_mode = ( + typing_module.ibm_db_dbi.SQL_AUTOCOMMIT_ON + if self._autocommit + else typing_module.ibm_db_dbi.SQL_AUTOCOMMIT_OFF + ) + connection = typing_module.ibm_db_dbi.connect( + self._dsn, "", "", "", "", {typing_module.ibm_db_dbi.SQL_ATTR_AUTOCOMMIT: autocommit_mode} + ) - if self._on_connection_create is not None: - try: + try: + schema = self._connection_parameters.get("current_schema") + if schema is not None: + _set_current_schema(connection, schema) + if self._on_connection_create is not None: self._on_connection_create(connection) - except BaseException: - with contextlib.suppress(Exception): - connection.close() - raise + except BaseException: + with contextlib.suppress(BaseException): + connection.close() + raise return connection @@ -146,7 +196,7 @@ def size(self) -> int: def checked_out(self) -> int: """Return the number of checked out connections from the perspective of this thread.""" - return 0 + return 1 if getattr(self._thread_local, "connection", None) is not None else 0 def _create_connection(self) -> Any: """Open a new connection and register it in the shutdown registry.""" @@ -337,7 +387,9 @@ def __init__( an awaitable. """ self._connection_parameters = connection_parameters - self._dsn = build_dsn_string(connection_parameters) + native_parameters = dict(connection_parameters) + native_parameters.pop("current_schema", None) + self._dsn = build_dsn_string(native_parameters) self._autocommit = bool(connection_parameters.get("autocommit", True)) self._max_size = max_size self._acquire_timeout = acquire_timeout @@ -368,25 +420,35 @@ async def new_connection(self) -> Any: The connection opens in the pool's autocommit mode and the creation hook runs on it. The result is owned by the caller and is not tracked by the pool. + The schema is applied with a bound statement and synchronized with the SDK catalog schema + before the hook, avoiding the CLI keyword parser's treatment of semicolons. + Returns: Any: A newly opened ``ibm_db_dbi.AsyncConnection``. Raises: MissingDependencyError: When ibm_db is not installed. """ - ibm_db_dbi = _require_ibm_db_dbi() - autocommit_mode = ibm_db_dbi.SQL_AUTOCOMMIT_ON if self._autocommit else ibm_db_dbi.SQL_AUTOCOMMIT_OFF - connection = await ibm_db_dbi.AsyncConnection.connect( - self._dsn, "", "", "", "", {ibm_db_dbi.SQL_ATTR_AUTOCOMMIT: autocommit_mode} + autocommit_mode = ( + typing_module.ibm_db_dbi.SQL_AUTOCOMMIT_ON + if self._autocommit + else typing_module.ibm_db_dbi.SQL_AUTOCOMMIT_OFF ) - if self._on_connection_create is not None: - try: + connection = await typing_module.ibm_db_dbi.AsyncConnection.connect( + self._dsn, "", "", "", "", {typing_module.ibm_db_dbi.SQL_ATTR_AUTOCOMMIT: autocommit_mode} + ) + try: + schema = self._connection_parameters.get("current_schema") + if schema is not None: + await _set_current_schema_async(connection, schema) + if self._on_connection_create is not None: result = self._on_connection_create(connection) if inspect.isawaitable(result): await result - except BaseException: - await self._close_connection(connection) - raise + except BaseException: + with contextlib.suppress(BaseException): + await connection.close() + raise return connection async def acquire(self) -> Any: @@ -534,36 +596,30 @@ async def _is_reusable(self, record: _Db2PooledConnection) -> bool: return True async def _is_connection_alive(self, connection: Any) -> bool: - """Ping the connection with ``SELECT 1 FROM SYSIBM.SYSDUMMY1``.""" + """Ping the connection with ``SELECT 1 FROM SYSIBM.SYSDUMMY1``. + + Close the cursor outside exception handling so compiled coroutine cleanup + does not retain the failed query as its caller's ambient exception. + """ + cursor: Any = None + pending_error: BaseException | None = None try: cursor = await connection.cursor() + await cursor.execute("SELECT 1 FROM SYSIBM.SYSDUMMY1") + await cursor.fetchone() + except BaseException as exc: + pending_error = exc + if cursor is not None: try: - await cursor.execute("SELECT 1 FROM SYSIBM.SYSDUMMY1") - await cursor.fetchone() - finally: await cursor.close() - except Exception: - return False - return True + except BaseException as exc: + if pending_error is None or isinstance(pending_error, Exception): + pending_error = exc + if pending_error is not None and not isinstance(pending_error, Exception): + raise pending_error + return pending_error is None async def _close_connection(self, connection: Any) -> None: """Close a connection, suppressing driver errors.""" with contextlib.suppress(Exception): await connection.close() - - -def _require_ibm_db_dbi() -> Any: - """Return the ``ibm_db_dbi`` module, importing it on first use. - - Returns: - Any: The ``ibm_db_dbi`` module. - - Raises: - MissingDependencyError: When ibm_db is not installed. - """ - global _IBM_DB_DBI - if _IBM_DB_DBI is None: - _IBM_DB_DBI = import_optional("ibm_db_dbi") - if _IBM_DB_DBI is None: - raise MissingDependencyError(package="ibm_db", install_package="db2") - return _IBM_DB_DBI diff --git a/sqlspec/adapters/duckdb/_typing.py b/sqlspec/adapters/duckdb/_typing.py index ab6d8232b..96193568e 100644 --- a/sqlspec/adapters/duckdb/_typing.py +++ b/sqlspec/adapters/duckdb/_typing.py @@ -9,14 +9,26 @@ import duckdb as duckdb_module from duckdb import BinderException as DuckDBBinderException from duckdb import CatalogException as DuckDBCatalogException +from duckdb import ConnectionException as DuckDBConnectionException from duckdb import ConstraintException as DuckDBConstraintException from duckdb import ConversionException as DuckDBConversionException +from duckdb import DataError as DuckDBDataError from duckdb import DuckDBPyConnection +from duckdb import InternalException as DuckDBInternalException from duckdb import InterruptException as DuckDBInterruptException +from duckdb import InvalidInputException as DuckDBInvalidInputException +from duckdb import InvalidTypeException as DuckDBInvalidTypeException from duckdb import IOException as DuckDBIOException +from duckdb import NotImplementedException as DuckDBNotImplementedException +from duckdb import OperationalError as DuckDBOperationalError +from duckdb import OutOfMemoryException as DuckDBOutOfMemoryException +from duckdb import OutOfRangeException as DuckDBOutOfRangeException from duckdb import ParserException as DuckDBParserException from duckdb import PermissionException as DuckDBPermissionException +from duckdb import SerializationException as DuckDBSerializationException +from duckdb import SyntaxException as DuckDBSyntaxException from duckdb import TransactionException as DuckDBTransactionException +from duckdb import TypeMismatchException as DuckDBTypeMismatchException _DuckDBConnection = DuckDBPyConnection @@ -36,15 +48,27 @@ "DuckDBBinderException", "DuckDBCatalogException", "DuckDBConnection", + "DuckDBConnectionException", "DuckDBConstraintException", "DuckDBConversionException", "DuckDBCursor", + "DuckDBDataError", "DuckDBIOException", + "DuckDBInternalException", "DuckDBInterruptException", + "DuckDBInvalidInputException", + "DuckDBInvalidTypeException", + "DuckDBNotImplementedException", + "DuckDBOperationalError", + "DuckDBOutOfMemoryException", + "DuckDBOutOfRangeException", "DuckDBParserException", "DuckDBPermissionException", + "DuckDBSerializationException", "DuckDBSessionContext", + "DuckDBSyntaxException", "DuckDBTransactionException", + "DuckDBTypeMismatchException", "duckdb_module", ) @@ -120,6 +144,9 @@ def __enter__(self) -> "DuckDBDriver": def __exit__( self, exc_type: "type[BaseException] | None", exc_val: "BaseException | None", exc_tb: "TracebackType | None" ) -> "bool | None": + if self._driver is not None: + self._driver.close() + self._driver = None if self._connection is not None: self._release_connection(self._connection, exc_type=exc_type, exc_val=exc_val, exc_tb=exc_tb) self._connection = None diff --git a/sqlspec/adapters/duckdb/adk/__init__.py b/sqlspec/adapters/duckdb/adk/__init__.py index b54354c1a..d7b8b0568 100644 --- a/sqlspec/adapters/duckdb/adk/__init__.py +++ b/sqlspec/adapters/duckdb/adk/__init__.py @@ -5,6 +5,6 @@ testing, and analytical workloads. """ -from sqlspec.adapters.duckdb.adk.store import DuckdbADKConfig, DuckdbADKFTSOptions, DuckdbADKMemoryStore, DuckdbADKStore +from sqlspec.adapters.duckdb.adk.store import DuckDBADKConfig, DuckDBADKFTSOptions, DuckDBADKMemoryStore, DuckDBADKStore -__all__ = ("DuckdbADKConfig", "DuckdbADKFTSOptions", "DuckdbADKMemoryStore", "DuckdbADKStore") +__all__ = ("DuckDBADKConfig", "DuckDBADKFTSOptions", "DuckDBADKMemoryStore", "DuckDBADKStore") diff --git a/sqlspec/adapters/duckdb/adk/store.py b/sqlspec/adapters/duckdb/adk/store.py index 0b634baa2..9e691e01b 100644 --- a/sqlspec/adapters/duckdb/adk/store.py +++ b/sqlspec/adapters/duckdb/adk/store.py @@ -10,15 +10,14 @@ import contextlib from collections.abc import Mapping from datetime import UTC, datetime, timedelta -from typing import TYPE_CHECKING, Any, Final, Literal, NotRequired, cast - -from typing_extensions import TypedDict +from typing import TYPE_CHECKING, Any, Final, Literal, NotRequired, TypedDict, cast from sqlspec.config import ADKConfig from sqlspec.extensions.adk import BaseSyncADKStore, StoredEvent, StoredSession, normalize_session_list_options from sqlspec.extensions.adk.memory.store import BaseSyncADKMemoryStore from sqlspec.utils.logging import get_logger from sqlspec.utils.serializers import from_json, to_json +from sqlspec.utils.text import split_qualified_identifier if TYPE_CHECKING: from collections.abc import Sequence @@ -27,7 +26,7 @@ from sqlspec.extensions.adk import SessionOrderBy, StoredMemory -__all__ = ("DuckdbADKConfig", "DuckdbADKFTSOptions", "DuckdbADKMemoryStore", "DuckdbADKStore") +__all__ = ("DuckDBADKConfig", "DuckDBADKFTSOptions", "DuckDBADKMemoryStore", "DuckDBADKStore") logger = get_logger("sqlspec.adapters.duckdb.adk.store") @@ -44,7 +43,7 @@ } -class DuckdbADKFTSOptions(TypedDict): +class DuckDBADKFTSOptions(TypedDict): """DuckDB FTS ``PRAGMA create_fts_index`` options for ADK memory search.""" stemmer: NotRequired[str] @@ -66,7 +65,7 @@ class DuckdbADKFTSOptions(TypedDict): """Whether initial FTS index creation overwrites an existing index. Defaults to ``0``.""" -class DuckdbADKConfig(ADKConfig): +class DuckDBADKConfig(ADKConfig): """DuckDB-specific ADK extension settings. Use these keys inside ``extension_config["adk"]`` with the DuckDB ADK stores. @@ -79,11 +78,36 @@ class DuckdbADKConfig(ADKConfig): enable_event_generated_column_indexes: NotRequired[bool] """Index generated event projection columns when generated columns are enabled.""" - memory_fts_options: NotRequired[DuckdbADKFTSOptions] + memory_fts_options: NotRequired[DuckDBADKFTSOptions] """DuckDB FTS ``PRAGMA create_fts_index`` options for ADK memory search.""" -class DuckdbADKStore(BaseSyncADKStore["DuckDBConfig"]): +def _duckdb_index_name(table_name: str, suffix: str) -> str: + """Generate a valid DuckDB index name for an optionally schema-qualified table.""" + parts = split_qualified_identifier(table_name.strip(), quote_chars='"') + bare_name = parts[-1] + if len(parts) > 1 and parts[0].lower() != "main": + return f"idx_{parts[0]}_{bare_name}_{suffix}" + return f"idx_{bare_name}_{suffix}" + + +def _split_duckdb_table_name(table_name: str) -> tuple[str, str]: + """Split an optionally schema-qualified table name into schema and bare table name. + + Returns: + Tuple of (schema, bare_table_name). If unqualified, schema defaults to 'main'. + """ + parts = split_qualified_identifier(table_name.strip(), quote_chars='"') + if not parts: + msg = "DuckDB table name must not be empty." + raise ValueError(msg) + first_part = parts[0] + if len(parts) > 1: + return first_part, parts[1] + return "main", first_part + + +class DuckDBADKStore(BaseSyncADKStore["DuckDBConfig"]): """DuckDB ADK store for Google Agent Development Kit. Implements session and event storage for Google Agent Development Kit @@ -101,6 +125,8 @@ class DuckdbADKStore(BaseSyncADKStore["DuckDBConfig"]): __slots__ = () + _max_table_name_parts = 2 + def __init__(self, config: "DuckDBConfig") -> None: """Initialize DuckDB ADK store. @@ -314,8 +340,8 @@ def _sessions_table_ddl(self) -> str: 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); + CREATE INDEX IF NOT EXISTS {_duckdb_index_name(self._session_table, "app_user")} ON {self._session_table}(app_name, user_id); + CREATE INDEX IF NOT EXISTS {_duckdb_index_name(self._session_table, "update_time")} ON {self._session_table}(update_time DESC); """ def _events_table_ddl(self) -> str: @@ -337,8 +363,8 @@ def _events_table_ddl(self) -> str: event_data JSON NOT NULL{generated_columns}, FOREIGN KEY (session_id) REFERENCES {self._session_table}(id) ); - CREATE INDEX IF NOT EXISTS idx_{self._events_table}_session ON {self._events_table}(session_id, timestamp ASC); - CREATE INDEX IF NOT EXISTS idx_{self._events_table}_app_timestamp ON {self._events_table}(app_name, timestamp ASC); + CREATE INDEX IF NOT EXISTS {_duckdb_index_name(self._events_table, "session")} ON {self._events_table}(session_id, timestamp ASC); + CREATE INDEX IF NOT EXISTS {_duckdb_index_name(self._events_table, "app_timestamp")} ON {self._events_table}(app_name, timestamp ASC); {generated_indexes} """ @@ -362,10 +388,10 @@ def _event_generated_column_indexes_sql(self) -> str: return "" return f""" - CREATE INDEX IF NOT EXISTS idx_{self._events_table}_author_gc + CREATE INDEX IF NOT EXISTS {_duckdb_index_name(self._events_table, "author_gc")} ON {self._events_table}(session_id, author_gc, timestamp ASC); - CREATE INDEX IF NOT EXISTS idx_{self._events_table}_node_path_gc + CREATE INDEX IF NOT EXISTS {_duckdb_index_name(self._events_table, "node_path_gc")} ON {self._events_table}(session_id, node_path_gc, timestamp ASC); """ @@ -872,7 +898,7 @@ def _set_metadata(self, key: str, value: str) -> None: conn.commit() -class DuckdbADKMemoryStore(BaseSyncADKMemoryStore["DuckDBConfig"]): +class DuckDBADKMemoryStore(BaseSyncADKMemoryStore["DuckDBConfig"]): """DuckDB ADK memory store using synchronous DuckDB driver with async wrappers. Implements memory entry storage for Google Agent Development Kit @@ -891,6 +917,8 @@ class DuckdbADKMemoryStore(BaseSyncADKMemoryStore["DuckDBConfig"]): __slots__ = () + _max_table_name_parts = 2 + def __init__(self, config: "DuckDBConfig") -> None: """Initialize DuckDB ADK memory store. @@ -931,6 +959,9 @@ def search_entries( When FTS is enabled, uses ``match_bm25()`` for BM25-ranked results. Falls back to ILIKE for simple substring matching. """ + if embedding is not None: + msg = "DuckDB ADK memory store does not support embedding-based search." + raise NotImplementedError(msg) return self._search_entries(query, app_name, user_id, limit, scope_filter) def delete_entries_by_session(self, session_id: str) -> int: @@ -1016,6 +1047,10 @@ def _memory_table_ddl(self) -> str: if self._owner_id_column_ddl: owner_id_line = f",\n {self._owner_id_column_ddl}" + idx_app_scope_user_time = _duckdb_index_name(self._memory_table, "app_scope_user_time") + idx_scope = _duckdb_index_name(self._memory_table, "scope") + idx_session = _duckdb_index_name(self._memory_table, "session") + return f""" CREATE TABLE IF NOT EXISTS {self._memory_table} ( id VARCHAR(128) PRIMARY KEY, @@ -1032,13 +1067,13 @@ def _memory_table_ddl(self) -> str: inserted_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ); - CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_app_scope_user_time + CREATE INDEX IF NOT EXISTS {idx_app_scope_user_time} ON {self._memory_table}(app_name, scope, user_id, timestamp DESC); - CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_scope + CREATE INDEX IF NOT EXISTS {idx_scope} ON {self._memory_table}(app_name, scope); - CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_session + CREATE INDEX IF NOT EXISTS {idx_session} ON {self._memory_table}(session_id); """ @@ -1162,12 +1197,14 @@ def _search_entries( use_fts = False if use_fts: + schema, bare_table = _split_duckdb_table_name(self._memory_table) + fts_macro = f"fts_{schema}_{bare_table}" where_scope, scope_params = _build_duckdb_scope_where(app_name, user_id, scope_filter, prefix="m") fts_sql = f""" SELECT m.* FROM {self._memory_table} m JOIN ( - SELECT id, fts_main_{self._memory_table}.match_bm25(id, ?, fields := 'content_text') AS score + SELECT id, {fts_macro}.match_bm25(id, ?, fields := 'content_text') AS score FROM {self._memory_table} ) fts ON m.id = fts.id WHERE {where_scope} AND fts.score IS NOT NULL @@ -1257,15 +1294,15 @@ def _raise_session_not_found(session_id: str) -> None: raise ValueError(msg) -def _adk_config(config: Any) -> DuckdbADKConfig: +def _adk_config(config: Any) -> DuckDBADKConfig: """Return DuckDB ADK extension settings from ``extension_config["adk"]``.""" extension_config = getattr(config, "extension_config", {}) if not isinstance(extension_config, dict): - return cast("DuckdbADKConfig", {}) + return cast("DuckDBADKConfig", {}) adk_config = extension_config.get("adk", {}) if not isinstance(adk_config, dict): - return cast("DuckdbADKConfig", {}) - return cast("DuckdbADKConfig", adk_config) + return cast("DuckDBADKConfig", {}) + return cast("DuckDBADKConfig", adk_config) def _validate_duckdb_fts_options(options: Mapping[str, object]) -> None: diff --git a/sqlspec/adapters/duckdb/core.py b/sqlspec/adapters/duckdb/core.py index c19a21d93..b155fc952 100644 --- a/sqlspec/adapters/duckdb/core.py +++ b/sqlspec/adapters/duckdb/core.py @@ -16,18 +16,31 @@ from sqlspec.adapters.duckdb._typing import ( DuckDBBinderException, DuckDBCatalogException, + DuckDBConnectionException, DuckDBConstraintException, DuckDBConversionException, + DuckDBDataError, + DuckDBInternalException, DuckDBInterruptException, + DuckDBInvalidInputException, + DuckDBInvalidTypeException, DuckDBIOException, + DuckDBNotImplementedException, + DuckDBOperationalError, + DuckDBOutOfMemoryException, + DuckDBOutOfRangeException, DuckDBParserException, DuckDBPermissionException, + DuckDBSerializationException, + DuckDBSyntaxException, DuckDBTransactionException, + DuckDBTypeMismatchException, duckdb_module, ) from sqlspec.core import DriverParameterProfile, ParameterStyle, StatementConfig, build_statement_config_from_profile from sqlspec.exceptions import ( CheckViolationError, + DatabaseConnectionError, DataError, ForeignKeyViolationError, ImproperConfigurationError, @@ -512,11 +525,23 @@ def configure_duckdb_statement_profiling(connection: Any, profile_path: Path) -> _EXCEPTION_MAPPING: Final[dict[type[BaseException], tuple[type[SQLSpecError], str]]] = { DuckDBCatalogException: (NotFoundError, "catalog error"), DuckDBParserException: (SQLParsingError, "SQL parsing error"), + DuckDBSyntaxException: (SQLParsingError, "SQL syntax error"), DuckDBBinderException: (SQLParsingError, "SQL parsing error"), DuckDBPermissionException: (PermissionDeniedError, "permission denied"), DuckDBInterruptException: (OperationCancelledError, "query interrupted"), DuckDBIOException: (OperationalError, "operational error"), DuckDBConversionException: (DataError, "data error"), + DuckDBOutOfRangeException: (DataError, "out of range error"), + DuckDBInvalidInputException: (DataError, "invalid input error"), + DuckDBTypeMismatchException: (DataError, "type mismatch error"), + DuckDBInvalidTypeException: (DataError, "invalid type error"), + DuckDBDataError: (DataError, "data error"), + DuckDBConnectionException: (DatabaseConnectionError, "connection error"), + DuckDBOperationalError: (OperationalError, "operational error"), + DuckDBOutOfMemoryException: (OperationalError, "out of memory error"), + DuckDBInternalException: (OperationalError, "internal error"), + DuckDBSerializationException: (SerializationConflictError, "serialization conflict"), + DuckDBNotImplementedException: (SQLSpecError, "not implemented"), } _EXCEPTION_MAPPING_CACHE: Final[dict[type[BaseException], tuple[type[SQLSpecError], str]]] = {} @@ -594,16 +619,28 @@ def build_statement_config(*, json_serializer: "Callable[[Any], str] | None" = N ) +_DUCKDB_VERSION_RE: Final[re.Pattern[str]] = re.compile(r"^v?(\d+)(?:\.(\d+))?(?:\.(\d+))?") + + +def _parse_duckdb_version(raw_version: str) -> tuple[int, int, int]: + """Parse installed DuckDB version string into a 3-tuple of integers.""" + match = _DUCKDB_VERSION_RE.match(raw_version.strip()) + if match is None: + return (0, 0, 0) + major = int(match.group(1)) if match.group(1) else 0 + minor = int(match.group(2)) if match.group(2) else 0 + patch = int(match.group(3)) if match.group(3) else 0 + return (major, minor, patch) + + +_INSTALLED_DUCKDB_VERSION: Final[tuple[int, int, int]] = _parse_duckdb_version( + str(getattr(duckdb_module, "__version__", "")) +) + + def duckdb_version_tuple() -> "tuple[int, int, int]": """Return the installed DuckDB version as a 3-tuple of integers.""" - try: - raw = getattr(duckdb_module, "__version__", "") - parts = [int(part) for part in raw.split(".") if part.isdigit()] - while len(parts) < _VERSION_COMPONENTS: - parts.append(0) - return parts[0], parts[1], parts[2] - except Exception: - return (0, 0, 0) + return _INSTALLED_DUCKDB_VERSION _DUCKDB_COPY_FILENAME_FOLLOWS_QUERY: Final[bool] = duckdb_version_tuple() >= (2, 0, 0) @@ -802,6 +839,11 @@ def _resolve_native_storage_target( return target +def _secret_scope_len(secret: "Mapping[str, Any]") -> int: + """Return the length of the secret scope string for sorting.""" + return len(secret.get("scope") or "") + + def _native_storage_eligible( uri: str, protocol: str, @@ -841,9 +883,9 @@ def _native_storage_eligible( ] if not matching: return False - matching.sort(key=lambda secret: len(secret.get("scope") or ""), reverse=True) + matching.sort(key=_secret_scope_len, reverse=True) selected = matching[0] - if len(matching) > 1 and len(selected.get("scope") or "") == len(matching[1].get("scope") or ""): + if len(matching) > 1 and _secret_scope_len(selected) == _secret_scope_len(matching[1]): return False provider = selected.get("provider", "config").lower() native = {key.lower(): value for key, value in selected.get("value", {}).items()} @@ -1022,7 +1064,14 @@ def _classify_duckdb_constraint(error: "BaseException") -> SQLSpecError: def _map_transaction_exception(error: "BaseException") -> SQLSpecError: - if "conflict on update" in str(error).lower(): + lowered = str(error).lower() + if ( + "conflict on update" in lowered + or "conflict on tuple deletion" in lowered + or "write-write conflict" in lowered + or "write conflict" in lowered + or "serialization" in lowered + ): return _create_duckdb_error(error, SerializationConflictError, "serialization conflict") return _create_duckdb_error(error, OperationalError, "transaction error") diff --git a/sqlspec/adapters/duckdb/data_dictionary.py b/sqlspec/adapters/duckdb/data_dictionary.py index f99a27405..0af956dc2 100644 --- a/sqlspec/adapters/duckdb/data_dictionary.py +++ b/sqlspec/adapters/duckdb/data_dictionary.py @@ -1,6 +1,6 @@ """DuckDB-specific data dictionary for metadata queries.""" -import re +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, ClassVar, Final, cast from mypy_extensions import mypyc_attr @@ -14,24 +14,76 @@ DependencyMetadata, ForeignKeyMetadata, IndexMetadata, + MetadataCapability, + MetadataCapabilityProfile, MetadataFidelity, MetadataResult, MetadataSource, MetadataSupport, ObjectIdentity, + SystemMetadataCapability, + SystemMetadataRequest, + SystemMetadataResult, TableMetadata, VersionInfo, + ensure_system_metadata_request, get_data_dictionary_loader, + system_metadata_gated_result, + unsupported_system_metadata_capability, ) from sqlspec.driver import SyncDataDictionaryBase if TYPE_CHECKING: - from collections.abc import Mapping - from sqlspec.adapters.duckdb.driver import DuckDBDriver __all__ = ("DuckDBDataDictionary",) +_DUCKDB_METADATA_DOMAINS: Final[tuple[str, ...]] = ( + "schemas", + "objects", + "tables", + "columns", + "constraints", + "indexes", + "views", + "routines", + "databases", + "extensions", + "system", + "dependencies", + "types", + "sequences", + "ddl", + "privileges", +) + +_DUCKDB_SUPPORTED_DOMAINS: Final[frozenset[str]] = frozenset(d for d in _DUCKDB_METADATA_DOMAINS if d != "privileges") + +_DUCKDB_SYSTEM_METADATA_QUERIES: Final[dict[str, str]] = {"logs": "logs", "memory": "memory", "settings": "settings"} + + +def _duckdb_capability_for_domain(domain: str) -> MetadataCapability: + if domain in _DUCKDB_SUPPORTED_DOMAINS: + return MetadataCapability( + domain=domain, + support=MetadataSupport.SUPPORTED, + fidelity=MetadataFidelity.NATIVE, + source=MetadataSource.CATALOG, + ) + return MetadataCapability.unsupported(domain) + + +def _duckdb_system_metadata_capability(domain: str) -> SystemMetadataCapability: + if domain not in _DUCKDB_SYSTEM_METADATA_QUERIES: + return unsupported_system_metadata_capability(domain) + return SystemMetadataCapability( + domain, MetadataSupport.SUPPORTED, fidelity=MetadataFidelity.NATIVE, source=MetadataSource.SYSTEM_VIEW + ) + + +def _rows_as_mappings(rows: "list[Any] | tuple[Any, ...]") -> "tuple[Mapping[str, object], ...]": + return tuple(dict(row) if not isinstance(row, Mapping) else row for row in rows) + @mypyc_attr(allow_interpreted_subclasses=True, native_class=False) class DuckDBDataDictionary(SyncDataDictionaryBase): @@ -42,6 +94,81 @@ class DuckDBDataDictionary(SyncDataDictionaryBase): def __init__(self) -> None: super().__init__() + def get_metadata_capabilities( + self, driver: "DuckDBDriver", domains: "Sequence[str] | None" = None + ) -> MetadataCapabilityProfile: + """Report DuckDB data-dictionary capability profile.""" + _ = driver + requested_domains = tuple(domains) if domains is not None else _DUCKDB_METADATA_DOMAINS + capabilities = tuple(_duckdb_capability_for_domain(domain) for domain in requested_domains) + return MetadataCapabilityProfile(self.dialect, adapter=type(self).__name__, capabilities=capabilities) + + def get_system_metadata_capabilities( + self, driver: "DuckDBDriver", domains: "Sequence[str] | None" = None + ) -> "tuple[SystemMetadataCapability, ...]": + """Get DuckDB system metadata capability disclosures.""" + _ = driver + requested_domains = tuple(_DUCKDB_SYSTEM_METADATA_QUERIES) if domains is None else tuple(domains) + return tuple(_duckdb_system_metadata_capability(domain) for domain in requested_domains) + + def get_schemas(self, driver: "DuckDBDriver") -> MetadataResult: + """Get DuckDB schema metadata.""" + query = self.get_query("schemas", "list") + rows = driver.select(query) + return MetadataResult("schemas", capability=_duckdb_capability_for_domain("schemas"), items=tuple(rows)) + + def get_objects(self, driver: "DuckDBDriver", schema: "str | None" = None) -> MetadataResult: + """Get database object metadata.""" + query = self.get_query("objects", "by_schema") + schema_name = self.resolve_schema(schema) + rows = driver.select(query, schema_name=schema_name) + return MetadataResult("objects", capability=_duckdb_capability_for_domain("objects"), items=tuple(rows)) + + def get_routines(self, driver: "DuckDBDriver", schema: "str | None" = None) -> MetadataResult: + """Get routine and function metadata.""" + query = self.get_query("functions", "by_schema") + schema_name = self.resolve_schema(schema) + rows = driver.select(query, schema_name=schema_name) + return MetadataResult("routines", capability=_duckdb_capability_for_domain("routines"), items=tuple(rows)) + + def get_databases(self, driver: "DuckDBDriver") -> MetadataResult: + """Get DuckDB databases metadata.""" + query = self.get_query("databases", "list") + rows = driver.select(query) + return MetadataResult("databases", capability=_duckdb_capability_for_domain("databases"), items=tuple(rows)) + + def get_extensions(self, driver: "DuckDBDriver") -> MetadataResult: + """Get DuckDB extensions metadata.""" + query = self.get_query("extensions", "list") + rows = driver.select(query) + return MetadataResult("extensions", capability=_duckdb_capability_for_domain("extensions"), items=tuple(rows)) + + def get_table_details(self, driver: "DuckDBDriver", table: str, schema: "str | None" = None) -> MetadataResult: + """Get rich table metadata.""" + query = self.get_query("tables", "by_schema") + schema_name = self.resolve_schema(schema) + rows = driver.select(query, schema_name=schema_name, table_name=table) + return MetadataResult("tables", capability=_duckdb_capability_for_domain("tables"), items=tuple(rows)) + + def get_system_metadata( + self, driver: "DuckDBDriver", request: "SystemMetadataRequest | str | None" = None, **kwargs: Any + ) -> SystemMetadataResult: + """Get opt-in DuckDB system metadata.""" + metadata_request = ensure_system_metadata_request(request, **kwargs) + capability = _duckdb_system_metadata_capability(metadata_request.domain) + gate_result = system_metadata_gated_result(metadata_request, capability) + if gate_result.capability.support != MetadataSupport.SUPPORTED: + return gate_result + + query_name = _DUCKDB_SYSTEM_METADATA_QUERIES.get(metadata_request.domain) + if query_name is None: + return gate_result + query = self.get_query("system", query_name) + rows = driver.select(query) + return SystemMetadataResult.from_rows( + metadata_request, capability, rows=_rows_as_mappings(rows), source=MetadataSource.SYSTEM_VIEW + ) + def get_version(self, driver: "DuckDBDriver") -> "VersionInfo | None": """Get DuckDB database version information. @@ -101,7 +228,9 @@ def get_tables(self, driver: "DuckDBDriver", schema: "str | None" = None) -> "li """Get tables sorted by topological dependency order using DuckDB catalog.""" schema_name = self.resolve_schema(schema) self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="tables") - return driver.select(self.get_query("tables", "by_schema"), schema_name=schema_name, schema_type=TableMetadata) + return driver.select( + self.get_query("tables", "by_schema"), schema_name=schema_name, table_name=None, schema_type=TableMetadata + ) def get_columns( self, driver: "DuckDBDriver", table: "str | None" = None, schema: "str | None" = None @@ -300,7 +429,8 @@ def _get_native_schema_ddl( source=MetadataSource.CATALOG, ) for ref_key, ref_identity in candidates.items() - if ref_key != item.identity.name.lower() and _references_identifier(ddl_text, ref_identity.name) + if ref_key != item.identity.name.lower() + and self._references_identifier(ddl_text, ref_identity.name) ) linked.append( DDLResult( @@ -421,8 +551,3 @@ def _duckdb_index_columns(value: Any) -> list[str]: "sequence": DependencyEdgeKind.SEQUENCE_USE, "type": DependencyEdgeKind.DEFAULT_EXPRESSION, } - - -def _references_identifier(sql_text: str, name: str) -> bool: - """Return True when ``name`` appears in ``sql_text`` as a whole identifier.""" - return re.search(rf"(? None: + """Close driver and clean up temporary statement profile files.""" + if self._created_profile_path is not None: + cleanup_duckdb_profile_path(self._created_profile_path) + self._created_profile_path = None + + def __del__(self) -> None: + self.close() + @property def data_dictionary(self) -> "DuckDBDataDictionary": """Get the data dictionary for this driver. @@ -202,8 +209,12 @@ def dispatch_execute_many(self, cursor: "DuckDBConnection", statement: SQL) -> " if prepared_parameters: parsed_expression = cached_statement.expression or statement._peek_expression() if isinstance(parsed_expression, exp.Insert) and not parsed_expression.args.get("returning"): - parsed_expression.meta["compiled_sql"] = sql - bulk_result = self._execute_bulk_insert_many(parsed_expression, prepared_parameters) + try: + bulk_result = self._execute_bulk_insert_many( + parsed_expression, prepared_parameters, compiled_sql=sql + ) + except TypeError: + bulk_result = self._execute_bulk_insert_many(parsed_expression, prepared_parameters) if bulk_result is not None: return self.create_execution_result( cursor, @@ -293,27 +304,26 @@ def dispatch_select_stream(self, statement: "SQL", chunk_size: int) -> "SyncRowS def begin(self) -> None: """Begin a database transaction.""" - try: - self.connection.execute("BEGIN TRANSACTION") - except duckdb.Error as e: - raise create_mapped_exception(e) from e + handler = self.handle_database_exceptions() + self._run_with_exception_handler(handler, self.connection.execute, "BEGIN TRANSACTION") + self._check_pending_exception(handler) self._transaction_active = True def commit(self) -> None: """Commit the current transaction.""" + handler = self.handle_database_exceptions() try: - self.connection.commit() - except duckdb.Error as e: - raise create_mapped_exception(e) from e + self._run_with_exception_handler(handler, self.connection.commit) + self._check_pending_exception(handler) finally: self._transaction_active = False def rollback(self) -> None: """Rollback the current transaction.""" + handler = self.handle_database_exceptions() try: - self.connection.rollback() - except duckdb.Error as e: - raise create_mapped_exception(e) from e + self._run_with_exception_handler(handler, self.connection.rollback) + self._check_pending_exception(handler) finally: self._transaction_active = False @@ -346,14 +356,25 @@ def rollback_to_savepoint(self, name: str) -> None: def set_migration_session_schema(self, schema: str) -> None: """Set DuckDB search_path for migration SQL.""" - self.connection.execute(f"SET search_path = {quote_identifier(schema)}") + handler = self.handle_database_exceptions() + self._run_with_exception_handler( + handler, self.connection.execute, f"SET search_path = {quote_identifier(schema)}" + ) + self._check_pending_exception(handler) def has_schema(self, schema: str) -> bool: """Return whether a DuckDB schema exists.""" - result = self.connection.execute( - "SELECT 1 FROM information_schema.schemata WHERE schema_name = ?", [schema] - ).fetchone() - return result is not None + + def _check_schema() -> bool: + result = self.connection.execute( + "SELECT 1 FROM information_schema.schemata WHERE schema_name = ?", [schema] + ).fetchone() + return result is not None + + handler = self.handle_database_exceptions() + res = self._run_with_exception_handler(handler, _check_schema) + self._check_pending_exception(handler) + return bool(res) def with_cursor(self, connection: "DuckDBConnection") -> "DuckDBCursor": """Create context manager for DuckDB cursor. @@ -498,7 +519,7 @@ def select_to_arrow( if arrow_result is None: msg = "Unreachable" - raise RuntimeError(msg) # pragma: no cover + raise RuntimeError(msg) return arrow_result @@ -599,7 +620,8 @@ def load_from_arrow( source_data = source.get_data() if hasattr(source, "get_data") else source arrow_table = None arrow_source: object - if isinstance(source_data, pa.RecordBatchReader): + is_reader = hasattr(source_data, "read_next_batch") + if is_reader: arrow_source = source_data else: arrow_table = self._coerce_arrow_table(source_data) @@ -631,7 +653,7 @@ def load_from_arrow( raise self._check_pending_exception(exc_handler) - if isinstance(source_data, pa.RecordBatchReader): + if is_reader: telemetry_payload: StorageTelemetry = { "rows_processed": inserted_rows, "bytes_processed": 0, @@ -726,6 +748,7 @@ def _ensure_statement_profiling(self, cursor: Any) -> None: profile_path = create_duckdb_profile_path() configure_duckdb_statement_profiling(cursor, profile_path) self.driver_features["_duckdb_profile_path"] = str(profile_path) + self._created_profile_path = profile_path def _capture_statement_profile(self, executed_sql: str) -> "dict[str, Any] | None": """Read the latest DuckDB JSON profile and stage it for StatementEvent emission.""" @@ -741,14 +764,15 @@ def _execute_and_materialize_select( hint = self._cache_holder.get_row_count_hint(cache_key) sql_method = getattr(cursor, "sql", None) if ( - pa is not None + PYARROW_INSTALLED and (hint is None or hint >= ARROW_MIN_ROWS) and not profiling_enabled and callable(sql_method) ): relation: Any = sql_method(wire_sql, params=parameters or None) if relation is None: - self._cache_holder.set_row_count_hint(cache_key, 0) + if hint != 0: + self._cache_holder.set_row_count_hint(cache_key, 0) return [], [], 0, "tuple" try: description = relation.description @@ -756,14 +780,16 @@ def _execute_and_materialize_select( fetched_data = relation.fetchall() data, column_names = collect_rows(fetched_data, description) row_count = len(data) - self._cache_holder.set_row_count_hint(cache_key, row_count) + if hint != row_count: + self._cache_holder.set_row_count_hint(cache_key, row_count) return data, column_names, row_count, "tuple" arrow_table = relation.to_arrow_table() data = arrow_table.to_pylist() _restore_uuid_columns(data, description) column_names = list(arrow_table.column_names) row_count = len(data) - self._cache_holder.set_row_count_hint(cache_key, row_count) + if hint != row_count: + self._cache_holder.set_row_count_hint(cache_key, row_count) return data, column_names, row_count, "dict" finally: close_method = getattr(relation, "close", None) @@ -779,10 +805,12 @@ def _execute_and_materialize_select( def _materialize_fetched_rows(self, cursor: Any, cache_key: str) -> "tuple[list[Any], list[str], int, RowFormat]": """Materialize already-executed cursor rows via fetchall.""" description = cursor.description + hint = self._cache_holder.get_row_count_hint(cache_key) fetched_data = cursor.fetchall() data, column_names = collect_rows(fetched_data, description) row_count = len(data) - self._cache_holder.set_row_count_hint(cache_key, row_count) + if hint != row_count: + self._cache_holder.set_row_count_hint(cache_key, row_count) return data, column_names, row_count, "tuple" def _materialize_select_rows(self, cursor: Any, compiled_sql: str) -> "tuple[list[Any], list[str], int, RowFormat]": @@ -790,7 +818,7 @@ def _materialize_select_rows(self, cursor: Any, compiled_sql: str) -> "tuple[lis description = cursor.description hint = self._cache_holder.get_row_count_hint(compiled_sql) if ( - pa is None + not PYARROW_INSTALLED or (hint is not None and hint < ARROW_MIN_ROWS) or description_has_arrow_divergent_types(description) or not hasattr(cursor, "to_arrow_table") @@ -803,7 +831,8 @@ def _materialize_select_rows(self, cursor: Any, compiled_sql: str) -> "tuple[lis _restore_uuid_columns(data, description) column_names = list(arrow_table.column_names) row_count = len(data) - self._cache_holder.set_row_count_hint(compiled_sql, row_count) + if hint != row_count: + self._cache_holder.set_row_count_hint(compiled_sql, row_count) return data, column_names, row_count, "dict" def _native_storage_job( @@ -906,7 +935,7 @@ def _execute_bulk_insert_many( @staticmethod def _build_arrow_table(rows: "list[Any]", column_names: "list[str]") -> Any | None: """Build a pyarrow table from batch rows when they share a stable shape.""" - if not rows or pa is None: + if not rows or not PYARROW_INSTALLED: return None first_row = rows[0] @@ -917,6 +946,8 @@ def _build_arrow_table(rows: "list[Any]", column_names: "list[str]") -> Any | No if any(not isinstance(row, (list, tuple)) or len(row) != len(column_names) for row in rows): return None try: + import pyarrow as pa + return pa.Table.from_arrays([pa.array(col) for col in zip(*rows, strict=True)], names=column_names) except Exception: return None @@ -948,23 +979,35 @@ def _open_stream_reader(self, sql: str, parameters: Any, chunk_size: int) -> "tu """Open a native Arrow reader through DuckDB's exception boundary.""" ensure_pyarrow() handler = self.handle_database_exceptions() - description: list[Any] | None = None - reader: Any | None = None stream_cursor, close_cursor = self._acquire_stream_cursor() - try: - with handler: - sql_method = getattr(stream_cursor, "sql", None) - if callable(sql_method): - relation: Any = sql_method(sql, params=parameters or None) - description = relation.description if relation is not None else None + close_cb = stream_cursor.close if close_cursor else None + + def _fetch_reader() -> "tuple[Any, list[Any] | None]": + reader: Any = None + sql_method = getattr(stream_cursor, "sql", None) + if callable(sql_method): + relation: Any = sql_method(sql, params=parameters or None) + description = relation.description if relation is not None else None + if description and description_has_arrow_divergent_types(description): + cursor_result = stream_cursor.execute(sql, normalize_execute_parameters(parameters)) + reader = _DuckDBFetchManyReader(cursor_result, chunk_size, description, close_callback=close_cb) + else: reader = relation.to_arrow_reader(chunk_size) if relation is not None else None + else: + result = stream_cursor.execute(sql, normalize_execute_parameters(parameters)) + description = result.description + if description and description_has_arrow_divergent_types(description): + reader = _DuckDBFetchManyReader(result, chunk_size, description, close_callback=close_cb) else: - result = stream_cursor.execute(sql, normalize_execute_parameters(parameters)) - description = result.description reader = result.to_arrow_reader(chunk_size) + return reader, description + + try: + reader_result = self._run_with_exception_handler(handler, _fetch_reader) self._check_pending_exception(handler) + reader, description = _ensure_stream_reader_result(reader_result) reader = _ensure_arrow_reader(reader) - if close_cursor: + if close_cursor and not isinstance(reader, _DuckDBFetchManyReader): reader = arrow_reader_with_deferred_close(reader, stream_cursor.close) except BaseException: if close_cursor: @@ -977,6 +1020,53 @@ def _open_stream_reader(self, sql: str, parameters: Any, chunk_size: int) -> "tu register_driver_profile("duckdb", driver_profile) +class _DuckDBFetchManyBatch: + __slots__ = ("_rows",) + + def __init__(self, rows: list[dict[str, Any]]) -> None: + self._rows = rows + + def to_pylist(self) -> list[dict[str, Any]]: + return self._rows + + +class _DuckDBFetchManyReader: + __slots__ = ("_chunk_size", "_close_callback", "_column_names", "_description", "_target") + + def __init__( + self, target: Any, chunk_size: int, description: list[Any], close_callback: "Callable[[], None] | None" = None + ) -> None: + self._target = target + self._chunk_size = chunk_size + self._description = description + self._column_names = [col[0] for col in description] + self._close_callback = close_callback + + def read_next_batch(self) -> _DuckDBFetchManyBatch: + target = self._target + if target is None: + raise StopIteration + rows = target.fetchmany(self._chunk_size) + if not rows: + self.close() + raise StopIteration + dict_rows = [dict(zip(self._column_names, row, strict=False)) for row in rows] + return _DuckDBFetchManyBatch(dict_rows) + + def close(self) -> None: + target = self._target + self._target = None + if self._close_callback is not None: + callback = self._close_callback + self._close_callback = None + callback() + elif target is not None: + # DuckDB 2 keeps execute results active until they are completed. + complete_method = getattr(target, "complete", None) + if callable(complete_method): + complete_method() + + def _ensure_arrow_reader(reader: Any) -> Any: if reader is None: msg = "DuckDB did not return an Arrow record batch reader." @@ -984,6 +1074,13 @@ def _ensure_arrow_reader(reader: Any) -> Any: return reader +def _ensure_stream_reader_result(result: "tuple[Any, list[Any] | None] | None") -> "tuple[Any, list[Any] | None]": + if result is None: + msg = "DuckDB did not return an Arrow record batch reader." + raise SQLSpecError(msg) + return result + + def _is_plain_values_insert(expression: exp.Insert, column_count: int) -> bool: if any(expression.args.get(key) for key in ("conflict", "alternative", "with_")): return False diff --git a/sqlspec/adapters/duckdb/litestar/__init__.py b/sqlspec/adapters/duckdb/litestar/__init__.py index c6b17c019..7b9ea3253 100644 --- a/sqlspec/adapters/duckdb/litestar/__init__.py +++ b/sqlspec/adapters/duckdb/litestar/__init__.py @@ -1,5 +1,5 @@ """Litestar integration for DuckDB adapter.""" -from sqlspec.adapters.duckdb.litestar.store import DuckdbStore +from sqlspec.adapters.duckdb.litestar.store import DuckDBStore -__all__ = ("DuckdbStore",) +__all__ = ("DuckDBStore",) diff --git a/sqlspec/adapters/duckdb/litestar/store.py b/sqlspec/adapters/duckdb/litestar/store.py index 3031bf7c6..421aacf91 100644 --- a/sqlspec/adapters/duckdb/litestar/store.py +++ b/sqlspec/adapters/duckdb/litestar/store.py @@ -10,10 +10,10 @@ from sqlspec.adapters.duckdb.config import DuckDBConfig -__all__ = ("DuckdbStore",) +__all__ = ("DuckDBStore",) -class DuckdbStore(BaseSQLSpecStore["DuckDBConfig"]): +class DuckDBStore(BaseSQLSpecStore["DuckDBConfig"]): """DuckDB session store using synchronous DuckDB driver. Implements server-side session storage for Litestar using DuckDB diff --git a/sqlspec/adapters/duckdb/pool.py b/sqlspec/adapters/duckdb/pool.py index 1d5d9084d..3b93801d0 100644 --- a/sqlspec/adapters/duckdb/pool.py +++ b/sqlspec/adapters/duckdb/pool.py @@ -31,9 +31,6 @@ logger = get_logger(POOL_LOGGER_NAME) _ADAPTER_NAME = "duckdb" -DEFAULT_MIN_POOL: Final[int] = 1 -DEFAULT_MAX_POOL: Final[int] = 4 -POOL_TIMEOUT: Final[float] = 30.0 POOL_RECYCLE: Final[int] = 86400 HEALTH_CHECK_INTERVAL: Final[float] = 30.0 SECRET_CONFLICT_ATTEMPTS: Final[int] = 5 @@ -77,6 +74,7 @@ class DuckDBConnectionPool: "_enable_statement_profiling", "_extensions", "_generation", + "_has_dynamic_secrets", "_health_check_interval", "_installed_signatures", "_is_memory_db", @@ -115,6 +113,7 @@ def __init__( self._health_check_interval = health_check_interval self._extensions = extensions or [] self._secrets = secrets or [] + self._has_dynamic_secrets: bool = any(callable(secret.get("value_provider")) for secret in self._secrets) self._on_connection_create = on_connection_create self._enable_statement_profiling = enable_statement_profiling self._profile_paths: dict[int, Path] = {} @@ -356,6 +355,9 @@ def _storage_settings(self, connection: DuckDBConnection) -> "dict[str, Any]": if profile_path is not None: return {"_duckdb_profile_path": str(profile_path)} return {} + cached_settings = cast("dict[str, Any] | None", state.get("_cached_storage_settings")) + if cached_settings is not None: + return cached_settings settings: dict[str, Any] = { "_duckdb_storage_extensions": state.get("storage_extensions", frozenset()), "_duckdb_storage_secrets": state.get("storage_secrets", ()), @@ -364,6 +366,7 @@ def _storage_settings(self, connection: DuckDBConnection) -> "dict[str, Any]": profile_path = state.get("profile_path") if profile_path is not None: settings["_duckdb_profile_path"] = str(profile_path) + state["_cached_storage_settings"] = settings return settings def _install_extension_once( @@ -452,7 +455,7 @@ def _get_thread_connection(self) -> DuckDBConnection: def _refresh_dynamic_secrets(self, connection: DuckDBConnection) -> None: """Refresh secrets configured with ``value_provider`` when ``refresh_interval_seconds`` has elapsed.""" - if not any(callable(secret.get("value_provider")) for secret in self._secrets): + if not self._has_dynamic_secrets: return thread_state = self._thread_local.__dict__ refresh_times = cast("dict[str, float]", thread_state.setdefault("secret_refresh_times", {})) @@ -483,6 +486,7 @@ def _refresh_dynamic_secrets(self, connection: DuckDBConnection) -> None: updated = True if updated: self._thread_local.storage_secrets = tuple(current_secrets) + thread_state.pop("_cached_storage_settings", None) def _retire_connection(self, connection: DuckDBConnection) -> None: """Close a pool-owned connection and drop it from the shutdown registry.""" @@ -500,6 +504,7 @@ def _close_thread_connection(self) -> None: thread_state.pop("storage_secrets", None) thread_state.pop("storage_protocols", None) thread_state.pop("profile_path", None) + thread_state.pop("_cached_storage_settings", None) if "connection" in thread_state: self._retire_connection(cast("DuckDBConnection", self._thread_local.connection)) del self._thread_local.connection diff --git a/sqlspec/adapters/mssql_python/_typing.py b/sqlspec/adapters/mssql_python/_typing.py index d5fc593ba..1102b0522 100644 --- a/sqlspec/adapters/mssql_python/_typing.py +++ b/sqlspec/adapters/mssql_python/_typing.py @@ -23,8 +23,7 @@ MssqlPythonConnection: TypeAlias = Connection MssqlPythonRawCursor: TypeAlias = Cursor - -if not TYPE_CHECKING: +else: MssqlPythonConnection = Connection MssqlPythonRawCursor = Cursor diff --git a/sqlspec/adapters/mssql_python/adk/store.py b/sqlspec/adapters/mssql_python/adk/store.py index 8340300d9..61f24f058 100644 --- a/sqlspec/adapters/mssql_python/adk/store.py +++ b/sqlspec/adapters/mssql_python/adk/store.py @@ -13,6 +13,7 @@ ) from sqlspec.adapters.mssql_python.data_dictionary import MssqlVersionInfo from sqlspec.config import ADKConfig +from sqlspec.exceptions import ImproperConfigurationError from sqlspec.extensions.adk import BaseSyncADKStore, StoredEvent, StoredSession, normalize_session_list_options from sqlspec.extensions.adk.memory.store import BaseSyncADKMemoryStore from sqlspec.utils.serializers import from_json, to_json @@ -29,8 +30,6 @@ __all__ = ("MssqlPythonADKConfig", "MssqlPythonADKMemoryStore", "MssqlPythonADKStore") MSSQL_TABLE_NOT_FOUND_ERROR: Final[int] = 208 -MSSQL_DUPLICATE_OBJECT_ERROR: Final[int] = 2714 -MSSQL_DUPLICATE_INDEX_ERROR: Final[int] = 1913 MSSQL_SCHEMA: Final[str] = "dbo" JSON_FALLBACK_COLUMN_TYPE: Final[str] = "NVARCHAR(MAX)" JSON_NATIVE_COLUMN_TYPE: Final[str] = "JSON" @@ -49,6 +48,8 @@ class MssqlPythonADKStore(BaseSyncADKStore["MssqlPythonConfig"]): connector_name: ClassVar[str] = "mssql_python" __slots__ = ("_json_column_type", "_native_json") + _max_table_name_parts = 2 + def __init__(self, config: "MssqlPythonConfig") -> None: super().__init__(config) adk_config = _adk_config(config) @@ -433,6 +434,8 @@ class MssqlPythonADKMemoryStore(BaseSyncADKMemoryStore["MssqlPythonConfig"]): __slots__ = ("_json_column_type", "_native_json") + _max_table_name_parts = 2 + def __init__(self, config: "MssqlPythonConfig") -> None: super().__init__(config) adk_config = _adk_config(config) @@ -485,26 +488,30 @@ def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object ); """ inserted = 0 + batch_params: list[tuple[Any, ...]] = [] + for entry in entries: + params: tuple[Any, ...] = ( + entry["id"], + entry["session_id"], + entry["app_name"], + entry["user_id"], + entry.get("scope", "user"), + entry["event_id"], + entry.get("author"), + entry["timestamp"], + to_json(entry["content_json"]), + entry["content_text"], + to_json(entry["metadata_json"]) if entry.get("metadata_json") is not None else None, + ) + if self._owner_id_column_name: + params = (*params, owner_id) + batch_params.append((*params, entry["event_id"])) + with self._config.provide_connection() as conn, MssqlPythonCursor(conn) as cursor: try: - for entry in entries: - params: tuple[Any, ...] = ( - entry["id"], - entry["session_id"], - entry["app_name"], - entry["user_id"], - entry.get("scope", "user"), - entry["event_id"], - entry.get("author"), - entry["timestamp"], - to_json(entry["content_json"]), - entry["content_text"], - to_json(entry["metadata_json"]) if entry.get("metadata_json") is not None else None, - ) - if self._owner_id_column_name: - params = (*params, owner_id) - cursor.execute(sql, (*params, entry["event_id"])) - inserted += resolve_rowcount(cursor) + cursor.executemany(sql, batch_params) + rowcount = resolve_rowcount(cursor) + inserted = rowcount if rowcount >= 0 else len(batch_params) except Exception: conn.rollback() raise @@ -524,6 +531,9 @@ def search_entries( if not self._enabled: msg = "ADK memory store is disabled" raise RuntimeError(msg) + if embedding is not None: + msg = "Vector embeddings are not configured for MssqlPythonADKMemoryStore" + raise ImproperConfigurationError(msg) limit_value = limit or self._max_results where_scope, scope_params = _build_mssql_scope_where(app_name, user_id, scope_filter) escaped_query = _escape_like_pattern(query) @@ -991,6 +1001,3 @@ def _session_list_query( ORDER BY {column} {direction}, id {direction}{page_clause} """ return sql, tuple(params) - - -_escape_sql_literal = escape_sql_literal diff --git a/sqlspec/adapters/mssql_python/config.py b/sqlspec/adapters/mssql_python/config.py index 18823161e..f6eb990ed 100644 --- a/sqlspec/adapters/mssql_python/config.py +++ b/sqlspec/adapters/mssql_python/config.py @@ -139,7 +139,7 @@ class MssqlPythonPoolParams(MssqlPythonConnectionParams): class MssqlPythonDriverFeatures(TypedDict): """mssql-python driver feature flags.""" - use_pool: NotRequired[bool] + enable_pooling: NotRequired[bool] """Enable SQLSpec connection pooling around ``mssql-python`` connections (defaults to ``True``).""" json_serializer: "NotRequired[Callable[[Any], str]]" """Custom JSON serializer applied to bound JSON parameters (defaults to ``to_json``).""" @@ -295,7 +295,7 @@ def _create_mssql_python_pool( ) -> "MssqlPythonConnectionPool": pool_size = int(connection_config.get("pool_size", 100)) pool_idle_timeout = int(connection_config.get("pool_idle_timeout", 600)) - pool_enabled = bool(connection_config.get("pool_enabled", driver_features.get("use_pool", True))) + pool_enabled = bool(connection_config.get("pool_enabled", driver_features.get("enable_pooling", True))) connection_string, connect_kwargs = build_connection_config(connection_config) return MssqlPythonConnectionPool( connection_string=connection_string, diff --git a/sqlspec/adapters/mssql_python/core.py b/sqlspec/adapters/mssql_python/core.py index 251e6227a..670bcd0fb 100644 --- a/sqlspec/adapters/mssql_python/core.py +++ b/sqlspec/adapters/mssql_python/core.py @@ -22,6 +22,7 @@ SerializationConflictError, SQLParsingError, SQLSpecError, + TransactionError, UniqueViolationError, ) from sqlspec.utils.config_tools import parse_odbc_connection_string @@ -49,8 +50,10 @@ "driver_profile", "escape_sql_literal", "extract_error_number", + "extract_sqlstate", "format_identifier", "materialize_tuple_rows", + "normalize_execute_parameters", "quote_bracket_identifier", "resolve_column_names", "resolve_many_rowcount", @@ -92,7 +95,16 @@ } _CONNECT_KWARG_KEYS: Final[set[str]] = {"autocommit", "attrs_before", "native_uuid", "token_provider"} _CONNECT_TIMEOUT_KEYS: Final[tuple[str, ...]] = ("timeout", "connection_timeout", "login_timeout", "command_timeout") -_POOL_CONFIG_KEYS: Final[set[str]] = {"pool_size", "pool_idle_timeout", "pool_enabled"} +_POOL_CONFIG_KEYS: Final[set[str]] = { + "pool_size", + "pool_idle_timeout", + "pool_enabled", + "pool_min_size", + "pool_max_size", + "min_size", + "max_size", + "pool_timeout", +} _IGNORED_CONNECTION_CONFIG_KEYS: Final[set[str]] = { "application_name", "app", @@ -100,9 +112,19 @@ "multiple_active_result_sets", "workstation_id", } +_ODBC_SQLSTATE_MAP: Final[dict[str, tuple[type[SQLSpecError], str]]] = { + "08S01": (DatabaseConnectionError, "communication link failure"), + "08001": (DatabaseConnectionError, "unable to connect to data source"), + "HYT00": (QueryTimeoutError, "query timeout expired"), + "HY008": (QueryTimeoutError, "operation canceled"), +} _ERROR_CODE_MAPPING: Final[dict[int, tuple[type[SQLSpecError], str]]] = { + 220: (DataError, "arithmetic overflow error"), + 258: (QueryTimeoutError, "timeout waiting for latch"), 2601: (UniqueViolationError, "unique constraint violation"), 2627: (UniqueViolationError, "unique constraint violation"), + 3617: (QueryTimeoutError, "query canceled by user or attention event"), + 3930: (TransactionError, "current transaction cannot be committed"), 515: (NotNullViolationError, "not-null constraint violation"), 18456: (PermissionDeniedError, "permission denied"), 229: (PermissionDeniedError, "permission denied"), @@ -127,6 +149,10 @@ 8115: (DataError, "arithmetic overflow error"), 8152: (DataError, "string or binary data would be truncated"), 1105: (OperationalError, "operational error"), + 20003: (QueryTimeoutError, "query timeout expired (FreeTDS)"), + 20002: (DatabaseConnectionError, "adaptive Server connection failed (FreeTDS)"), + 20009: (DatabaseConnectionError, "unable to connect: host unreachable or server down (FreeTDS)"), + 20047: (DatabaseConnectionError, "connection closed or broken (FreeTDS)"), } @@ -157,6 +183,15 @@ def split_schema_and_table(identifier: str, *, default_schema: str = "dbo") -> t return parts[-2], parts[-1] +def normalize_execute_parameters(parameters: Any) -> Any: + """Normalize parameters for mssql-python execute calls.""" + if parameters is None: + return None + if isinstance(parameters, list): + return tuple(parameters) + return parameters + + def format_identifier(identifier: str) -> str: """Format a T-SQL identifier with bracket quoting.""" cleaned = identifier.strip() @@ -164,13 +199,13 @@ def format_identifier(identifier: str) -> str: msg = "Table name must not be empty" raise SQLSpecError(msg) parts = split_qualified_identifier(cleaned, quote_chars='"', allow_bracket_quotes=True) - return ".".join(quote_bracket_identifier(part) for part in parts) + return ".".join([quote_bracket_identifier(part) for part in parts]) def build_insert_statement(table: str, columns: "list[str]") -> str: """Build an mssql-python INSERT statement with positional qmark placeholders.""" - column_clause = ", ".join(quote_bracket_identifier(column) for column in columns) - placeholders = ", ".join("?" for _ in columns) + column_clause = ", ".join([quote_bracket_identifier(column) for column in columns]) + placeholders = ", ".join(["?"] * len(columns)) return f"INSERT INTO {format_identifier(table)} ({column_clause}) VALUES ({placeholders})" @@ -276,8 +311,32 @@ def extract_error_number(exc: BaseException | None) -> int | None: return int(last_match[0] or last_match[1]) +_SQLSTATE_LENGTH: Final[int] = 5 + + +def extract_sqlstate(error: BaseException) -> str | None: + """Extract an ODBC SQLSTATE code from a driver error.""" + val = getattr(error, "sqlstate", None) + if isinstance(val, str) and val: + return val.strip().upper() + if hasattr(error, "args") and len(error.args) > 1: + second = error.args[1] + if isinstance(second, str) and len(second.strip()) == _SQLSTATE_LENGTH: + return second.strip().upper() + message = str(error).upper() + for code in _ODBC_SQLSTATE_MAP: + if code in message: + return code + return None + + def create_mapped_exception(error: BaseException, *, logger: "Logger | None" = None) -> SQLSpecError: """Map a SQL Server exception to SQLSpec's exception hierarchy.""" + sqlstate = extract_sqlstate(error) + if sqlstate is not None and sqlstate in _ODBC_SQLSTATE_MAP: + error_class, description = _ODBC_SQLSTATE_MAP[sqlstate] + return error_class(f"SQL Server SQLSTATE {sqlstate}: {description}. Original error: {error}") + error_number = extract_error_number(error) if error_number == _MSSQL_CONSTRAINT_547: @@ -300,6 +359,8 @@ def create_mapped_exception(error: BaseException, *, logger: "Logger | None" = N return constraint_exception exc_name = type(error).__name__ + if exc_name == "InterfaceError": + return DatabaseConnectionError(f"SQL Server interface error. Original error: {error}") if exc_name == "IntegrityError": return IntegrityError(f"SQL Server integrity error. Original error: {error}") if exc_name == "OperationalError": @@ -329,7 +390,7 @@ def apply_driver_features( ) -> "tuple[StatementConfig, dict[str, Any]]": """Merge mssql-python driver-feature defaults with caller overrides.""" features: dict[str, Any] = dict(driver_features) if driver_features else {} - features.setdefault("use_pool", True) + features.setdefault("enable_pooling", True) json_serializer = features.setdefault("json_serializer", to_json) json_deserializer = features.setdefault("json_deserializer", from_json) @@ -430,7 +491,7 @@ def build_connection_config(params: dict[str, Any]) -> tuple[str, dict[str, Any] extra = config.pop("extra", None) if isinstance(extra, dict): - parts.extend(f"{key}={_format_connection_value(value)}" for key, value in extra.items() if value is not None) + parts.extend([f"{key}={_format_connection_value(value)}" for key, value in extra.items() if value is not None]) for remaining_key, remaining_value in list(config.items()): if remaining_value is not None: diff --git a/sqlspec/adapters/mssql_python/data_dictionary.py b/sqlspec/adapters/mssql_python/data_dictionary.py index a9a9e0658..ddfac4775 100644 --- a/sqlspec/adapters/mssql_python/data_dictionary.py +++ b/sqlspec/adapters/mssql_python/data_dictionary.py @@ -44,7 +44,7 @@ from sqlspec.adapters.mssql_python.driver import MssqlPythonDriver from sqlspec.core import SQL - from sqlspec.data_dictionary._types import DialectConfig, MetadataCapabilityProfile + from sqlspec.data_dictionary import DialectConfig, MetadataCapabilityProfile __all__ = ("MssqlPythonSyncDataDictionary", "MssqlVersionInfo") @@ -277,6 +277,51 @@ def get_constraints( ) return MetadataResult("constraints", capability=capability, items=tuple(rows)) + def get_schemas(self, driver: "MssqlPythonDriver") -> MetadataResult: + """Get schema metadata.""" + capability = self.get_metadata_capabilities(driver, ("schemas",)).get("schemas") + self._log_schema_introspect(driver, schema_name=None, table_name=None, operation="schemas") + rows = driver.select(self.get_domain_query("schemas", "list")) + return MetadataResult("schemas", capability=capability, items=tuple(rows)) + + def get_objects(self, driver: "MssqlPythonDriver", schema: str | None = None) -> MetadataResult: + """Get database object metadata.""" + schema_name = self.resolve_connection_schema(driver, schema) + capability = self.get_metadata_capabilities(driver, ("objects",)).get("objects") + self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="objects") + rows = driver.select(self.get_domain_query("objects", "by_schema"), schema_name=schema_name) + return MetadataResult("objects", capability=capability, items=tuple(rows)) + + def get_table_details(self, driver: "MssqlPythonDriver", table: str, schema: str | None = None) -> MetadataResult: + """Get rich table metadata.""" + schema_name = self.resolve_connection_schema(driver, schema) + capability = self.get_metadata_capabilities(driver, ("tables",)).get("tables") + self._log_schema_introspect(driver, schema_name=schema_name, table_name=table, operation="table_details") + rows = driver.select( + self.get_domain_query("tables", "details_by_table"), schema_name=schema_name, table_name=table + ) + return MetadataResult("tables", capability=capability, items=tuple(rows)) + + def get_routines(self, driver: "MssqlPythonDriver", schema: str | None = None) -> MetadataResult: + """Get routine metadata.""" + schema_name = self.resolve_connection_schema(driver, schema) + capability = self.get_metadata_capabilities(driver, ("routines",)).get("routines") + self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="routines") + rows = driver.select(self.get_domain_query("routines", "by_schema"), schema_name=schema_name) + return MetadataResult("routines", capability=capability, items=tuple(rows)) + + def get_privileges( + self, driver: "MssqlPythonDriver", object_name: str | None = None, schema: str | None = None + ) -> MetadataResult: + """Get privilege metadata.""" + schema_name = self.resolve_connection_schema(driver, schema) + capability = self.get_metadata_capabilities(driver, ("privileges",)).get("privileges") + self._log_schema_introspect(driver, schema_name=schema_name, table_name=object_name, operation="privileges") + rows = driver.select(self.get_domain_query("permissions", "by_schema"), schema_name=schema_name) + if object_name is not None: + rows = [r for r in rows if _row_value(r, "object_name") == object_name] + return MetadataResult("privileges", capability=capability, items=tuple(rows)) + def get_views(self, driver: "MssqlPythonDriver", schema: str | None = None) -> MetadataResult: """Get view metadata for a schema.""" schema_name = self.resolve_connection_schema(driver, schema) @@ -297,11 +342,12 @@ def get_dependencies( self, driver: "MssqlPythonDriver", object_name: str | None = None, schema: str | None = None ) -> MetadataResult: """Get SQL expression dependency metadata for a schema.""" - _ = object_name schema_name = self.resolve_connection_schema(driver, schema) capability = self.get_metadata_capabilities(driver, ("dependencies",)).get("dependencies") - self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="dependencies") + self._log_schema_introspect(driver, schema_name=schema_name, table_name=object_name, operation="dependencies") rows = driver.select(self.get_domain_query("dependencies", "by_schema"), schema_name=schema_name) + if object_name is not None: + rows = [r for r in rows if _row_value(r, "referencing_object_name") == object_name] return MetadataResult("dependencies", capability=capability, items=tuple(rows)) def get_system_metadata( diff --git a/sqlspec/adapters/mssql_python/driver.py b/sqlspec/adapters/mssql_python/driver.py index 7809336fb..6674d99ae 100644 --- a/sqlspec/adapters/mssql_python/driver.py +++ b/sqlspec/adapters/mssql_python/driver.py @@ -1,6 +1,7 @@ """mssql-python sync driver.""" import contextlib +import sys from collections.abc import Sized from typing import TYPE_CHECKING, Any, Literal, NotRequired, TypedDict, cast @@ -17,6 +18,7 @@ driver_profile, format_identifier, materialize_tuple_rows, + normalize_execute_parameters, quote_bracket_identifier, resolve_column_names, resolve_many_rowcount, @@ -99,8 +101,11 @@ def start(self) -> None: try: cursor = cursor_manager.__enter__() handler = self._driver.handle_database_exceptions() - with handler: - _execute_cursor(cursor, self._sql, self._parameters) + if sys.exc_info()[1] is None: + with handler: + _execute_cursor(cursor, self._sql, self._parameters) + else: + self._driver._run_with_exception_handler(handler, _execute_cursor, cursor, self._sql, self._parameters) self._driver._check_pending_exception(handler) except BaseException: with contextlib.suppress(Exception): @@ -115,8 +120,11 @@ def fetch_chunk(self) -> "list[dict[str, Any]]": cursor = cursor_manager.cursor handler = self._driver.handle_database_exceptions() rows: Any = [] - with handler: - rows = cursor.fetchmany(self._chunk_size) + if sys.exc_info()[1] is None: + with handler: + rows = cursor.fetchmany(self._chunk_size) + else: + rows = self._driver._run_with_exception_handler(handler, cursor.fetchmany, self._chunk_size) self._driver._check_pending_exception(handler) if not rows: return [] @@ -140,6 +148,7 @@ class MssqlPythonDriver(SyncDriverAdapterBase): __slots__ = ( "_column_name_cache", "_data_dictionary", + "_implicit_transaction", "_migration_schema_restore", "_restore_autocommit", "_transaction_active", @@ -162,6 +171,7 @@ def __init__( self._migration_schema_restore: tuple[str, str] | None = None self._restore_autocommit = False self._transaction_active = False + self._implicit_transaction = False @property def data_dictionary(self) -> "MssqlPythonSyncDataDictionary": @@ -170,10 +180,12 @@ def data_dictionary(self) -> "MssqlPythonSyncDataDictionary": return self._data_dictionary def dispatch_execute(self, cursor: "MssqlPythonRawCursor", statement: "SQL") -> "ExecutionResult": + if not self._transaction_active and not bool(getattr(self.connection, "autocommit", True)): + self._implicit_transaction = True sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) _execute_cursor(cursor, sql, prepared_parameters) - if statement.returns_rows(): + if cursor.description is not None: fetched, column_names, row_format = collect_rows( cursor.fetchall(), cursor.description, self._column_name_cache ) @@ -189,6 +201,8 @@ def dispatch_execute(self, cursor: "MssqlPythonRawCursor", statement: "SQL") -> return self.create_execution_result(cursor, rowcount_override=resolve_rowcount(cursor)) def dispatch_execute_many(self, cursor: "MssqlPythonRawCursor", statement: "SQL") -> "ExecutionResult": + if not self._transaction_active and not bool(getattr(self.connection, "autocommit", True)): + self._implicit_transaction = True sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) parameter_count = len(prepared_parameters) if isinstance(prepared_parameters, Sized) else None cursor.executemany(sql, cast("Any", prepared_parameters)) @@ -196,15 +210,7 @@ def dispatch_execute_many(self, cursor: "MssqlPythonRawCursor", statement: "SQL" return self.create_execution_result(cursor, rowcount_override=affected_rows, is_many_result=True) def dispatch_execute_script(self, cursor: "MssqlPythonRawCursor", statement: "SQL") -> "ExecutionResult": - sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) - statements = self.split_script_statements(sql, statement.statement_config, strip_trailing_semicolon=True) - successful_count = 0 - for stmt in statements: - _execute_cursor(cursor, stmt, prepared_parameters) - successful_count += 1 - return self.create_execution_result( - cursor, statement_count=len(statements), successful_statements=successful_count, is_script_result=True - ) + return super().dispatch_execute_script(cursor, statement) def collect_rows(self, cursor: "MssqlPythonRawCursor", fetched: "list[Any]") -> "tuple[list[Any], list[str], int]": column_names = resolve_column_names(cursor.description, self._column_name_cache) @@ -222,8 +228,9 @@ def begin(self) -> None: if restore_autocommit: self.connection.autocommit = False except MssqlPythonError as exc: - msg = f"Failed to begin transaction: {exc}" - raise SQLSpecError(msg) from exc + mapped = create_mapped_exception(exc) + msg = f"Failed to begin transaction: {mapped}" + raise type(mapped)(msg) from exc self._restore_autocommit = restore_autocommit self._transaction_active = True @@ -231,19 +238,23 @@ def commit(self) -> None: try: self.connection.commit() except MssqlPythonError as exc: - msg = f"Failed to commit transaction: {exc}" - raise SQLSpecError(msg) from exc + mapped = create_mapped_exception(exc) + msg = f"Failed to commit transaction: {mapped}" + raise type(mapped)(msg) from exc self._transaction_active = False + self._implicit_transaction = False self._restore_connection_autocommit() def rollback(self) -> None: try: self.connection.rollback() except MssqlPythonError as exc: - msg = f"Failed to rollback transaction: {exc}" - raise SQLSpecError(msg) from exc + mapped = create_mapped_exception(exc) + msg = f"Failed to rollback transaction: {mapped}" + raise type(mapped)(msg) from exc finally: self._transaction_active = False + self._implicit_transaction = False self._restore_connection_autocommit() def with_cursor(self, connection: "MssqlPythonConnection") -> "MssqlPythonCursor": @@ -270,21 +281,29 @@ def rollback_to_savepoint(self, name: str) -> None: def set_migration_session_schema(self, schema: str) -> None: """Point the database user's default schema at the migration schema, remembering the prior one.""" + + def _set_schema() -> None: + with self.with_cursor(self.connection) as cursor: + if self._migration_schema_restore is None: + _execute_cursor(cursor, "SELECT USER_NAME() AS user_name, SCHEMA_NAME() AS schema_name;", None) + row: Any = cursor.fetchone() + if row is None: + return + user_name, current_schema = ( + (row["user_name"], row["schema_name"]) if isinstance(row, dict) else (row[0], row[1]) + ) + resolved_schema = str(current_schema or "dbo") + _execute_cursor(cursor, _alter_default_schema_sql(str(user_name), schema), None) + self._migration_schema_restore = (str(user_name), resolved_schema) + else: + _execute_cursor(cursor, _alter_default_schema_sql(self._migration_schema_restore[0], schema), None) + exc_handler = self.handle_database_exceptions() - with exc_handler, self.with_cursor(self.connection) as cursor: - if self._migration_schema_restore is None: - _execute_cursor(cursor, "SELECT USER_NAME() AS user_name, SCHEMA_NAME() AS schema_name;", None) - row: Any = cursor.fetchone() - if row is None: - return - user_name, current_schema = ( - (row["user_name"], row["schema_name"]) if isinstance(row, dict) else (row[0], row[1]) - ) - resolved_schema = str(current_schema or "dbo") - _execute_cursor(cursor, _alter_default_schema_sql(str(user_name), schema), None) - self._migration_schema_restore = (str(user_name), resolved_schema) - else: - _execute_cursor(cursor, _alter_default_schema_sql(self._migration_schema_restore[0], schema), None) + if sys.exc_info()[1] is None: + with exc_handler: + _set_schema() + else: + self._run_with_exception_handler(exc_handler, _set_schema) self._check_pending_exception(exc_handler) def reset_migration_session_schema(self) -> None: @@ -292,21 +311,36 @@ def reset_migration_session_schema(self) -> None: if self._migration_schema_restore is None: return user_name, previous_schema = self._migration_schema_restore - exc_handler = self.handle_database_exceptions() - with exc_handler: + + def _reset_schema() -> None: with self.with_cursor(self.connection) as cursor: _execute_cursor(cursor, _alter_default_schema_sql(user_name, previous_schema), None) self.connection.commit() + + exc_handler = self.handle_database_exceptions() + if sys.exc_info()[1] is None: + with exc_handler: + _reset_schema() + else: + self._run_with_exception_handler(exc_handler, _reset_schema) self._check_pending_exception(exc_handler) self._migration_schema_restore = None def has_schema(self, schema: str) -> bool: """Return whether the specified schema exists.""" + + def _check_schema() -> bool: + with self.with_cursor(self.connection) as cursor: + _execute_cursor(cursor, "SELECT 1 FROM sys.schemas WHERE name = ?", (schema,)) + return cursor.fetchone() is not None + result = False exc_handler = self.handle_database_exceptions() - with exc_handler, self.with_cursor(self.connection) as cursor: - _execute_cursor(cursor, "SELECT 1 FROM sys.schemas WHERE name = ?", (schema,)) - result = cursor.fetchone() is not None + if sys.exc_info()[1] is None: + with exc_handler: + result = _check_schema() + else: + result = bool(self._run_with_exception_handler(exc_handler, _check_schema)) self._check_pending_exception(exc_handler) return result @@ -336,7 +370,9 @@ def select_to_arrow( cursor = None reader: object | None = None exc_handler = self.handle_database_exceptions() - with exc_handler: + + def _open_reader() -> None: + nonlocal cursor, reader, table cursor = cursor_manager.__enter__() try: _execute_cursor(cursor, sql, prepared_parameters) @@ -349,6 +385,12 @@ def select_to_arrow( cursor_manager.__exit__(None, None, None) cursor = None raise + + if sys.exc_info()[1] is None: + with exc_handler: + _open_reader() + else: + self._run_with_exception_handler(exc_handler, _open_reader) self._check_pending_exception(exc_handler) if table is not None: return build_arrow_result_from_table( @@ -369,28 +411,40 @@ def select_to_arrow( arrow_schema=arrow_schema, ) + def _fetch_arrow() -> "ArrowResult": + with self.with_cursor(self.connection) as cursor: + _execute_cursor(cursor, sql, prepared_parameters) + if return_format == "batches": + reader = _cursor_arrow_reader(cursor, arrow_kwargs) + if reader is not None: + return build_arrow_result_from_reader( + prepared_statement, + arrow_reader_with_deferred_close(reader, cursor.close), + return_format=return_format, + batch_size=batch_size, + arrow_schema=arrow_schema, + ) + table = cursor.arrow(**arrow_kwargs) + if table is None: + msg = "mssql-python did not return an Arrow table." + raise SQLSpecError(msg) + return build_arrow_result_from_table( + prepared_statement, + table, + return_format=return_format, + batch_size=batch_size, + arrow_schema=arrow_schema, + ) + + result: ArrowResult | None = None exc_handler = self.handle_database_exceptions() - with exc_handler, self.with_cursor(self.connection) as cursor: - _execute_cursor(cursor, sql, prepared_parameters) - if return_format == "batches": - reader = _cursor_arrow_reader(cursor, arrow_kwargs) - if reader is not None: - return build_arrow_result_from_reader( - prepared_statement, - reader, - return_format=return_format, - batch_size=batch_size, - arrow_schema=arrow_schema, - ) - table = cursor.arrow(**arrow_kwargs) + if sys.exc_info()[1] is None: + with exc_handler: + result = _fetch_arrow() + else: + result = self._run_with_exception_handler(exc_handler, _fetch_arrow) self._check_pending_exception(exc_handler) - - if table is None: - msg = "mssql-python did not return an Arrow table." - raise SQLSpecError(msg) - return build_arrow_result_from_table( - prepared_statement, table, return_format=return_format, batch_size=batch_size, arrow_schema=arrow_schema - ) + return cast("ArrowResult", result) def bulk_copy( self, @@ -408,25 +462,33 @@ def bulk_copy( use_internal_transaction: bool = False, ) -> MssqlPythonBulkCopyResult: """Bulk insert rows via mssql-python cursor.bulkcopy().""" - result: MssqlPythonBulkCopyResult = {"rows_copied": 0} + + def _do_bulkcopy() -> Any: + with self.with_cursor(self.connection) as cursor: + raw_result = cursor.bulkcopy( + target_table, + rows, + batch_size=batch_size, + timeout=timeout, + column_mappings=column_mappings, + keep_identity=keep_identity, + check_constraints=check_constraints, + table_lock=table_lock, + keep_nulls=keep_nulls, + fire_triggers=fire_triggers, + use_internal_transaction=use_internal_transaction, + ) + return _coerce_bulk_copy_result(raw_result, cursor) + + result: MssqlPythonBulkCopyResult | None = None exc_handler = self.handle_database_exceptions() - with exc_handler, self.with_cursor(self.connection) as cursor: - raw_result = cursor.bulkcopy( - target_table, - rows, - batch_size=batch_size, - timeout=timeout, - column_mappings=column_mappings, - keep_identity=keep_identity, - check_constraints=check_constraints, - table_lock=table_lock, - keep_nulls=keep_nulls, - fire_triggers=fire_triggers, - use_internal_transaction=use_internal_transaction, - ) - result = _coerce_bulk_copy_result(raw_result, cursor) + if sys.exc_info()[1] is None: + with exc_handler: + result = _do_bulkcopy() + else: + result = self._run_with_exception_handler(exc_handler, _do_bulkcopy) self._check_pending_exception(exc_handler) - return result + return cast("MssqlPythonBulkCopyResult", result) def load_from_arrow( self, @@ -542,19 +604,28 @@ def load_from_arrow( } raw_result: Any = None use_fallback = False - exc_handler = self.handle_database_exceptions() - with exc_handler, self.with_cursor(self.connection) as cursor: - native_bulkcopy = getattr(cursor, "bulkcopy_arrow", None) - if is_stream and not callable(native_bulkcopy): - msg = "This mssql-python version does not support native Arrow stream bulk copy" - raise SQLSpecError(msg) - if overwrite: - cursor.execute(f"DELETE FROM {_quote_mssql_table(table)}") - if has_rows: - if callable(native_bulkcopy): - raw_result = native_bulkcopy(table, arrow_source, **options) - else: + + def _do_bulkcopy_arrow() -> Any: + nonlocal use_fallback + with self.with_cursor(self.connection) as cursor: + native_bulkcopy = getattr(cursor, "bulkcopy_arrow", None) + if is_stream and not callable(native_bulkcopy): + msg = "This mssql-python version does not support native Arrow stream bulk copy" + raise SQLSpecError(msg) + if overwrite: + cursor.execute(f"DELETE FROM {format_identifier(table)}") + if has_rows: + if callable(native_bulkcopy): + return native_bulkcopy(table, arrow_source, **options) use_fallback = True + return None + + exc_handler = self.handle_database_exceptions() + if sys.exc_info()[1] is None: + with exc_handler: + raw_result = _do_bulkcopy_arrow() + else: + raw_result = self._run_with_exception_handler(exc_handler, _do_bulkcopy_arrow) self._check_pending_exception(exc_handler) if use_fallback: _, records = self._arrow_table_to_rows(cast("Any", arrow_source)) @@ -594,25 +665,18 @@ def load_from_storage( **kwargs: Any, ) -> "StorageBridgeJob": """Load staged artifacts from storage into SQL Server via BulkCopy.""" - if kwargs: - return super().load_from_storage( - table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs - ) - arrow_table, inbound = self._read_storage_arrow(source, file_format=file_format) - return self.load_from_arrow(table, arrow_table, partitioner=partitioner, overwrite=overwrite, telemetry=inbound) + return super().load_from_storage( + table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs + ) def _execute_cache_hit( self, sql: str, params: "tuple[Any, ...] | list[Any] | dict[str, Any]", cached: "CachedQuery" ) -> "SQLResult": - if not params: - statement = self._cached_statement( - sql, params, cached, None, params_are_simple=True, compiled_sql=cached.compiled_sql - ) - return self._execute_cached_statement(statement) - return super()._execute_cache_hit(sql, params, cached) + return super()._execute_cache_hit(sql, normalize_execute_parameters(params) or (), cached) def _connection_in_transaction(self) -> bool: - return self._transaction_active + """Return True for explicit transactions and for open implicit transactions.""" + return self._transaction_active or self._implicit_transaction def _restore_connection_autocommit(self) -> None: restore_autocommit = self._restore_autocommit @@ -622,18 +686,9 @@ def _restore_connection_autocommit(self) -> None: try: self.connection.autocommit = True except MssqlPythonError as exc: - msg = f"Failed to restore autocommit: {exc}" - raise SQLSpecError(msg) from exc - - -def _quote_mssql_table(table: str) -> str: - """Bracket-quote each part of a qualified table name. - - Bracket quoting stays valid whatever the session's ``QUOTED_IDENTIFIER`` - setting is, unlike double quotes which parse as a string literal when it is - off. - """ - return format_identifier(table) + mapped = create_mapped_exception(exc) + msg = f"Failed to restore autocommit: {mapped}" + raise type(mapped)(msg) from exc def _execute_cursor(cursor: "MssqlPythonRawCursor", sql: str, parameters: Any) -> None: diff --git a/sqlspec/adapters/mssql_python/litestar/store.py b/sqlspec/adapters/mssql_python/litestar/store.py index 33c7a9261..fba56967c 100644 --- a/sqlspec/adapters/mssql_python/litestar/store.py +++ b/sqlspec/adapters/mssql_python/litestar/store.py @@ -4,7 +4,12 @@ from typing import TYPE_CHECKING, Any from sqlspec.adapters.mssql_python._typing import MssqlPythonCursor -from sqlspec.adapters.mssql_python.core import escape_sql_literal, quote_bracket_identifier, split_schema_and_table +from sqlspec.adapters.mssql_python.core import ( + escape_sql_literal, + quote_bracket_identifier, + resolve_rowcount, + split_schema_and_table, +) from sqlspec.extensions.litestar.store import BaseSQLSpecStore from sqlspec.utils.sync_tools import async_ @@ -133,7 +138,7 @@ def _set(self, key: str, value: "str | bytes", expires_in: "int | timedelta | No expires_at = self._calculate_expires_at(expires_in) table_ref = _table_ref(self._table_name) sql = f""" - MERGE INTO {table_ref} AS target + MERGE INTO {table_ref} WITH (HOLDLOCK) AS target USING (SELECT ? AS session_id, CONVERT(VARBINARY(MAX), ?) AS data, ? AS expires_at) AS src ON target.session_id = src.session_id WHEN MATCHED THEN @@ -197,7 +202,7 @@ def _delete_expired(self) -> int: """ with self._config.provide_connection() as conn, MssqlPythonCursor(conn) as cursor: cursor.execute(sql) - count = int(getattr(cursor, "rowcount", 0) or 0) + count = resolve_rowcount(cursor) conn.commit() if count > 0: self._log_delete_expired(count) diff --git a/sqlspec/adapters/mssql_python/pool.py b/sqlspec/adapters/mssql_python/pool.py index 27a84199d..7bb572e27 100644 --- a/sqlspec/adapters/mssql_python/pool.py +++ b/sqlspec/adapters/mssql_python/pool.py @@ -1,5 +1,6 @@ """mssql-python pool facade.""" +import contextlib import threading import warnings from typing import TYPE_CHECKING, Any, cast @@ -67,8 +68,10 @@ def acquire(self) -> "MssqlPythonConnection": connection = cast( "MssqlPythonConnection", MSSQL_PYTHON_MODULE.connect(self.connection_string, **self.connect_kwargs) ) - if self.on_connection_create is not None: + if self.on_connection_create is not None and not getattr(connection, "_sqlspec_initialized", False): self.on_connection_create(connection) + with contextlib.suppress(AttributeError): + setattr(connection, "_sqlspec_initialized", True) return connection def release(self, connection: "MssqlPythonConnection") -> None: diff --git a/sqlspec/adapters/mssql_python/type_converter.py b/sqlspec/adapters/mssql_python/type_converter.py index e5fc68363..858707086 100644 --- a/sqlspec/adapters/mssql_python/type_converter.py +++ b/sqlspec/adapters/mssql_python/type_converter.py @@ -1,4 +1,9 @@ -"""Type converters for mssql-python parameter binding.""" +"""Type converters and Apache Arrow mappings for mssql-python. + +Provides public utility functions and converter classes for mapping SQL Server +data types to PyArrow types (``mssql_type_to_arrow``) and converting Python +runtime values for mssql-python driver operations (``MssqlPythonTypeConverter``). +""" from typing import TYPE_CHECKING, Any, Final, cast from uuid import UUID diff --git a/sqlspec/adapters/mysqlconnector/_typing.py b/sqlspec/adapters/mysqlconnector/_typing.py index 054fea63f..f91e7a715 100644 --- a/sqlspec/adapters/mysqlconnector/_typing.py +++ b/sqlspec/adapters/mysqlconnector/_typing.py @@ -5,17 +5,18 @@ """ import contextlib +import importlib from typing import TYPE_CHECKING, Any -import mysql as _mysql +import mysql import mysql.connector as _mysql_connector from mysql.connector import MySQLConnection as _MysqlConnectorSyncConnection from mysql.connector import aio as _mysql_connector_aio +from mysql.connector import constants as _mysql_connector_constants +from mysql.connector import cursor as _mysql_connector_cursor from mysql.connector import pooling as _mysql_connector_pooling -from mysql.connector.constants import FieldType as _MysqlConnectorFieldType -from mysql.connector.cursor import MySQLCursor as _MysqlConnectorSyncRawCursor - -from sqlspec.typing import import_optional_attr +from mysql.connector.aio import MySQLConnection as _MysqlConnectorAsyncConnection +from mysql.connector.aio import cursor as _mysql_connector_aio_cursor if TYPE_CHECKING: from collections.abc import Awaitable, Callable @@ -76,24 +77,28 @@ class MysqlConnectorMysqlModuleProtocol(Protocol): MysqlConnectorSyncConnection: TypeAlias = _MysqlConnectorSyncConnection MysqlConnectorAio: TypeAlias = MysqlConnectorAioModuleProtocol MysqlConnectorAsyncConnection: TypeAlias = MysqlConnectorAsyncConnectionProtocol - MysqlConnectorSyncRawCursor: TypeAlias = _MysqlConnectorSyncRawCursor + MysqlConnectorSyncRawCursor: TypeAlias = _mysql_connector_cursor.MySQLCursor MysqlConnectorConnectionPool: TypeAlias = _mysql_connector_pooling.MySQLConnectionPool MysqlConnectorError: TypeAlias = _mysql_connector.Error MysqlConnectorFieldType: TypeAlias = MysqlConnectorFieldTypeProtocol MysqlConnectorMysqlModule: TypeAlias = MysqlConnectorMysqlModuleProtocol MysqlConnectorAsyncRawCursor: TypeAlias = MysqlConnectorAsyncRawCursorProtocol - -if not TYPE_CHECKING: - MysqlConnectorAsyncPool = import_optional_attr("mysql.connector.aio.pooling", "MySQLConnectionPool") +else: + try: + from mysql.connector.aio.pooling import MySQLConnectionPool as MysqlConnectorAsyncPool + except ModuleNotFoundError as exc: + if exc.name != "mysql.connector.aio.pooling": + raise + MysqlConnectorAsyncPool = None MysqlConnectorAio = _mysql_connector_aio MysqlConnectorSyncConnection = _MysqlConnectorSyncConnection - MysqlConnectorAsyncConnection = import_optional_attr("mysql.connector.aio", "MySQLConnection") + MysqlConnectorAsyncConnection = _MysqlConnectorAsyncConnection MysqlConnectorConnectionPool = _mysql_connector_pooling.MySQLConnectionPool MysqlConnectorError = _mysql_connector.Error - MysqlConnectorFieldType = _MysqlConnectorFieldType - MysqlConnectorMysqlModule = _mysql - MysqlConnectorSyncRawCursor = _MysqlConnectorSyncRawCursor - MysqlConnectorAsyncRawCursor = import_optional_attr("mysql.connector.aio.cursor", "MySQLCursor") + MysqlConnectorFieldType = _mysql_connector_constants.FieldType + MysqlConnectorMysqlModule = mysql + MysqlConnectorSyncRawCursor = _mysql_connector_cursor.MySQLCursor + MysqlConnectorAsyncRawCursor = _mysql_connector_aio_cursor.MySQLCursor __all__ = ( "MysqlConnectorAio", @@ -126,10 +131,13 @@ def __init__( ) -> None: self.connection = connection self.cursor: MysqlConnectorSyncRawCursor | None = None - self.cursor_options = cursor_options or {} + self.cursor_options = cursor_options def __enter__(self) -> "MysqlConnectorSyncRawCursor": - self.cursor = self.connection.cursor(**self.cursor_options) + if self.cursor_options is None: + self.cursor = self.connection.cursor() + else: + self.cursor = self.connection.cursor(**self.cursor_options) return self.cursor def __exit__(self, *_: object) -> None: @@ -152,10 +160,13 @@ def __init__( ) -> None: self.connection = connection self.cursor: MysqlConnectorAsyncRawCursor | None = None - self.cursor_options = cursor_options or {} + self.cursor_options = cursor_options async def __aenter__(self) -> "MysqlConnectorAsyncRawCursor": - self.cursor = await self.connection.cursor(**self.cursor_options) + if self.cursor_options is None: + self.cursor = await self.connection.cursor() + else: + self.cursor = await self.connection.cursor(**self.cursor_options) return self.cursor async def __aexit__(self, *_: object) -> None: @@ -265,7 +276,8 @@ def _get_mysqlconnector_sync_driver_type() -> "type[MysqlConnectorSyncDriver]": global _resolved_sync_driver_type driver_cls = _resolved_sync_driver_type if driver_cls is None: - driver_cls = import_optional_attr("sqlspec.adapters.mysqlconnector.driver", "MysqlConnectorSyncDriver") + module = importlib.import_module("sqlspec.adapters.mysqlconnector.driver") + driver_cls = module.MysqlConnectorSyncDriver _resolved_sync_driver_type = driver_cls return driver_cls @@ -275,6 +287,7 @@ def _get_mysqlconnector_async_driver_type() -> "type[MysqlConnectorAsyncDriver]" global _resolved_async_driver_type driver_cls = _resolved_async_driver_type if driver_cls is None: - driver_cls = import_optional_attr("sqlspec.adapters.mysqlconnector.driver", "MysqlConnectorAsyncDriver") + module = importlib.import_module("sqlspec.adapters.mysqlconnector.driver") + driver_cls = module.MysqlConnectorAsyncDriver _resolved_async_driver_type = driver_cls return driver_cls diff --git a/sqlspec/adapters/mysqlconnector/adk/store.py b/sqlspec/adapters/mysqlconnector/adk/store.py index b6ef15019..96806dd03 100644 --- a/sqlspec/adapters/mysqlconnector/adk/store.py +++ b/sqlspec/adapters/mysqlconnector/adk/store.py @@ -759,12 +759,12 @@ def append_event_and_update_state( _mysql_upsert_user_state_sql(self._user_state_table), (app_name, user_id, self._json_serializer(user_state)), ) + conn.commit() except Exception: conn.rollback() raise finally: cursor.close() - conn.commit() return _session_record_from_row(row, deserializer=self._json_deserializer) @@ -1045,14 +1045,19 @@ async def search_entries( """ params = (*scope_params, f"%{query}%", limit_value) - async with self._config.provide_connection() as conn: - cursor = await conn.cursor() - try: - await cursor.execute(sql, params) - rows = await cursor.fetchall() - columns = [col[0] for col in cursor.description or []] - finally: - await cursor.close() + try: + async with self._config.provide_connection() as conn: + cursor = await conn.cursor() + try: + await cursor.execute(sql, params) + rows = await cursor.fetchall() + columns = [col[0] for col in cursor.description or []] + finally: + await cursor.close() + except MysqlConnectorError as exc: + if _is_mysql_table_missing(exc): + return [] + raise records: list[StoredMemory] = [] for row in rows: @@ -1072,14 +1077,19 @@ async def delete_entries_by_session(self, session_id: str) -> int: raise RuntimeError(msg) sql = f"DELETE FROM {self._memory_table} WHERE session_id = %s" - async with self._config.provide_connection() as conn: - cursor = await conn.cursor() - try: - await cursor.execute(sql, (session_id,)) - await conn.commit() - return cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0 - finally: - await cursor.close() + try: + async with self._config.provide_connection() as conn: + cursor = await conn.cursor() + try: + await cursor.execute(sql, (session_id,)) + await conn.commit() + return cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0 + finally: + await cursor.close() + except MysqlConnectorError as exc: + if _is_mysql_table_missing(exc): + return 0 + raise async def delete_entries_older_than( self, days: int, app_name: "str | None" = None, scope: "str | None" = None @@ -1099,14 +1109,19 @@ async def delete_entries_older_than( where_sql = " AND ".join(clauses) sql = f"DELETE FROM {self._memory_table} WHERE {where_sql}" - async with self._config.provide_connection() as conn: - cursor = await conn.cursor() - try: - await cursor.execute(sql, tuple(params)) - await conn.commit() - return cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0 - finally: - await cursor.close() + try: + async with self._config.provide_connection() as conn: + cursor = await conn.cursor() + try: + await cursor.execute(sql, tuple(params)) + await conn.commit() + return cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0 + finally: + await cursor.close() + except MysqlConnectorError as exc: + if _is_mysql_table_missing(exc): + return 0 + raise async def _memory_table_ddl(self) -> str: adk_config = _adk_config(self._config) @@ -1279,14 +1294,19 @@ def search_entries( """ params = (*scope_params, f"%{query}%", limit_value) - with self._config.provide_connection() as conn: - cursor = conn.cursor() - try: - cursor.execute(sql, params) - rows = cursor.fetchall() - columns = [col[0] for col in cursor.description or []] - finally: - cursor.close() + try: + with self._config.provide_connection() as conn: + cursor = conn.cursor() + try: + cursor.execute(sql, params) + rows = cursor.fetchall() + columns = [col[0] for col in cursor.description or []] + finally: + cursor.close() + except MysqlConnectorError as exc: + if _is_mysql_table_missing(exc): + return [] + raise records: list[StoredMemory] = [] for row in rows: @@ -1307,14 +1327,19 @@ def delete_entries_by_session(self, session_id: str) -> int: raise RuntimeError(msg) sql = f"DELETE FROM {self._memory_table} WHERE session_id = %s" - with self._config.provide_connection() as conn: - cursor = conn.cursor() - try: - cursor.execute(sql, (session_id,)) - conn.commit() - return cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0 - finally: - cursor.close() + try: + with self._config.provide_connection() as conn: + cursor = conn.cursor() + try: + cursor.execute(sql, (session_id,)) + conn.commit() + return cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0 + finally: + cursor.close() + except MysqlConnectorError as exc: + if _is_mysql_table_missing(exc): + return 0 + raise def delete_entries_older_than(self, days: int, app_name: "str | None" = None, scope: "str | None" = None) -> int: """Delete memory entries older than specified days.""" @@ -1333,14 +1358,19 @@ def delete_entries_older_than(self, days: int, app_name: "str | None" = None, sc where_sql = " AND ".join(clauses) sql = f"DELETE FROM {self._memory_table} WHERE {where_sql}" - with self._config.provide_connection() as conn: - cursor = conn.cursor() - try: - cursor.execute(sql, tuple(params)) - conn.commit() - return cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0 - finally: - cursor.close() + try: + with self._config.provide_connection() as conn: + cursor = conn.cursor() + try: + cursor.execute(sql, tuple(params)) + conn.commit() + return cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0 + finally: + cursor.close() + except MysqlConnectorError as exc: + if _is_mysql_table_missing(exc): + return 0 + raise def _memory_table_ddl(self) -> str: adk_config = _adk_config(self._config) diff --git a/sqlspec/adapters/mysqlconnector/config.py b/sqlspec/adapters/mysqlconnector/config.py index 707c5823a..b3b73f26c 100644 --- a/sqlspec/adapters/mysqlconnector/config.py +++ b/sqlspec/adapters/mysqlconnector/config.py @@ -592,7 +592,7 @@ def __init__( observability_config: "ObservabilityConfig | None" = None, **kwargs: Any, ) -> None: - self.connection_config = build_connection_config(connection_config) + connection_config = build_connection_config(connection_config) statement_config = statement_config or default_statement_config statement_config, driver_features = apply_driver_features(statement_config, driver_features) @@ -601,15 +601,15 @@ def __init__( features_dict.pop("on_connection_create", None) ) - features_dict.setdefault("enable_local_infile_bulk_load", self.connection_config["allow_local_infile"]) - if features_dict.get("enable_local_infile_bulk_load") and not self.connection_config.get("allow_local_infile"): + features_dict.setdefault("enable_local_infile_bulk_load", connection_config["allow_local_infile"]) + if features_dict.get("enable_local_infile_bulk_load") and not connection_config.get("allow_local_infile"): msg = "enable_local_infile_bulk_load requires local_infile=True or allow_local_infile=True in connection_config." raise ImproperConfigurationError(msg) self._initialized_connections: WeakSet[Any] = WeakSet() super().__init__( - connection_config=self.connection_config, + connection_config=connection_config, connection_instance=connection_instance, migration_config=migration_config, statement_config=statement_config, @@ -635,23 +635,6 @@ async def create_connection(self) -> MysqlConnectorAsyncConnection: raise return connection - def provide_connection(self, *args: Any, **kwargs: Any) -> "MysqlConnectorAsyncConnectionContext": - return MysqlConnectorAsyncConnectionContext(self) - - def provide_session( - self, *_args: Any, statement_config: "StatementConfig | None" = None, **_kwargs: Any - ) -> "MysqlConnectorAsyncSessionContext": - statement_config = statement_config or self.statement_config or default_statement_config - handler = _MysqlConnectorAsyncSessionConnectionHandler(self) - - return MysqlConnectorAsyncSessionContext( - acquire_connection=handler.acquire_connection, - release_connection=handler.release_connection, - statement_config=statement_config, - driver_features=self.driver_features, - prepare_driver=self._prepare_driver, - ) - def get_signature_namespace(self) -> "dict[str, Any]": namespace = super().get_signature_namespace() namespace.update({ diff --git a/sqlspec/adapters/mysqlconnector/core.py b/sqlspec/adapters/mysqlconnector/core.py index 5103ad39f..cb5b89387 100644 --- a/sqlspec/adapters/mysqlconnector/core.py +++ b/sqlspec/adapters/mysqlconnector/core.py @@ -3,7 +3,7 @@ import contextlib from collections.abc import Callable, Mapping, Sequence, Sized from datetime import timedelta -from typing import TYPE_CHECKING, Any, Final, Literal, cast +from typing import Any, Final, Literal from sqlspec.core import DriverParameterProfile, ParameterStyle, StatementConfig, build_statement_config_from_profile from sqlspec.driver import rows_to_dicts @@ -16,6 +16,7 @@ ForeignKeyViolationError, IntegrityError, NotNullViolationError, + OperationalError, PermissionDeniedError, QueryTimeoutError, SQLParsingError, @@ -27,14 +28,12 @@ from sqlspec.utils.serializers import from_json, to_json from sqlspec.utils.text import quote_backtick_identifier, split_qualified_identifier from sqlspec.utils.type_converters import build_uuid_coercions -from sqlspec.utils.type_guards import has_cursor_metadata, has_lastrowid, has_rowcount, has_sqlstate, has_type_code - -if TYPE_CHECKING: - from sqlspec.core.parameters import ParameterValidator +from sqlspec.utils.type_guards import has_lastrowid, has_rowcount, has_sqlstate, has_type_code __all__ = ( "MysqlConnectorAsyncStreamSource", "MysqlConnectorSyncStreamSource", + "_resolve_column_types", "apply_driver_features", "build_insert_statement", "build_load_data_statement", @@ -45,16 +44,12 @@ "collect_stream_rows", "create_mapped_exception", "default_statement_config", - "detect_json_columns", - "detect_json_columns_from_description", "driver_profile", "encode_records_for_local_infile", - "escape_literal_percent", "format_identifier", "normalize_execute_many_parameters", "normalize_execute_parameters", "normalize_lastrowid", - "resolve_column_names", "resolve_many_rowcount", "resolve_row_plan", "resolve_rowcount", @@ -80,12 +75,17 @@ MYSQL_SYNTAX_ERROR_MIN: Final[int] = 1064 MYSQL_SYNTAX_ERROR_MAX_EXCLUSIVE: Final[int] = 1100 -_MYSQL_MIGRATION_ERROR_CODES: Final[frozenset[int]] = frozenset((1061, 1091)) +_MYSQL_DDL_ERROR_DISPATCH: dict[int, tuple[type[SQLSpecError], str]] = { + 1061: (OperationalError, "duplicate key name"), + 1091: (OperationalError, "cannot drop column or key"), +} _MYSQL_SQLSTATE_EXACT_DISPATCH: Final[dict[str, tuple[type[SQLSpecError], str]]] = { - "23505": (UniqueViolationError, "unique constraint violation"), - "23503": (ForeignKeyViolationError, "foreign key constraint violation"), - "23502": (NotNullViolationError, "not-null constraint violation"), - "23514": (CheckViolationError, "check constraint violation"), + "42S02": (OperationalError, "table does not exist"), + "42S22": (OperationalError, "column does not exist"), +} +_MYSQL_OPERATIONAL_ERROR_DISPATCH: Final[dict[int, tuple[type[SQLSpecError], str]]] = { + 1146: (OperationalError, "table does not exist"), + 1054: (OperationalError, "column does not exist"), } _MYSQL_SQLSTATE_PREFIX_DISPATCH: Final[dict[str, tuple[type[SQLSpecError], str]]] = { "23": (IntegrityError, "integrity constraint violation"), @@ -152,6 +152,7 @@ class MysqlConnectorSyncStreamSource: "_chunk_size", "_cursor", "_cursor_options", + "_deserializer", "_driver", "_json_type_codes", "_parameters", @@ -177,36 +178,35 @@ def __init__( self._cursor_options = dict(cursor_options or {}) self._cursor_options["buffered"] = False self._row_plan: tuple[list[str], list[int] | None] | None = None + deserializer = getattr(driver, "_json_deserializer", None) + if deserializer is None and hasattr(driver, "driver_features") and isinstance(driver.driver_features, dict): + deserializer = driver.driver_features.get("json_deserializer") + self._deserializer: Callable[[Any], Any] = deserializer or from_json def start(self) -> None: handler = self._driver.handle_database_exceptions() - with handler: - try: - cursor = self._driver.connection.cursor(**self._cursor_options) - self._cursor = cursor - cursor.execute(self._sql, normalize_execute_parameters(self._parameters)) - self._row_plan = resolve_row_plan(self._cursor.description, self._json_type_codes) - except BaseException: - self.close(error=True) - raise + self._driver._run_with_exception_handler(handler, self._start) self._driver._check_pending_exception(handler) + def _start(self) -> None: + try: + cursor = self._driver.connection.cursor(**self._cursor_options) + self._cursor = cursor + cursor.execute(self._sql, normalize_execute_parameters(self._parameters)) + self._row_plan = resolve_row_plan(self._cursor.description, self._json_type_codes) + except BaseException: + self.close(error=True) + raise + def fetch_chunk(self) -> "list[dict[str, Any]]": handler = self._driver.handle_database_exceptions() - rows: list[Any] = [] - with handler: - rows = self._cursor.fetchmany(self._chunk_size) + rows = self._driver._run_with_exception_handler(handler, self._cursor.fetchmany, self._chunk_size) self._driver._check_pending_exception(handler) if not rows: return [] if self._row_plan is None: self._row_plan = resolve_row_plan(self._cursor.description, self._json_type_codes) - deserializer = cast( - "Callable[[Any], Any]", - getattr(self._driver, "_json_deserializer", None) - or self._driver.driver_features.get("json_deserializer", from_json), - ) - return collect_stream_rows(rows, self._row_plan, deserializer) + return collect_stream_rows(rows, self._row_plan, self._deserializer) def close(self, error: bool = False) -> None: cursor = self._cursor @@ -230,6 +230,7 @@ class MysqlConnectorAsyncStreamSource: "_chunk_size", "_cursor", "_cursor_options", + "_deserializer", "_driver", "_json_type_codes", "_parameters", @@ -255,6 +256,10 @@ def __init__( self._cursor_options = dict(cursor_options or {}) self._cursor_options["buffered"] = False self._row_plan: tuple[list[str], list[int] | None] | None = None + deserializer = getattr(driver, "_json_deserializer", None) + if deserializer is None and hasattr(driver, "driver_features") and isinstance(driver.driver_features, dict): + deserializer = driver.driver_features.get("json_deserializer") + self._deserializer: Callable[[Any], Any] = deserializer or from_json async def start(self) -> None: handler = self._driver.handle_database_exceptions() @@ -269,12 +274,7 @@ async def fetch_chunk(self) -> "list[dict[str, Any]]": return [] if self._row_plan is None: self._row_plan = resolve_row_plan(self._cursor.description, self._json_type_codes) - deserializer = cast( - "Callable[[Any], Any]", - getattr(self._driver, "_json_deserializer", None) - or self._driver.driver_features.get("json_deserializer", from_json), - ) - return collect_stream_rows(rows, self._row_plan, deserializer) + return collect_stream_rows(rows, self._row_plan, self._deserializer) async def close(self, error: bool = False) -> None: cursor = self._cursor @@ -303,7 +303,7 @@ async def _start(self) -> None: def build_profile() -> "DriverParameterProfile": """Create the mysql-connector driver parameter profile.""" - coercions: dict[type, Callable[[Any], Any]] = {bool: bool_to_int, **build_uuid_coercions()} + coercions: dict[type, Callable[[Any], Any]] = {bool: int, **build_uuid_coercions()} return DriverParameterProfile( name="mysql-connector", default_style=ParameterStyle.QMARK, @@ -441,13 +441,6 @@ def normalize_execute_many_parameters(parameters: Any) -> Any: return parameters -def resolve_column_names(description: "Sequence[Any] | None") -> "list[str]": - """Extract column names from cursor description.""" - if not description: - return [] - return [col[0] for col in description] - - def resolve_row_plan( description: "Sequence[Any] | None", json_type_codes: "set[int]" ) -> "tuple[list[str], list[int] | None]": @@ -473,24 +466,6 @@ def resolve_row_plan( return column_names, json_indexes or None -def detect_json_columns_from_description( - description: "Sequence[Any] | None", json_type_codes: "set[int]" -) -> "list[int]": - """Identify JSON column indexes from pre-fetched cursor description metadata.""" - return resolve_row_plan(description, json_type_codes)[1] or [] - - -def detect_json_columns( - cursor: Any, json_type_codes: "set[int]", description: "Sequence[Any] | None" = None -) -> "list[int]": - """Identify JSON column indexes from cursor metadata.""" - if description is None: - if not has_cursor_metadata(cursor): - return [] - description = cursor.description - return detect_json_columns_from_description(description, json_type_codes) - - def collect_stream_rows( rows: "Sequence[Any]", row_plan: "tuple[list[str], list[int] | None]", @@ -572,13 +547,15 @@ def normalize_lastrowid(cursor: Any) -> int | None: return last_id if isinstance(last_id, int) else None -def bool_to_int(value: bool) -> int: - """Coerce boolean values to integer for MySQL parameter compatibility.""" - return int(value) - - -def create_mapped_exception(error: Any, *, logger: Any | None = None) -> "SQLSpecError | bool": +def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpecError: """Map MySQL driver exceptions to SQLSpec errors.""" + _ = logger + type_name = type(error).__name__ + if type_name == "PoolError" or "PoolError" in str(type(error)): + return ConnectionTimeoutError(str(error)) + if type_name == "InterfaceError" or getattr(error, "errno", None) == -1: + return DatabaseConnectionError(str(error)) + error_code = getattr(error, "errno", None) if error_code is None and isinstance(error, Exception) and error.args: value = error.args[0] @@ -587,10 +564,9 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> "SQLSpe sqlstate = error.sqlstate if has_sqlstate(error) else None sqlstate_prefix = sqlstate[:2] if isinstance(sqlstate, str) and sqlstate else None - if error_code in _MYSQL_MIGRATION_ERROR_CODES: - if logger is not None: - logger.warning("MySQL expected migration error (ignoring): %s", error) - return True + dispatch = _MYSQL_DDL_ERROR_DISPATCH.get(error_code) if error_code is not None else None + if dispatch is not None: + return _create_mysql_error(error, sqlstate, error_code, dispatch[0], dispatch[1]) dispatch = _MYSQL_SQLSTATE_EXACT_DISPATCH.get(sqlstate) if sqlstate is not None else None if dispatch is not None: @@ -618,6 +594,10 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> "SQLSpe dispatch = _MYSQL_SQLSTATE_PREFIX_DISPATCH["40"] return _create_mysql_error(error, sqlstate, error_code, dispatch[0], dispatch[1]) + dispatch = _MYSQL_OPERATIONAL_ERROR_DISPATCH.get(error_code) if error_code is not None else None + if dispatch is not None: + return _create_mysql_error(error, sqlstate, error_code, dispatch[0], dispatch[1]) + if sqlstate_prefix == "42": dispatch = _MYSQL_SQLSTATE_PREFIX_DISPATCH["42"] return _create_mysql_error(error, sqlstate, error_code, dispatch[0], dispatch[1]) @@ -641,14 +621,6 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> "SQLSpe return _create_mysql_error(error, sqlstate, error_code, SQLSpecError, "database error") -def escape_literal_percent(sql: str, parameters: Any, validator: "ParameterValidator") -> str: - """Escape literal percent characters before driver-side interpolation.""" - if not parameters or "%" not in sql: - return sql - keep = {info.position for info in validator.extract_parameters(sql)} - return "".join("%%" if char == "%" and index not in keep else char for index, char in enumerate(sql)) - - def build_load_data_statement( table: str, columns: "list[str]", placeholder: str = "%s", *, escape_percent: bool = False ) -> str: @@ -752,4 +724,39 @@ def _create_mysql_error( driver_profile = build_profile() +_MYSQL_TYPE_CODE_TOKENS: Final[dict[int, str]] = { + 0: "decimal", + 1: "int32", + 2: "int32", + 3: "int64", + 4: "float32", + 5: "float64", + 7: "timestamp", + 8: "int64", + 10: "date", + 11: "time", + 12: "timestamp", + 246: "decimal", + 252: "binary", + 253: "string", + 254: "string", +} + + +def _resolve_column_types(description: Any) -> "dict[str, str] | None": + """Map MySQL cursor column FIELD_TYPE codes to neutral Arrow type tokens. + + Returns ``None`` when the cursor has no description or reports no + recognizable type codes. + """ + if not description: + return None + column_types: dict[str, str] = {} + for col in description: + token = _MYSQL_TYPE_CODE_TOKENS.get(col[1]) + if token is not None: + column_types[col[0]] = token + return column_types or None + + default_statement_config = build_statement_config() diff --git a/sqlspec/adapters/mysqlconnector/driver.py b/sqlspec/adapters/mysqlconnector/driver.py index b19cb450f..693d08d97 100644 --- a/sqlspec/adapters/mysqlconnector/driver.py +++ b/sqlspec/adapters/mysqlconnector/driver.py @@ -22,6 +22,7 @@ from sqlspec.adapters.mysqlconnector.core import ( MysqlConnectorAsyncStreamSource, MysqlConnectorSyncStreamSource, + _resolve_column_types, build_insert_statement, build_load_data_statement, can_use_local_infile_records, @@ -77,24 +78,6 @@ json_type_value = MysqlConnectorFieldType.JSON if supports_json_type(MysqlConnectorFieldType) else None MYSQLCONNECTOR_JSON_TYPE_CODES: Final[set[int]] = {json_type_value} if json_type_value is not None else set() -_MYSQL_TYPE_CODE_TOKENS: Final[dict[int, str]] = { - 0: "decimal", - 1: "int32", - 2: "int32", - 3: "int64", - 4: "float32", - 5: "float64", - 7: "timestamp", - 8: "int64", - 10: "date", - 11: "time", - 12: "timestamp", - 246: "decimal", - 252: "binary", - 253: "string", - 254: "string", -} - class MysqlConnectorSyncExceptionHandler(BaseSyncExceptionHandler): """Context manager for handling mysql-connector sync exceptions.""" @@ -105,10 +88,7 @@ def _handle_exception(self, exc_type: "type[BaseException] | None", exc_val: "Ba if exc_type is None: return False if issubclass(exc_type, MysqlConnectorError): - result = create_mapped_exception(exc_val, logger=logger) - if result is True: - return True - self.pending_exception = cast("Exception", result) + self.pending_exception = create_mapped_exception(exc_val) return True return False @@ -116,7 +96,7 @@ def _handle_exception(self, exc_type: "type[BaseException] | None", exc_val: "Ba class MysqlConnectorSyncDriver(SyncDriverAdapterBase): """MySQL/MariaDB database driver using mysql-connector sync library.""" - __slots__ = ("_data_dictionary", "_json_deserializer") + __slots__ = ("_cursor_options", "_data_dictionary", "_json_deserializer") dialect = "mysql" def __init__( @@ -135,6 +115,8 @@ def __init__( self._json_deserializer: Callable[[Any], Any] = cast( "Callable[[Any], Any]", self.driver_features.get("json_deserializer", from_json) ) + cursor_options = cast("dict[str, Any] | None", self.driver_features.get("cursor_options")) + self._cursor_options: dict[str, Any] | None = dict(cursor_options) if cursor_options else None @property def data_dictionary(self) -> "MysqlConnectorSyncDataDictionary": @@ -143,7 +125,7 @@ def data_dictionary(self) -> "MysqlConnectorSyncDataDictionary": return self._data_dictionary def dispatch_execute(self, cursor: Any, statement: "SQL") -> "ExecutionResult": - sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) + sql, prepared_parameters = self._compiled_sql(statement, statement.statement_config) cursor.execute(sql, normalize_execute_parameters(prepared_parameters)) if statement.returns_rows() or getattr(cursor, "with_rows", False): @@ -169,7 +151,7 @@ def dispatch_execute(self, cursor: Any, statement: "SQL") -> "ExecutionResult": return self.create_execution_result(cursor, rowcount_override=affected_rows, last_inserted_id=last_id) def dispatch_execute_many(self, cursor: Any, statement: "SQL") -> "ExecutionResult": - sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) + sql, prepared_parameters = self._compiled_sql(statement, statement.statement_config) prepared_parameters = normalize_execute_many_parameters(prepared_parameters) parameter_count = len(prepared_parameters) if isinstance(prepared_parameters, Sized) else None @@ -179,63 +161,47 @@ def dispatch_execute_many(self, cursor: Any, statement: "SQL") -> "ExecutionResu return self.create_execution_result(cursor, rowcount_override=affected_rows, is_many_result=True) def dispatch_execute_script(self, cursor: Any, statement: "SQL") -> "ExecutionResult": - sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) - statements = self.split_script_statements(sql, statement.statement_config, strip_trailing_semicolon=True) - - successful_count = 0 - last_cursor = cursor + """Execute SQL script with statement splitting and parameter handling. - for stmt in statements: - cursor.execute(stmt, normalize_execute_parameters(prepared_parameters)) - if getattr(cursor, "with_rows", False): - cursor.fetchall() - successful_count += 1 + Args: + cursor: mysql-connector cursor object + statement: SQL script to execute - return self.create_execution_result( - last_cursor, statement_count=len(statements), successful_statements=successful_count, is_script_result=True - ) + Returns: + ExecutionResult: Script execution results with statement count + """ + return super().dispatch_execute_script(cursor, statement) def begin(self) -> None: try: with MysqlConnectorSyncCursor(self.connection) as cursor: cursor.execute("BEGIN") except MysqlConnectorError as e: - mapped = create_mapped_exception(e, logger=logger) - if isinstance(mapped, bool): - return - raise mapped from e + raise create_mapped_exception(e) from e def commit(self) -> None: try: self.connection.commit() except MysqlConnectorError as e: - mapped = create_mapped_exception(e, logger=logger) - if isinstance(mapped, bool): - return - raise mapped from e + raise create_mapped_exception(e) from e def rollback(self) -> None: try: self.connection.rollback() except MysqlConnectorError as e: - mapped = create_mapped_exception(e, logger=logger) - if isinstance(mapped, bool): - return - raise mapped from e + raise create_mapped_exception(e) from e def with_cursor(self, connection: "MysqlConnectorSyncConnection") -> "MysqlConnectorSyncCursor": - cursor_options = cast("dict[str, Any]", self.driver_features.get("cursor_options") or {}) - return MysqlConnectorSyncCursor(connection, dict(cursor_options)) + return MysqlConnectorSyncCursor(connection, self._cursor_options) def dispatch_select_stream(self, statement: "SQL", chunk_size: int) -> "SyncRowStream[dict[str, Any]] | None": """Return a native mysql-connector row stream backed by an unbuffered cursor.""" if not statement.returns_rows(): return None sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) - cursor_options = cast("dict[str, Any]", self.driver_features.get("cursor_options") or {}) return SyncRowStream( MysqlConnectorSyncStreamSource( - self, sql, prepared_parameters, chunk_size, MYSQLCONNECTOR_JSON_TYPE_CODES, cursor_options + self, sql, prepared_parameters, chunk_size, MYSQLCONNECTOR_JSON_TYPE_CODES, self._cursor_options ) ) @@ -302,10 +268,8 @@ def load_from_arrow( if overwrite: statement = f"TRUNCATE TABLE {format_identifier(table)}" exc_handler = self.handle_database_exceptions() - with exc_handler, self.with_cursor(self.connection) as cursor: - cursor.execute(statement) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + self._run_with_exception_handler(exc_handler, self._execute_cursor_statement_unprepared, statement) + self._check_pending_exception(exc_handler) columns, records = self._arrow_table_to_rows(arrow_table) if records: @@ -321,12 +285,12 @@ def load_from_arrow( try: with os.fdopen(fd, "wb") as tmp: tmp.write(payload) - load_sql = build_load_data_statement(table, columns) + load_sql = build_load_data_statement(table, columns, escape_percent=True) exc_handler = self.handle_database_exceptions() - with exc_handler, self.with_cursor(self.connection) as cursor: - cursor.execute(load_sql, (tmp_name,)) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + self._run_with_exception_handler( + exc_handler, self._execute_cursor_statement_unprepared, load_sql, (tmp_name,) + ) + self._check_pending_exception(exc_handler) finally: Path(tmp_name).unlink(missing_ok=True) else: @@ -337,30 +301,16 @@ def load_from_arrow( else records ) exc_handler = self.handle_database_exceptions() - with exc_handler, self.with_cursor(self.connection) as cursor: - cursor.executemany(insert_sql, cast("Any", prepared_records)) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + self._run_with_exception_handler( + exc_handler, self._executemany_cursor_statement, insert_sql, prepared_records + ) + self._check_pending_exception(exc_handler) telemetry_payload = self._ingest_telemetry(arrow_table) telemetry_payload["destination"] = table self._attach_partition_telemetry(telemetry_payload, partitioner) return self._storage_job(telemetry_payload, telemetry) - def load_from_storage( - self, - table: str, - source: "StorageDestination", - *, - file_format: "StorageFormat", - partitioner: "dict[str, object] | None" = None, - overwrite: bool = False, - **kwargs: Any, - ) -> "StorageBridgeJob": - return super().load_from_storage( - table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs - ) - def collect_rows(self, cursor: Any, fetched: "list[Any]") -> "tuple[list[Any], list[str], int]": """Collect mysql-connector sync rows for the direct execution path.""" description = cursor.description or None @@ -373,6 +323,27 @@ def resolve_rowcount(self, cursor: Any) -> int: """Resolve rowcount from mysql-connector cursor for the direct execution path.""" return resolve_rowcount(cursor) + def _execute_cursor_statement(self, statement: str, parameters: Any = None) -> None: + with self.with_cursor(self.connection) as cursor: + if parameters is None: + cursor.execute(statement) + else: + cursor.execute(statement, parameters) + + def _execute_cursor_statement_unprepared(self, statement: str, parameters: Any = None) -> None: + options = self._cursor_options + if options and options.get("prepared"): + options = {k: v for k, v in options.items() if k != "prepared"} + with MysqlConnectorSyncCursor(self.connection, options or None) as cursor: + if parameters is None: + cursor.execute(statement) + else: + cursor.execute(statement, parameters) + + def _executemany_cursor_statement(self, statement: str, parameters: Any) -> None: + with self.with_cursor(self.connection) as cursor: + cursor.executemany(statement, parameters) + def _connection_in_transaction(self) -> bool: """Report whether the connector has an open transaction.""" return bool(self.connection.in_transaction) @@ -387,10 +358,7 @@ def _handle_exception(self, exc_type: "type[BaseException] | None", exc_val: "Ba if exc_type is None: return False if issubclass(exc_type, MysqlConnectorError): - result = create_mapped_exception(exc_val, logger=logger) - if result is True: - return True - self.pending_exception = cast("Exception", result) + self.pending_exception = create_mapped_exception(exc_val) return True return False @@ -398,7 +366,7 @@ def _handle_exception(self, exc_type: "type[BaseException] | None", exc_val: "Ba class MysqlConnectorAsyncDriver(AsyncDriverAdapterBase): """MySQL/MariaDB database driver using mysql-connector async library.""" - __slots__ = ("_data_dictionary", "_json_deserializer") + __slots__ = ("_cursor_options", "_data_dictionary", "_json_deserializer") dialect = "mysql" def __init__( @@ -417,6 +385,8 @@ def __init__( self._json_deserializer: Callable[[Any], Any] = cast( "Callable[[Any], Any]", self.driver_features.get("json_deserializer", from_json) ) + cursor_options = cast("dict[str, Any] | None", self.driver_features.get("cursor_options")) + self._cursor_options: dict[str, Any] | None = dict(cursor_options) if cursor_options else None @property def data_dictionary(self) -> "MysqlConnectorAsyncDataDictionary": @@ -425,7 +395,7 @@ def data_dictionary(self) -> "MysqlConnectorAsyncDataDictionary": return self._data_dictionary async def dispatch_execute(self, cursor: Any, statement: "SQL") -> "ExecutionResult": - sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) + sql, prepared_parameters = self._compiled_sql(statement, statement.statement_config) await cursor.execute(sql, normalize_execute_parameters(prepared_parameters)) if statement.returns_rows() or getattr(cursor, "with_rows", False): @@ -451,7 +421,7 @@ async def dispatch_execute(self, cursor: Any, statement: "SQL") -> "ExecutionRes return self.create_execution_result(cursor, rowcount_override=affected_rows, last_inserted_id=last_id) async def dispatch_execute_many(self, cursor: Any, statement: "SQL") -> "ExecutionResult": - sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) + sql, prepared_parameters = self._compiled_sql(statement, statement.statement_config) prepared_parameters = normalize_execute_many_parameters(prepared_parameters) parameter_count = len(prepared_parameters) if isinstance(prepared_parameters, Sized) else None @@ -461,63 +431,47 @@ async def dispatch_execute_many(self, cursor: Any, statement: "SQL") -> "Executi return self.create_execution_result(cursor, rowcount_override=affected_rows, is_many_result=True) async def dispatch_execute_script(self, cursor: Any, statement: "SQL") -> "ExecutionResult": - sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) - statements = self.split_script_statements(sql, statement.statement_config, strip_trailing_semicolon=True) - - successful_count = 0 - last_cursor = cursor + """Execute SQL script with statement splitting and parameter handling. - for stmt in statements: - await cursor.execute(stmt, normalize_execute_parameters(prepared_parameters)) - if getattr(cursor, "with_rows", False): - await cursor.fetchall() - successful_count += 1 + Args: + cursor: mysql-connector async cursor object + statement: SQL script to execute - return self.create_execution_result( - last_cursor, statement_count=len(statements), successful_statements=successful_count, is_script_result=True - ) + Returns: + ExecutionResult: Script execution results with statement count + """ + return await super().dispatch_execute_script(cursor, statement) async def begin(self) -> None: try: async with MysqlConnectorAsyncCursor(self.connection) as cursor: await cursor.execute("BEGIN") except MysqlConnectorError as e: - mapped = create_mapped_exception(e, logger=logger) - if isinstance(mapped, bool): - return - raise mapped from e + raise create_mapped_exception(e) from e async def commit(self) -> None: try: await self.connection.commit() except MysqlConnectorError as e: - mapped = create_mapped_exception(e, logger=logger) - if isinstance(mapped, bool): - return - raise mapped from e + raise create_mapped_exception(e) from e async def rollback(self) -> None: try: await self.connection.rollback() except MysqlConnectorError as e: - mapped = create_mapped_exception(e, logger=logger) - if isinstance(mapped, bool): - return - raise mapped from e + raise create_mapped_exception(e) from e def with_cursor(self, connection: "MysqlConnectorAsyncConnection") -> "MysqlConnectorAsyncCursor": - cursor_options = cast("dict[str, Any]", self.driver_features.get("cursor_options") or {}) - return MysqlConnectorAsyncCursor(connection, dict(cursor_options)) + return MysqlConnectorAsyncCursor(connection, self._cursor_options) def dispatch_select_stream(self, statement: "SQL", chunk_size: int) -> "AsyncRowStream[dict[str, Any]] | None": """Return a native mysql-connector row stream backed by an unbuffered cursor.""" if not statement.returns_rows(): return None sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) - cursor_options = cast("dict[str, Any]", self.driver_features.get("cursor_options") or {}) return AsyncRowStream( MysqlConnectorAsyncStreamSource( - self, sql, prepared_parameters, chunk_size, MYSQLCONNECTOR_JSON_TYPE_CODES, cursor_options + self, sql, prepared_parameters, chunk_size, MYSQLCONNECTOR_JSON_TYPE_CODES, self._cursor_options ) ) @@ -584,10 +538,8 @@ async def load_from_arrow( if overwrite: statement = f"TRUNCATE TABLE {format_identifier(table)}" exc_handler = self.handle_database_exceptions() - async with exc_handler, self.with_cursor(self.connection) as cursor: - await cursor.execute(statement) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + await self._run_with_exception_handler(exc_handler, self._execute_cursor_statement_unprepared, statement) + self._check_pending_exception(exc_handler) columns, records = self._arrow_table_to_rows(arrow_table) if records: @@ -603,12 +555,12 @@ async def load_from_arrow( try: with os.fdopen(fd, "wb") as tmp: tmp.write(payload) - load_sql = build_load_data_statement(table, columns) + load_sql = build_load_data_statement(table, columns, escape_percent=True) exc_handler = self.handle_database_exceptions() - async with exc_handler, self.with_cursor(self.connection) as cursor: - await cursor.execute(load_sql, (tmp_name,)) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + await self._run_with_exception_handler( + exc_handler, self._execute_cursor_statement_unprepared, load_sql, (tmp_name,) + ) + self._check_pending_exception(exc_handler) finally: await asyncio.to_thread(Path(tmp_name).unlink, True) else: @@ -619,30 +571,16 @@ async def load_from_arrow( else records ) exc_handler = self.handle_database_exceptions() - async with exc_handler, self.with_cursor(self.connection) as cursor: - await cursor.executemany(insert_sql, cast("Any", prepared_records)) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + await self._run_with_exception_handler( + exc_handler, self._executemany_cursor_statement, insert_sql, prepared_records + ) + self._check_pending_exception(exc_handler) telemetry_payload = self._ingest_telemetry(arrow_table) telemetry_payload["destination"] = table self._attach_partition_telemetry(telemetry_payload, partitioner) return self._storage_job(telemetry_payload, telemetry) - async def load_from_storage( - self, - table: str, - source: "StorageDestination", - *, - file_format: "StorageFormat", - partitioner: "dict[str, object] | None" = None, - overwrite: bool = False, - **kwargs: Any, - ) -> "StorageBridgeJob": - return await super().load_from_storage( - table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs - ) - def collect_rows(self, cursor: Any, fetched: "list[Any]") -> "tuple[list[Any], list[str], int]": """Collect mysql-connector async rows for the direct execution path.""" description = cursor.description or None @@ -655,25 +593,30 @@ def resolve_rowcount(self, cursor: Any) -> int: """Resolve rowcount from mysql-connector cursor for the direct execution path.""" return resolve_rowcount(cursor) + async def _execute_cursor_statement(self, statement: str, parameters: Any = None) -> None: + async with self.with_cursor(self.connection) as cursor: + if parameters is None: + await cursor.execute(statement) + else: + await cursor.execute(statement, parameters) + + async def _execute_cursor_statement_unprepared(self, statement: str, parameters: Any = None) -> None: + options = self._cursor_options + if options and options.get("prepared"): + options = {k: v for k, v in options.items() if k != "prepared"} + async with MysqlConnectorAsyncCursor(self.connection, options or None) as cursor: + if parameters is None: + await cursor.execute(statement) + else: + await cursor.execute(statement, parameters) + + async def _executemany_cursor_statement(self, statement: str, parameters: Any) -> None: + async with self.with_cursor(self.connection) as cursor: + await cursor.executemany(statement, parameters) + def _connection_in_transaction(self) -> bool: """Report whether the connector has an open transaction.""" return bool(self.connection.in_transaction) -def _resolve_column_types(description: Any) -> "dict[str, str] | None": - """Map MySQL cursor column FIELD_TYPE codes to neutral Arrow type tokens. - - Returns ``None`` when the cursor has no description or reports no - recognizable type codes. - """ - if not description: - return None - column_types: dict[str, str] = {} - for col in description: - token = _MYSQL_TYPE_CODE_TOKENS.get(col[1]) - if token is not None: - column_types[col[0]] = token - return column_types or None - - register_driver_profile("mysql-connector", driver_profile) diff --git a/sqlspec/adapters/mysqlconnector/litestar/store.py b/sqlspec/adapters/mysqlconnector/litestar/store.py index 77961f2e1..69ae7d4a5 100644 --- a/sqlspec/adapters/mysqlconnector/litestar/store.py +++ b/sqlspec/adapters/mysqlconnector/litestar/store.py @@ -103,10 +103,10 @@ async def set(self, key: str, value: "str | bytes", expires_in: "int | timedelta sql = f""" INSERT INTO {self._table_name} (session_id, data, expires_at) - VALUES (%s, %s, %s) AS new + VALUES (%s, %s, %s) ON DUPLICATE KEY UPDATE - data = new.data, - expires_at = new.expires_at, + data = VALUES(data), + expires_at = VALUES(expires_at), updated_at = UTC_TIMESTAMP(6) """ @@ -237,11 +237,7 @@ def _table_ddl(self) -> str: """ def _drop_table_sql(self) -> "list[str]": - safe_index_table = self._table_name.replace(".", "_") - return [ - f"DROP INDEX idx_{safe_index_table}_expires_at ON {self._table_name}", - f"DROP TABLE IF EXISTS {self._table_name}", - ] + return [f"DROP TABLE IF EXISTS {self._table_name}"] class MysqlConnectorSyncStore(BaseSQLSpecStore["MysqlConnectorSyncConfig"]): @@ -298,17 +294,12 @@ def _table_ddl(self) -> str: """ def _drop_table_sql(self) -> "list[str]": - safe_index_table = self._table_name.replace(".", "_") - return [ - f"DROP INDEX idx_{safe_index_table}_expires_at ON {self._table_name}", - f"DROP TABLE IF EXISTS {self._table_name}", - ] + return [f"DROP TABLE IF EXISTS {self._table_name}"] def _create_table(self) -> None: sql = self._table_ddl() with self._config.provide_session() as driver: driver.execute_script(sql) - driver.commit() self._log_table_created() def _get(self, key: str, renew_for: "int | timedelta | None" = None) -> "bytes | None": @@ -358,10 +349,10 @@ def _set(self, key: str, value: "str | bytes", expires_in: "int | timedelta | No sql = f""" INSERT INTO {self._table_name} (session_id, data, expires_at) - VALUES (%s, %s, %s) AS new + VALUES (%s, %s, %s) ON DUPLICATE KEY UPDATE - data = new.data, - expires_at = new.expires_at, + data = VALUES(data), + expires_at = VALUES(expires_at), updated_at = UTC_TIMESTAMP(6) """ diff --git a/sqlspec/adapters/oracledb/__init__.py b/sqlspec/adapters/oracledb/__init__.py index 4e6361480..9b995ce8b 100644 --- a/sqlspec/adapters/oracledb/__init__.py +++ b/sqlspec/adapters/oracledb/__init__.py @@ -24,8 +24,6 @@ build_session_release_plsql, build_session_reset_plsql, build_tns_description, - coerce_oracle_df_to_arrow_batches, - coerce_oracle_df_to_arrow_table, default_statement_config, execute_session_checkout_async, execute_session_checkout_sync, @@ -35,8 +33,6 @@ extract_duality_view_etag, extract_oracle_wallet_zip, inject_duality_view_etag, - oracle_df_batches_to_arrow_reader, - oracle_df_to_arrow_table, parse_sqlnet_ora, parse_tnsnames_ora, sanitize_oracle_text_query, @@ -53,6 +49,12 @@ OracleSyncExceptionHandler, ) from sqlspec.adapters.oracledb.type_converter import OracleOutputConverter +from sqlspec.utils.arrow_helpers import ( + coerce_oracle_df_to_arrow_batches, + coerce_oracle_df_to_arrow_table, + oracle_df_batches_to_arrow_reader, + oracle_df_to_arrow_table, +) __all__ = ( "VALID_ORACLE_DRIVER_FEATURE_KEYS", diff --git a/sqlspec/adapters/oracledb/_json_handlers.py b/sqlspec/adapters/oracledb/_json_handlers.py index 3db1f0be5..0e9795232 100644 --- a/sqlspec/adapters/oracledb/_json_handlers.py +++ b/sqlspec/adapters/oracledb/_json_handlers.py @@ -38,7 +38,7 @@ from functools import partial from typing import TYPE_CHECKING, Any -from sqlspec.adapters.oracledb._param_types import OracleJson +from sqlspec.adapters.oracledb._param_types import OracleJson, is_json_payload from sqlspec.adapters.oracledb._typing import ( DB_TYPE_BLOB, DB_TYPE_CHAR, @@ -50,11 +50,17 @@ DB_TYPE_NVARCHAR, DB_TYPE_VARCHAR, ) -from sqlspec.adapters.oracledb.data_dictionary import ( +from sqlspec.adapters.oracledb.core import ( + chain_input_handler, + chain_output_handler, resolve_oracle_connection_compatible_major, resolve_oracle_connection_major, ) -from sqlspec.data_dictionary.dialects.oracle import ORACLE_JSON_STORAGE_NATIVE, resolve_oracle_json_storage +from sqlspec.data_dictionary.dialects.oracle import ( + ORACLE_JSON_STORAGE_BLOB_PLAIN, + ORACLE_JSON_STORAGE_NATIVE, + resolve_oracle_json_storage, +) from sqlspec.utils.serializers import from_json, to_json if TYPE_CHECKING: @@ -69,7 +75,6 @@ "is_json_payload", "json_converter_in_blob", "json_converter_in_clob", - "json_converter_in_native", "json_converter_out_blob", "json_converter_out_clob", "json_converter_out_oson", @@ -78,6 +83,7 @@ "register_json_handlers", ) + _JSON_STRING_TYPE_CODES = (DB_TYPE_VARCHAR, DB_TYPE_CHAR, DB_TYPE_NVARCHAR, DB_TYPE_NCHAR) @@ -152,52 +158,6 @@ def register_json_handlers(connection: "Connection | AsyncConnection") -> None: connection.outputtypehandler = chain_output_handler(_output_type_handler, existing_output) -def chain_input_handler(inner: Any, fallback: "Any | None") -> Any: - """Build an input type handler that chains ``inner`` to ``fallback``.""" - if fallback is not None and _has_handler(fallback, inner, _chained_input_handler): - return fallback - return partial(_chained_input_handler, inner, fallback) - - -def chain_output_handler(inner: Any, fallback: "Any | None") -> Any: - """Build an output type handler that chains ``inner`` to ``fallback``.""" - if fallback is not None and _has_handler(fallback, inner, _chained_output_handler): - return fallback - return partial(_chained_output_handler, inner, fallback) - - -def is_json_payload(value: Any) -> bool: - """Return True if the value should be claimed by the JSON input handler. - - ``OracleJson``, ``dict``, empty ``list``/``tuple``, and ``list``/``tuple`` - containing any non-numeric element are claimed. Non-empty sequences where - all elements are ``int``/``float`` (excluding ``bool``) are reserved for - the vector handler unless wrapped in ``OracleJson``. - """ - if isinstance(value, OracleJson): - return True - if isinstance(value, dict): - return True - if isinstance(value, (list, tuple)): - if not value: - return True - return not all(isinstance(v, (int, float)) and not isinstance(v, bool) for v in value) - return False - - -def _has_handler(handler: Any, target_inner: Any, chain_function: Any) -> bool: - current = handler - while current is not None: - if current is target_inner: - return True - if not isinstance(current, partial) or current.func is not chain_function: - return False - inner, current = current.args - if inner is target_inner: - return True - return False - - def _input_type_handler(cursor: "Cursor | AsyncCursor", value: Any, arraysize: int) -> Any: """Oracle input type handler for JSON-shaped Python values.""" if not is_json_payload(value): @@ -206,9 +166,7 @@ def _input_type_handler(cursor: "Cursor | AsyncCursor", value: Any, arraysize: i server_major = resolve_oracle_connection_major(cursor.connection) if server_major is None: - if isinstance(value, OracleJson): - return cursor.var(DB_TYPE_JSON, arraysize=arraysize, inconverter=json_converter_in_native) - return cursor.var(DB_TYPE_JSON, arraysize=arraysize) + return cursor.var(DB_TYPE_BLOB, arraysize=arraysize, inconverter=json_converter_in_blob) compatible_major = resolve_oracle_connection_compatible_major(cursor.connection) storage = resolve_oracle_json_storage(server_major, compatible_major=compatible_major) @@ -216,6 +174,8 @@ def _input_type_handler(cursor: "Cursor | AsyncCursor", value: Any, arraysize: i if isinstance(value, OracleJson): return cursor.var(DB_TYPE_JSON, arraysize=arraysize, inconverter=json_converter_in_native) return cursor.var(DB_TYPE_JSON, arraysize=arraysize) + if storage == ORACLE_JSON_STORAGE_BLOB_PLAIN: + return cursor.var(DB_TYPE_CLOB, arraysize=arraysize, inconverter=json_converter_in_clob) return cursor.var(DB_TYPE_BLOB, arraysize=arraysize, inconverter=json_converter_in_blob) @@ -246,25 +206,3 @@ def _output_type_handler(cursor: "Cursor | AsyncCursor", metadata: Any) -> Any: return cursor.var(type_code, arraysize=cursor.arraysize, outconverter=json_converter_out_clob) return None - - -def _chained_input_handler( - inner: Any, fallback: "Any | None", cursor: "Cursor | AsyncCursor", value: Any, arraysize: int -) -> Any: - """Run ``inner`` input handler, falling back to ``fallback`` when it abstains.""" - result = inner(cursor, value, arraysize) - if result is not None: - return result - if fallback is not None: - return fallback(cursor, value, arraysize) - return None - - -def _chained_output_handler(inner: Any, fallback: "Any | None", cursor: "Cursor | AsyncCursor", metadata: Any) -> Any: - """Run ``inner`` output handler, falling back to ``fallback`` when it abstains.""" - result = inner(cursor, metadata) - if result is not None: - return result - if fallback is not None: - return fallback(cursor, metadata) - return None diff --git a/sqlspec/adapters/oracledb/_param_types.py b/sqlspec/adapters/oracledb/_param_types.py index 34017e726..ffc534f5b 100644 --- a/sqlspec/adapters/oracledb/_param_types.py +++ b/sqlspec/adapters/oracledb/_param_types.py @@ -45,3 +45,21 @@ class OracleJson: def __init__(self, value: Any) -> None: self.value = value + + +def is_json_payload(value: Any) -> bool: + """Return True if the value should be claimed by the JSON input handler. + + OracleJson, dict, empty list/tuple, and list/tuple containing any non-numeric + element are claimed. Non-empty sequences where all elements are int/float + (excluding bool) are reserved for the vector handler unless wrapped in OracleJson. + """ + if isinstance(value, OracleJson): + return True + if isinstance(value, dict): + return True + if isinstance(value, (list, tuple)): + if not value: + return True + return not all(isinstance(v, (int, float)) and not isinstance(v, bool) for v in value) + return False diff --git a/sqlspec/adapters/oracledb/_storage.py b/sqlspec/adapters/oracledb/_storage.py index 7700ba7cf..2563a8244 100644 --- a/sqlspec/adapters/oracledb/_storage.py +++ b/sqlspec/adapters/oracledb/_storage.py @@ -8,17 +8,14 @@ error instead. """ +import re from collections.abc import Mapping -from typing import Any - -from typing_extensions import TypedDict - -from sqlspec.utils.logging import get_logger - -logger = get_logger("sqlspec.adapters.oracledb.storage") +from typing import Any, TypedDict __all__ = ("oracle_table_feature_report", "validate_oracle_identifier") +_UNQUOTED_ORACLE_IDENTIFIER_PATTERN: "re.Pattern[str]" = re.compile(r"^[A-Za-z][A-Za-z0-9_$#]*$") + ORACLE_DEFAULT_HASH_PARTITIONS = 16 ORACLE_MIN_HASH_PARTITIONS = 2 ORACLE_RANGE_INTERVALS = { @@ -43,6 +40,7 @@ class _OracleStorageFeatureReport(TypedDict): applied: tuple[str, ...] clause: str + plsql_clause: str def oracle_table_feature_report( @@ -92,13 +90,16 @@ def oracle_table_feature_report( applied.append("partitioning") clause = "" + plsql_clause = "" if clauses: - clause = " " + " ".join(clauses).replace("'", "''") - return {"applied": tuple(applied), "clause": clause} + clause = " " + " ".join(clauses) + plsql_clause = " " + " ".join(clauses).replace("'", "''") + return {"applied": tuple(applied), "clause": clause, "plsql_clause": plsql_clause} def validate_oracle_identifier(value: str, field_name: str) -> str: - if not value or not value.replace("_", "").isalnum() or value[0].isdigit(): + """Validate that a string is a valid unquoted Oracle SQL identifier.""" + if not value or not _UNQUOTED_ORACLE_IDENTIFIER_PATTERN.fullmatch(value): msg = f"Oracle {field_name} must be a simple SQL identifier" raise ValueError(msg) return value diff --git a/sqlspec/adapters/oracledb/_uuid_handlers.py b/sqlspec/adapters/oracledb/_uuid_handlers.py index 49da1484f..57896377d 100644 --- a/sqlspec/adapters/oracledb/_uuid_handlers.py +++ b/sqlspec/adapters/oracledb/_uuid_handlers.py @@ -5,10 +5,10 @@ """ import uuid -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Final -from sqlspec.adapters.oracledb._json_handlers import chain_input_handler, chain_output_handler from sqlspec.adapters.oracledb._typing import DB_TYPE_RAW +from sqlspec.adapters.oracledb.core import chain_input_handler, chain_output_handler from sqlspec.utils.logging import get_logger from sqlspec.utils.uuids import uuid_from_bytes @@ -31,6 +31,7 @@ UUID_BINARY_SIZE = 16 +_MIN_FETCH_INFO_LEN: Final = 4 def uuid_converter_in(value: uuid.UUID) -> bytes: @@ -131,7 +132,11 @@ def _output_type_handler(cursor: "Cursor | AsyncCursor", metadata: Any) -> Any: Returns: Cursor variable with UUID converter if column is RAW(16), None otherwise. """ - _name, type_code, _display_size, internal_size, _precision, _scale, _null_ok = metadata + type_code = getattr(metadata, "type_code", None) + internal_size = getattr(metadata, "internal_size", None) + if type_code is None and isinstance(metadata, (list, tuple)) and len(metadata) >= _MIN_FETCH_INFO_LEN: + type_code = metadata[1] + internal_size = metadata[3] if type_code is DB_TYPE_RAW and internal_size == UUID_BINARY_SIZE: return cursor.var(type_code, arraysize=cursor.arraysize, outconverter=uuid_converter_out) diff --git a/sqlspec/adapters/oracledb/_vector_handlers.py b/sqlspec/adapters/oracledb/_vector_handlers.py index 7f0856ae2..e3fa3793a 100644 --- a/sqlspec/adapters/oracledb/_vector_handlers.py +++ b/sqlspec/adapters/oracledb/_vector_handlers.py @@ -13,9 +13,8 @@ import array from typing import TYPE_CHECKING, Any, cast -from sqlspec.adapters.oracledb._json_handlers import chain_input_handler, chain_output_handler from sqlspec.adapters.oracledb._typing import DB_TYPE_VECTOR, OracleSparseVector -from sqlspec.adapters.oracledb.data_dictionary import resolve_oracle_connection_major +from sqlspec.adapters.oracledb.core import chain_input_handler, chain_output_handler, resolve_oracle_connection_major from sqlspec.typing import NUMPY_INSTALLED from sqlspec.utils.logging import get_logger from sqlspec.utils.module_loader import import_optional @@ -27,7 +26,7 @@ from sqlspec.adapters.oracledb._typing import OracleSyncRawCursor as Cursor SPARSE_VECTOR_TYPE = OracleSparseVector -np: Any = import_optional("numpy") + __all__ = ( "DTYPE_TO_ARRAY_CODE", @@ -50,8 +49,6 @@ _TYPECODE_INT16 = "h" _TYPECODE_INT32 = "i" -_INT8_MIN = -128 -_INT8_MAX = 127 _ORACLE_VECTOR_MIN_MAJOR = 23 _VECTOR_RETURN_NUMPY = "numpy" @@ -75,6 +72,11 @@ DTYPE_TO_ARRAY_CODE["float16"] = _TYPECODE_FLOAT16 +def _numpy_module() -> Any: + """Return the ``numpy`` module when installed, otherwise ``None``.""" + return import_optional("numpy") + + def numpy_converter_in(value: Any) -> "array.array[Any]": """Convert NumPy array to Oracle array for VECTOR insertion. @@ -115,6 +117,7 @@ def numpy_converter_out(value: "array.array[Any]") -> Any: if not NUMPY_INSTALLED: return value + np = _numpy_module() return np.array(value, copy=True, dtype=value.typecode) @@ -164,7 +167,7 @@ def _is_vector_payload(value: Any) -> bool: return False if isinstance(value, array.array): return True - if NUMPY_INSTALLED and isinstance(value, np.ndarray): + if NUMPY_INSTALLED and isinstance(value, _numpy_module().ndarray): return True if isinstance(value, (list, tuple)) and value: return all(isinstance(v, (int, float)) and not isinstance(v, bool) for v in value) @@ -174,12 +177,8 @@ def _is_vector_payload(value: Any) -> bool: def _pack_python_sequence(value: "list[Any] | tuple[Any, ...]") -> "array.array[Any]": """Pack a Python sequence into an ``array.array`` for VECTOR binding. - Integer sequences entirely within ``[-128, 127]`` use the int8 typecode for - a cheaper bind; everything else falls back to float32 (the 23ai default and - the most common LLM embedding dtype). + Always packs as float32, the standard 23ai default for embedding vectors. """ - if all(isinstance(v, int) and not isinstance(v, bool) and _INT8_MIN <= v <= _INT8_MAX for v in value): - return array.array(_TYPECODE_INT8, value) return array.array(_TYPECODE_FLOAT32, [float(v) for v in value]) @@ -209,7 +208,7 @@ def _input_type_handler(cursor: "Cursor | AsyncCursor", value: Any, arraysize: i if server_major is not None and server_major < _ORACLE_VECTOR_MIN_MAJOR: return None - if NUMPY_INSTALLED and isinstance(value, np.ndarray): + if NUMPY_INSTALLED and isinstance(value, _numpy_module().ndarray): return cursor.var(DB_TYPE_VECTOR, arraysize=arraysize, inconverter=numpy_converter_in) if isinstance(value, array.array): diff --git a/sqlspec/adapters/oracledb/adk/store.py b/sqlspec/adapters/oracledb/adk/store.py index bbf7a7f64..ce41281dd 100644 --- a/sqlspec/adapters/oracledb/adk/store.py +++ b/sqlspec/adapters/oracledb/adk/store.py @@ -1,6 +1,7 @@ """Oracle ADK store for Google Agent Development Kit session/event storage.""" -from decimal import Decimal +import array +import contextlib from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, NotRequired, cast from typing_extensions import TypedDict @@ -11,11 +12,15 @@ from sqlspec.adapters.oracledb.core import ( coerce_large_parameters_async, coerce_large_parameters_sync, + read_oracle_lob_async, + read_oracle_lob_sync, sanitize_oracle_text_query, ) from sqlspec.adapters.oracledb.data_dictionary import JSONStorageType, OracleVersionInfo, storage_type_from_version from sqlspec.builder._vector_distance import build_oracle_hnsw_subquery, build_oracle_hybrid_rrf_sql from sqlspec.config import ADKConfig +from sqlspec.data_dictionary.dialects.oracle import parse_oracle_version_components +from sqlspec.exceptions import ImproperConfigurationError from sqlspec.extensions.adk import ( BaseAsyncADKStore, BaseSyncADKStore, @@ -26,6 +31,7 @@ from sqlspec.extensions.adk.memory.store import BaseAsyncADKMemoryStore, BaseSyncADKMemoryStore from sqlspec.utils.logging import get_logger from sqlspec.utils.serializers import from_json, to_json +from sqlspec.utils.type_converters import build_nested_decimal_normalizer from sqlspec.utils.type_guards import is_async_readable, is_readable if TYPE_CHECKING: @@ -57,6 +63,7 @@ ORACLE_DEFAULT_APP_STATE_TABLE: Final = "adk_app_state" ORACLE_DEFAULT_USER_STATE_TABLE: Final = "adk_user_state" ORACLE_DEFAULT_METADATA_TABLE: Final = "adk_internal_metadata" +_EMBEDDING_COLUMN_INDEX: Final = 12 _ADK_METADATA_TABLE_DDL_TEMPLATE = ( @@ -245,6 +252,7 @@ " event_id VARCHAR2(128) NOT NULL UNIQUE,\n" " author VARCHAR2(256){1},\n" " timestamp TIMESTAMP WITH TIME ZONE NOT NULL,\n" + " {11}" " {2},\n" " content_text CLOB NOT NULL,\n" " inserted_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL\n" @@ -346,6 +354,9 @@ class OracleADKPartitionConfig(TypedDict): class OracleADKConfig(ADKConfig): """Oracle ADK extension settings consumed by Oracle ADK stores.""" + vector_dimensions: NotRequired[int] + """Dimensionality of embedding vectors (e.g. 768 for gemini-embedding-001 with MRL). Default: 768.""" + in_memory: NotRequired[bool] """Enable Oracle INMEMORY clauses on ADK tables.""" @@ -371,6 +382,36 @@ class OracleADKConfig(ADKConfig): """Raw Oracle table options for the ADK memory table.""" +def _rollback_sync(conn: Any) -> None: + with contextlib.suppress(Exception): + conn.rollback() + + +async def _rollback_async(conn: Any) -> None: + with contextlib.suppress(Exception): + await conn.rollback() + + +def _memory_supports_vectors(version: "OracleVersionInfo | None") -> bool: + if version is None or (version.major, version.minor, version.patch) < (23, 4, 0): + return False + compatible = parse_oracle_version_components(version.compatible) if version.compatible else None + return compatible is None or compatible >= (23, 4, 0) + + +def _require_memory_vectors(version: "OracleVersionInfo | None") -> None: + if not _memory_supports_vectors(version): + msg = "Oracle memory embeddings require Oracle Database and COMPATIBLE version 23.4 or later." + raise ImproperConfigurationError(msg) + + +def _memory_insert_count(entry_count: int, batch_errors: "list[Any]") -> int: + for error in batch_errors: + if error.code != ORACLE_DUPLICATE_KEY_ERROR: + raise OracleDatabaseError(error) + return entry_count - len(batch_errors) + + def coerce_decimal_values(value: Any) -> Any: return _coerce_decimal_values(value) @@ -492,10 +533,14 @@ async def create_session( params = {"id": session_id, "app_name": app_name, "user_id": user_id, "state": state_data} async with self._config.provide_connection() as conn: - coerced_params = await coerce_large_parameters_async(conn, params) - with conn.cursor() as cursor: - await cursor.execute(sql, coerced_params) - await conn.commit() + try: + coerced_params = await coerce_large_parameters_async(conn, params) + with conn.cursor() as cursor: + await cursor.execute(sql, coerced_params) + await conn.commit() + except BaseException: + await _rollback_async(conn) + raise result = await self.get_session(app_name, user_id, session_id) if result is None: @@ -669,9 +714,13 @@ async def delete_session(self, app_name: str, user_id: str, session_id: str) -> sql = f"DELETE FROM {self._session_table} WHERE app_name = :app_name AND user_id = :user_id AND id = :id" async with self._config.provide_connection() as conn: - with conn.cursor() as cursor: - await cursor.execute(sql, {"app_name": app_name, "user_id": user_id, "id": session_id}) - await conn.commit() + try: + with conn.cursor() as cursor: + await cursor.execute(sql, {"app_name": app_name, "user_id": user_id, "id": session_id}) + await conn.commit() + except BaseException: + await _rollback_async(conn) + raise async def append_event(self, event_record: StoredEvent) -> None: """Append an event to a session. @@ -688,21 +737,25 @@ async def append_event(self, event_record: StoredEvent) -> None: """ async with self._config.provide_connection() as conn: - coerced_params = await coerce_large_parameters_async( - conn, - { - "id": event_record["id"], - "app_name": event_record["app_name"], - "user_id": event_record["user_id"], - "session_id": event_record["session_id"], - "invocation_id": event_record["invocation_id"], - "timestamp": event_record["timestamp"], - "event_data": await self._serialize_event_data(event_record["event_data"]), - }, - ) - with conn.cursor() as cursor: - await cursor.execute(sql, coerced_params) - await conn.commit() + try: + coerced_params = await coerce_large_parameters_async( + conn, + { + "id": event_record["id"], + "app_name": event_record["app_name"], + "user_id": event_record["user_id"], + "session_id": event_record["session_id"], + "invocation_id": event_record["invocation_id"], + "timestamp": event_record["timestamp"], + "event_data": await self._serialize_event_data(event_record["event_data"]), + }, + ) + with conn.cursor() as cursor: + await cursor.execute(sql, coerced_params) + await conn.commit() + except BaseException: + await _rollback_async(conn) + raise async def append_event_and_update_state( self, @@ -1684,9 +1737,14 @@ def delete_session(self, app_name: str, user_id: str, session_id: str) -> None: """ sql = f"DELETE FROM {self._session_table} WHERE app_name = :app_name AND user_id = :user_id AND id = :id" - with self._config.provide_connection() as conn, conn.cursor() as cursor: - cursor.execute(sql, {"app_name": app_name, "user_id": user_id, "id": session_id}) - conn.commit() + with self._config.provide_connection() as conn: + try: + with conn.cursor() as cursor: + cursor.execute(sql, {"app_name": app_name, "user_id": user_id, "id": session_id}) + conn.commit() + except BaseException: + _rollback_sync(conn) + raise def append_event(self, event_record: StoredEvent) -> None: """Append an event to a session.""" @@ -1698,21 +1756,26 @@ def append_event(self, event_record: StoredEvent) -> None: ) """ - with self._config.provide_connection() as conn, conn.cursor() as cursor: - params = coerce_large_parameters_sync( - conn, - { - "id": event_record["id"], - "app_name": event_record["app_name"], - "user_id": event_record["user_id"], - "session_id": event_record["session_id"], - "invocation_id": event_record["invocation_id"], - "timestamp": event_record["timestamp"], - "event_data": self._serialize_event_data(event_record["event_data"]), - }, - ) - cursor.execute(sql, params) - conn.commit() + with self._config.provide_connection() as conn: + try: + params = coerce_large_parameters_sync( + conn, + { + "id": event_record["id"], + "app_name": event_record["app_name"], + "user_id": event_record["user_id"], + "session_id": event_record["session_id"], + "invocation_id": event_record["invocation_id"], + "timestamp": event_record["timestamp"], + "event_data": self._serialize_event_data(event_record["event_data"]), + }, + ) + with conn.cursor() as cursor: + cursor.execute(sql, params) + conn.commit() + except BaseException: + _rollback_sync(conn) + raise def append_event_and_update_state( self, @@ -2431,46 +2494,61 @@ async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: " if not entries: return 0 + version = await self._get_version_info() + supports_vectors = _memory_supports_vectors(version) + if any(entry.get("embedding") is not None for entry in entries): + _require_memory_vectors(version) + embedding_column = ", embedding" if supports_vectors else "" + embedding_param = ", :embedding" if supports_vectors else "" owner_column = f", {self._owner_id_column_name}" if self._owner_id_column_name else "" owner_param = ", :owner_id" if self._owner_id_column_name else "" sql = f""" INSERT INTO {self._memory_table} ( id, session_id, app_name, user_id, scope, event_id, author{owner_column}, - timestamp, content_json, content_text, metadata_json, inserted_at + timestamp{embedding_column}, content_json, content_text, metadata_json, inserted_at ) VALUES ( :id, :session_id, :app_name, :user_id, :scope, :event_id, :author{owner_param}, - :timestamp, :content_json, :content_text, :metadata_json, :inserted_at + :timestamp{embedding_param}, :content_json, :content_text, :metadata_json, :inserted_at ) """ - inserted_count = 0 - async with self._config.provide_connection() as conn: - with conn.cursor() as cursor: - for entry in entries: - content_json = await self._serialize_json_field(entry["content_json"]) - metadata_json = await self._serialize_json_field(entry["metadata_json"]) - params: dict[str, Any] = { - "id": entry["id"], - "session_id": entry["session_id"], - "app_name": entry["app_name"], - "user_id": entry["user_id"], - "scope": entry.get("scope", "user"), - "event_id": entry["event_id"], - "author": entry["author"], - "timestamp": entry["timestamp"], - "content_json": content_json, - "content_text": entry["content_text"], - "metadata_json": metadata_json, - "inserted_at": entry["inserted_at"], - } - if self._owner_id_column_name: - params["owner_id"] = owner_id - params = await coerce_large_parameters_async(conn, params) - if await self._execute_insert_entry(cursor, sql, params): - inserted_count += 1 - await conn.commit() + batch_params: list[dict[str, Any]] = [] + for entry in entries: + content_json = await self._serialize_json_field(entry["content_json"]) + metadata_json = await self._serialize_json_field(entry["metadata_json"]) + embedding_val = entry.get("embedding") + bound_embedding = array.array("f", embedding_val) if embedding_val is not None else None + p: dict[str, Any] = { + "id": entry["id"], + "session_id": entry["session_id"], + "app_name": entry["app_name"], + "user_id": entry["user_id"], + "scope": entry.get("scope", "user"), + "event_id": entry["event_id"], + "author": entry["author"], + "timestamp": entry["timestamp"], + "content_json": content_json, + "content_text": entry["content_text"], + "metadata_json": metadata_json, + "inserted_at": entry["inserted_at"], + } + if supports_vectors: + p["embedding"] = bound_embedding + if self._owner_id_column_name: + p["owner_id"] = owner_id + batch_params.append(p) - return inserted_count + async with self._config.provide_connection() as conn: + try: + coerced_batch = [await coerce_large_parameters_async(conn, p) for p in batch_params] + with conn.cursor() as cursor: + await cursor.executemany(sql, coerced_batch, batcherrors=True) + inserted_count = _memory_insert_count(len(entries), cursor.getbatcherrors()) + await conn.commit() + return inserted_count + except BaseException: + await _rollback_async(conn) + raise async def search_entries( self, @@ -2486,6 +2564,8 @@ async def search_entries( raise RuntimeError(msg) effective_limit = limit if limit is not None else self._max_results + if embedding is not None: + _require_memory_vectors(await self._get_version_info()) try: if embedding is not None and self._use_fts: @@ -2573,10 +2653,12 @@ async def _deserialize_json_field(self, data: Any) -> "dict[str, Any] | None": return _extract_json_value(data) async def _memory_table_ddl(self) -> str: - storage_type = await self._detect_json_storage_type() - return self._memory_table_ddl_for_type(storage_type) + version = await self._get_version_info() + return self._memory_table_ddl_for_type( + storage_type_from_version(version), supports_vectors=_memory_supports_vectors(version) + ) - def _memory_table_ddl_for_type(self, storage_type: "JSONStorageType") -> str: + def _memory_table_ddl_for_type(self, storage_type: "JSONStorageType", *, supports_vectors: bool = False) -> str: if storage_type == JSONStorageType.JSON_NATIVE: json_columns = """ content_json JSON, @@ -2621,6 +2703,7 @@ def _memory_table_ddl_for_type(self, storage_type: "JSONStorageType") -> str: self._memory_table, self._memory_table, fts_index, + f"embedding VECTOR({self._vector_dimensions}, FLOAT32),\n " if supports_vectors else "", ) def _drop_memory_table_sql(self) -> "list[str]": @@ -2708,12 +2791,15 @@ async def _search_entries_fts( if not sanitized_query: return [] where_scope, scope_params = _build_oracle_scope_where(app_name, user_id, scope_filter) + embedding_column = ( + "embedding" if _memory_supports_vectors(await self._get_version_info()) else "NULL AS embedding" + ) sql = f""" SELECT id, session_id, app_name, user_id, scope, event_id, author, - timestamp, content_json, content_text, metadata_json, inserted_at + timestamp, content_json, content_text, metadata_json, inserted_at, embedding FROM ( SELECT id, session_id, app_name, user_id, scope, event_id, author, - timestamp, content_json, content_text, metadata_json, inserted_at, + timestamp, content_json, content_text, metadata_json, inserted_at, {embedding_column}, SCORE(1) AS score FROM {self._memory_table} WHERE {where_scope} @@ -2759,6 +2845,7 @@ async def _search_entries_hybrid_rrf( "content_text", "metadata_json", "inserted_at", + "embedding", ) if not sanitized_query: sql = build_oracle_hnsw_subquery( @@ -2804,12 +2891,15 @@ async def _search_entries_simple( self, query: str, app_name: str, user_id: str, limit: int, scope_filter: Literal["all", "user", "app"] = "all" ) -> "list[StoredMemory]": where_scope, scope_params = _build_oracle_scope_where(app_name, user_id, scope_filter) + embedding_column = ( + "embedding" if _memory_supports_vectors(await self._get_version_info()) else "NULL AS embedding" + ) sql = f""" SELECT id, session_id, app_name, user_id, scope, event_id, author, - timestamp, content_json, content_text, metadata_json, inserted_at + timestamp, content_json, content_text, metadata_json, inserted_at, embedding FROM ( SELECT id, session_id, app_name, user_id, scope, event_id, author, - timestamp, content_json, content_text, metadata_json, inserted_at + timestamp, content_json, content_text, metadata_json, inserted_at, {embedding_column} FROM {self._memory_table} WHERE {where_scope} AND LOWER(content_text) LIKE :pattern @@ -2833,6 +2923,8 @@ async def _rows_to_records(self, rows: "list[Any]") -> "list[StoredMemory]": content_text = row[9] if is_async_readable(content_text) or is_readable(content_text): content_text = await _read_lob_async(content_text) + embedding_val = row[_EMBEDDING_COLUMN_INDEX] if len(row) > _EMBEDDING_COLUMN_INDEX else None + embedding = [float(x) for x in embedding_val] if embedding_val is not None else None records.append({ "id": row[0], "session_id": row[1], @@ -2846,7 +2938,7 @@ async def _rows_to_records(self, rows: "list[Any]") -> "list[StoredMemory]": "content_text": str(content_text), "metadata_json": metadata_json, "inserted_at": row[11], - "embedding": None, + "embedding": embedding, }) return records @@ -2905,45 +2997,61 @@ def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object if not entries: return 0 + version = self._get_version_info() + supports_vectors = _memory_supports_vectors(version) + if any(entry.get("embedding") is not None for entry in entries): + _require_memory_vectors(version) + embedding_column = ", embedding" if supports_vectors else "" + embedding_param = ", :embedding" if supports_vectors else "" owner_column = f", {self._owner_id_column_name}" if self._owner_id_column_name else "" owner_param = ", :owner_id" if self._owner_id_column_name else "" sql = f""" INSERT INTO {self._memory_table} ( id, session_id, app_name, user_id, scope, event_id, author{owner_column}, - timestamp, content_json, content_text, metadata_json, inserted_at + timestamp{embedding_column}, content_json, content_text, metadata_json, inserted_at ) VALUES ( :id, :session_id, :app_name, :user_id, :scope, :event_id, :author{owner_param}, - :timestamp, :content_json, :content_text, :metadata_json, :inserted_at + :timestamp{embedding_param}, :content_json, :content_text, :metadata_json, :inserted_at ) """ - inserted_count = 0 - with self._config.provide_connection() as conn, conn.cursor() as cursor: - for entry in entries: - content_json = self._serialize_json_field(entry["content_json"]) - metadata_json = self._serialize_json_field(entry["metadata_json"]) - params: dict[str, Any] = { - "id": entry["id"], - "session_id": entry["session_id"], - "app_name": entry["app_name"], - "user_id": entry["user_id"], - "scope": entry.get("scope", "user"), - "event_id": entry["event_id"], - "author": entry["author"], - "timestamp": entry["timestamp"], - "content_json": content_json, - "content_text": entry["content_text"], - "metadata_json": metadata_json, - "inserted_at": entry["inserted_at"], - } - if self._owner_id_column_name: - params["owner_id"] = owner_id - params = coerce_large_parameters_sync(conn, params) - if self._execute_insert_entry(cursor, sql, params): - inserted_count += 1 - conn.commit() + batch_params: list[dict[str, Any]] = [] + for entry in entries: + content_json = self._serialize_json_field(entry["content_json"]) + metadata_json = self._serialize_json_field(entry["metadata_json"]) + embedding_val = entry.get("embedding") + bound_embedding = array.array("f", embedding_val) if embedding_val is not None else None + p: dict[str, Any] = { + "id": entry["id"], + "session_id": entry["session_id"], + "app_name": entry["app_name"], + "user_id": entry["user_id"], + "scope": entry.get("scope", "user"), + "event_id": entry["event_id"], + "author": entry["author"], + "timestamp": entry["timestamp"], + "content_json": content_json, + "content_text": entry["content_text"], + "metadata_json": metadata_json, + "inserted_at": entry["inserted_at"], + } + if supports_vectors: + p["embedding"] = bound_embedding + if self._owner_id_column_name: + p["owner_id"] = owner_id + batch_params.append(p) - return inserted_count + with self._config.provide_connection() as conn: + try: + coerced_batch = [coerce_large_parameters_sync(conn, p) for p in batch_params] + with conn.cursor() as cursor: + cursor.executemany(sql, coerced_batch, batcherrors=True) + inserted_count = _memory_insert_count(len(entries), cursor.getbatcherrors()) + conn.commit() + return inserted_count + except BaseException: + _rollback_sync(conn) + raise def search_entries( self, @@ -2960,6 +3068,8 @@ def search_entries( raise RuntimeError(msg) effective_limit = limit if limit is not None else self._max_results + if embedding is not None: + _require_memory_vectors(self._get_version_info()) try: if embedding is not None and self._use_fts: @@ -3045,10 +3155,12 @@ def _deserialize_json_field(self, data: Any) -> "dict[str, Any] | None": return _extract_json_value(data) def _memory_table_ddl(self) -> str: - storage_type = self._detect_json_storage_type() - return self._memory_table_ddl_for_type(storage_type) + version = self._get_version_info() + return self._memory_table_ddl_for_type( + storage_type_from_version(version), supports_vectors=_memory_supports_vectors(version) + ) - def _memory_table_ddl_for_type(self, storage_type: "JSONStorageType") -> str: + def _memory_table_ddl_for_type(self, storage_type: "JSONStorageType", *, supports_vectors: bool = False) -> str: if storage_type == JSONStorageType.JSON_NATIVE: json_columns = """ content_json JSON, @@ -3093,6 +3205,7 @@ def _memory_table_ddl_for_type(self, storage_type: "JSONStorageType") -> str: self._memory_table, self._memory_table, fts_index, + f"embedding VECTOR({self._vector_dimensions}, FLOAT32),\n " if supports_vectors else "", ) def _drop_memory_table_sql(self) -> "list[str]": @@ -3180,12 +3293,13 @@ def _search_entries_fts( if not sanitized_query: return [] where_scope, scope_params = _build_oracle_scope_where(app_name, user_id, scope_filter) + embedding_column = "embedding" if _memory_supports_vectors(self._get_version_info()) else "NULL AS embedding" sql = f""" SELECT id, session_id, app_name, user_id, scope, event_id, author, - timestamp, content_json, content_text, metadata_json, inserted_at + timestamp, content_json, content_text, metadata_json, inserted_at, embedding FROM ( SELECT id, session_id, app_name, user_id, scope, event_id, author, - timestamp, content_json, content_text, metadata_json, inserted_at, + timestamp, content_json, content_text, metadata_json, inserted_at, {embedding_column}, SCORE(1) AS score FROM {self._memory_table} WHERE {where_scope} @@ -3230,6 +3344,7 @@ def _search_entries_hybrid_rrf( "content_text", "metadata_json", "inserted_at", + "embedding", ) if not sanitized_query: sql = build_oracle_hnsw_subquery( @@ -3274,12 +3389,13 @@ def _search_entries_simple( self, query: str, app_name: str, user_id: str, limit: int, scope_filter: Literal["all", "user", "app"] = "all" ) -> "list[StoredMemory]": where_scope, scope_params = _build_oracle_scope_where(app_name, user_id, scope_filter) + embedding_column = "embedding" if _memory_supports_vectors(self._get_version_info()) else "NULL AS embedding" sql = f""" SELECT id, session_id, app_name, user_id, scope, event_id, author, - timestamp, content_json, content_text, metadata_json, inserted_at + timestamp, content_json, content_text, metadata_json, inserted_at, embedding FROM ( SELECT id, session_id, app_name, user_id, scope, event_id, author, - timestamp, content_json, content_text, metadata_json, inserted_at + timestamp, content_json, content_text, metadata_json, inserted_at, {embedding_column} FROM {self._memory_table} WHERE {where_scope} AND LOWER(content_text) LIKE :pattern @@ -3302,6 +3418,8 @@ def _rows_to_records(self, rows: "list[Any]") -> "list[StoredMemory]": content_text = row[9] if is_readable(content_text): content_text = _read_lob_sync(content_text) + embedding_val = row[_EMBEDDING_COLUMN_INDEX] if len(row) > _EMBEDDING_COLUMN_INDEX else None + embedding = [float(x) for x in embedding_val] if embedding_val is not None else None records.append({ "id": row[0], "session_id": row[1], @@ -3315,7 +3433,7 @@ def _rows_to_records(self, rows: "list[Any]") -> "list[StoredMemory]": "content_text": str(content_text), "metadata_json": metadata_json, "inserted_at": row[11], - "embedding": None, + "embedding": embedding, }) return records @@ -3351,20 +3469,11 @@ def _raise_session_not_found(session_id: str) -> NoReturn: raise ValueError(msg) +_normalize_decimals = build_nested_decimal_normalizer(mode="float") + + def _coerce_decimal_values(value: Any) -> Any: - if isinstance(value, Decimal): - return float(value) - if isinstance(value, dict): - return {key: _coerce_decimal_values(val) for key, val in value.items()} - if isinstance(value, list): - return [_coerce_decimal_values(item) for item in value] - if isinstance(value, tuple): - return tuple(_coerce_decimal_values(item) for item in value) - if isinstance(value, set): - return {_coerce_decimal_values(item) for item in value} - if isinstance(value, frozenset): - return frozenset(_coerce_decimal_values(item) for item in value) - return value + return _normalize_decimals(value) def _oracle_text_value(value: Any) -> str: @@ -3438,21 +3547,15 @@ def _adk_table_feature_clause( hash_partition_key=hash_partition_key, range_partition_key=range_partition_key, ) - return report["clause"] + return report["plsql_clause"] async def _read_lob_async(data: Any) -> Any: - if is_async_readable(data): - return await data.read() - if is_readable(data): - return data.read() - return data + return await read_oracle_lob_async(data) def _read_lob_sync(data: Any) -> Any: - if is_readable(data): - return data.read() - return data + return read_oracle_lob_sync(data) def _build_oracle_scope_where( diff --git a/sqlspec/adapters/oracledb/config.py b/sqlspec/adapters/oracledb/config.py index 150678834..0cb129a41 100644 --- a/sqlspec/adapters/oracledb/config.py +++ b/sqlspec/adapters/oracledb/config.py @@ -302,6 +302,129 @@ class OracleDriverFeatures(TypedDict): """Default ``DBMS_CLOUD`` credential name used for server-side ``DBMS_CLOUD.COPY_DATA`` object-store loads.""" +def _process_oracle_wallet_and_config( + connection_config: "Mapping[str, Any] | None", +) -> "tuple[dict[str, Any], Path | None]": + """Extract wallet zip and normalize connection configuration.""" + normalized_conn = dict(normalize_connection_config(connection_config or {})) + pending_wallet_dir: Path | None = None + if normalized_conn.get("wallet_zip") is not None: + raw_wallet_source = normalized_conn.pop("wallet_zip") + source_is_existing_dir = isinstance(raw_wallet_source, (str, Path)) and Path(raw_wallet_source).is_dir() + wallet_dir = extract_oracle_wallet_zip( + raw_wallet_source, wallet_password=normalized_conn.get("wallet_password") + ) + if not source_is_existing_dir or wallet_dir.resolve() != Path(raw_wallet_source).resolve(): + pending_wallet_dir = wallet_dir + normalized_conn.setdefault("config_dir", str(wallet_dir)) + normalized_conn.setdefault("wallet_location", str(wallet_dir)) + built_config = build_connection_config(normalized_conn) + return built_config, pending_wallet_dir + + +def _build_standalone_connection_params(connection_config: "Mapping[str, Any]") -> dict[str, Any]: + """Build connection parameters for standalone (unpooled) connect calls.""" + config = dict(connection_config) + config.pop("thick_mode", None) + config.pop("lib_dir", None) + config.pop("soda_metadata_cache", None) + config.pop("min", None) + config.pop("max", None) + config.pop("increment", None) + config.pop("homogeneous", None) + config.pop("timeout", None) + config.pop("wait_timeout", None) + config.pop("max_lifetime_session", None) + config.pop("session_callback", None) + config.pop("threaded", None) + config.pop("getmode", None) + config.pop("ping_interval", None) + config.pop("ping_timeout", None) + config.pop("max_sessions_per_shard", None) + config.pop("pool_class", None) + connection_type = config.pop("connectiontype", None) + if connection_type is not None: + config["conn_class"] = connection_type + apply_wallet_ssl_context(config) + return config + + +def _initialize_oracle_client(connection_config: "Mapping[str, Any]") -> None: + """Initialize requested Thick mode before creating a pool or connection.""" + lib_dir = connection_config.get("lib_dir") + if ( + not (connection_config.get("thick_mode") or lib_dir is not None or connection_config.get("soda_metadata_cache")) + or not oracledb.is_thin_mode() + ): + return + client_config = {} + if lib_dir is not None: + client_config["lib_dir"] = lib_dir + if connection_config.get("config_dir") is not None: + client_config["config_dir"] = connection_config["config_dir"] + try: + oracledb.init_oracle_client(**client_config) + except oracledb.Error as exc: + msg = f"Failed to initialize Oracle client in thick mode: {exc}" + raise ImproperConfigurationError(msg) from exc + + +def _validate_async_oracle_mode(connection_config: "Mapping[str, Any]") -> None: + """Reject Thick mode settings unsupported by async connections.""" + if ( + connection_config.get("thick_mode") + or connection_config.get("lib_dir") is not None + or connection_config.get("soda_metadata_cache") + ): + msg = "OracleAsyncConfig only supports Thin mode; Thick mode options (thick_mode, lib_dir, soda_metadata_cache) require OracleSyncConfig." + raise ImproperConfigurationError(msg) + + +def _derive_oracle_pool_kwargs( + source_config: Any, + connection_config_overrides: "Mapping[str, Any] | None", + driver_features_overrides: "Mapping[str, Any] | None", + statement_config: "StatementConfig | None", + migration_config: "dict[str, Any] | MigrationConfig | None", + extension_config: "ExtensionConfigs | None", + observability_config: "ObservabilityConfig | None", + bind_key: "str | None", + overrides: dict[str, Any], +) -> dict[str, Any]: + """Derive common configuration keyword arguments for cloned Oracle pools.""" + new_conn_cfg: dict[str, Any] = dict(source_config.connection_config) + new_driver_features: dict[str, Any] = dict(source_config.driver_features) + if source_config._pool_session_callback is not None and "session_callback" not in new_conn_cfg: + new_conn_cfg["session_callback"] = source_config._pool_session_callback + if source_config._user_connection_hook is not None and "on_connection_create" not in new_driver_features: + new_driver_features["on_connection_create"] = source_config._user_connection_hook + if connection_config_overrides: + new_conn_cfg.update(connection_config_overrides) + if driver_features_overrides: + new_driver_features.update(driver_features_overrides) + for key, val in overrides.items(): + if key in VALID_ORACLE_DRIVER_FEATURE_KEYS: + new_driver_features[key] = val + else: + new_conn_cfg[key] = val + return { + "connection_config": new_conn_cfg, + "connection_instance": None, + "migration_config": cast("dict[str, Any] | None", dict(source_config.migration_config)) + if migration_config is None + else cast("dict[str, Any] | None", migration_config), + "statement_config": source_config.statement_config if statement_config is None else statement_config, + "driver_features": new_driver_features, + "bind_key": source_config.bind_key if bind_key is None else bind_key, + "extension_config": cast("ExtensionConfigs", dict(source_config.extension_config)) + if extension_config is None + else extension_config, + "observability_config": source_config.observability_config + if observability_config is None + else observability_config, + } + + class OracleSyncConnectionContext(SyncPoolConnectionContext): """Context manager for Oracle sync connections.""" @@ -311,7 +434,7 @@ def __init__(self, config: "OracleSyncConfig") -> None: super().__init__(config) self._conn: OracleSyncConnection | None = None - def __enter__(self) -> "OracleSyncConnection": + def acquire_connection(self) -> "OracleSyncConnection": if self._config.connection_instance is None: self._config.connection_instance = self._config.create_pool() conn = self._config.connection_instance.acquire() @@ -326,56 +449,44 @@ def __enter__(self) -> "OracleSyncConnection": raise return cast("OracleSyncConnection", conn) + def release_connection(self, _conn: "OracleSyncConnection | None" = None, **kwargs: Any) -> None: + if self._conn is None: + return + conn = self._conn + self._conn = None + session_hook = self._config.driver_features.get("session_context_hook") + try: + if session_hook is not None: + execute_session_release_sync(conn, session_hook) + finally: + if self._config.connection_instance: + self._config.connection_instance.release(conn) + + def __enter__(self) -> "OracleSyncConnection": + return self.acquire_connection() + def __exit__( self, exc_type: "type[BaseException] | None", exc_val: "BaseException | None", exc_tb: "TracebackType | None" ) -> bool | None: - if self._conn is not None: - conn = self._conn - self._conn = None - session_hook = self._config.driver_features.get("session_context_hook") - try: - if session_hook is not None: - execute_session_release_sync(conn, session_hook) - finally: - if self._config.connection_instance: - self._config.connection_instance.release(conn) + if self._conn is not None and exc_type is not None: + with contextlib.suppress(Exception): + self._conn.rollback() + self.release_connection() return None class _OracleSyncSessionConnectionHandler(SyncPoolSessionFactory): - __slots__ = ("_conn",) + __slots__ = () def __init__(self, config: "OracleSyncConfig") -> None: super().__init__(config) - self._conn: OracleSyncConnection | None = None + self._ctx = OracleSyncConnectionContext(config) def acquire_connection(self) -> "OracleSyncConnection": - if self._config.connection_instance is None: - self._config.connection_instance = self._config.create_pool() - conn = self._config.connection_instance.acquire() - self._conn = conn - session_hook = self._config.driver_features.get("session_context_hook") - if session_hook is not None: - try: - execute_session_checkout_sync(conn, session_hook) - except BaseException: - self._conn = None - self._config.connection_instance.release(conn) - raise - return cast("OracleSyncConnection", conn) + return cast("OracleSyncConnection", self._ctx.acquire_connection()) def release_connection(self, _conn: "OracleSyncConnection", **kwargs: Any) -> None: - if self._conn is None: - return - conn = self._conn - self._conn = None - session_hook = self._config.driver_features.get("session_context_hook") - try: - if session_hook is not None: - execute_session_release_sync(conn, session_hook) - finally: - if self._config.connection_instance: - self._config.connection_instance.release(conn) + self._ctx.release_connection(_conn, **kwargs) class OracleSyncConfig(SyncDatabaseConfig[OracleSyncConnection, "OracleSyncConnectionPool", OracleSyncDriver]): @@ -426,20 +537,8 @@ def __init__( extension_config: Extension-specific configuration. **kwargs: Additional keyword arguments. """ - normalized_conn = dict(normalize_connection_config(connection_config)) + connection_config, pending_wallet_dir = _process_oracle_wallet_and_config(connection_config) self._extracted_wallet_dir: Path | None = None - pending_wallet_dir: Path | None = None - if normalized_conn.get("wallet_zip") is not None: - raw_wallet_source = normalized_conn.pop("wallet_zip") - source_is_existing_dir = isinstance(raw_wallet_source, (str, Path)) and Path(raw_wallet_source).is_dir() - wallet_dir = extract_oracle_wallet_zip( - raw_wallet_source, wallet_password=normalized_conn.get("wallet_password") - ) - if not source_is_existing_dir or wallet_dir.resolve() != Path(raw_wallet_source).resolve(): - pending_wallet_dir = wallet_dir - normalized_conn.setdefault("config_dir", str(wallet_dir)) - normalized_conn.setdefault("wallet_location", str(wallet_dir)) - connection_config = build_connection_config(normalized_conn) self._oracle_version_cache = OracleVersionCache() self._pool_session_callback = cast( "Callable[[OracleSyncConnection, str], None] | None", connection_config.pop("session_callback", None) @@ -471,9 +570,16 @@ def create_connection(self) -> "OracleSyncConnection": Returns: An Oracle Connection instance. """ - if self.connection_instance is None: - self.connection_instance = self.create_pool() - return self.connection_instance.acquire() + _initialize_oracle_client(self.connection_config) + config = _build_standalone_connection_params(self.connection_config) + conn = oracledb.connect(**config) + try: + self._init_connection(conn, "") + except BaseException: + with contextlib.suppress(BaseException): + conn.close() + raise + return conn def provide_pool(self) -> "OracleSyncConnectionPool": """Provide pool instance. @@ -498,35 +604,18 @@ def derive_pool( **overrides: Any, ) -> Self: """Clone this OracleSyncConfig with overridden connection/pool or feature settings and ``connection_instance=None``.""" - new_conn_cfg: dict[str, Any] = dict(self.connection_config) - new_driver_features: dict[str, Any] = dict(self.driver_features) - if self._pool_session_callback is not None and "session_callback" not in new_conn_cfg: - new_conn_cfg["session_callback"] = self._pool_session_callback - if self._user_connection_hook is not None and "on_connection_create" not in new_driver_features: - new_driver_features["on_connection_create"] = self._user_connection_hook - if connection_config_overrides: - new_conn_cfg.update(connection_config_overrides) - if driver_features_overrides: - new_driver_features.update(driver_features_overrides) - for key, val in overrides.items(): - if key in VALID_ORACLE_DRIVER_FEATURE_KEYS: - new_driver_features[key] = val - else: - new_conn_cfg[key] = val - derived = type(self)( - connection_config=new_conn_cfg, - connection_instance=None, - migration_config=cast("dict[str, Any] | None", dict(self.migration_config)) - if migration_config is None - else cast("dict[str, Any] | None", migration_config), - statement_config=self.statement_config if statement_config is None else statement_config, - driver_features=new_driver_features, - bind_key=self.bind_key if bind_key is None else bind_key, - extension_config=cast("ExtensionConfigs", dict(self.extension_config)) - if extension_config is None - else extension_config, - observability_config=self.observability_config if observability_config is None else observability_config, + derived_kwargs = _derive_oracle_pool_kwargs( + self, + connection_config_overrides, + driver_features_overrides, + statement_config, + migration_config, + extension_config, + observability_config, + bind_key, + overrides, ) + derived = type(self)(**derived_kwargs) if ( self._extracted_wallet_dir is not None and derived._extracted_wallet_dir is None @@ -576,16 +665,10 @@ def get_event_runtime_hints(self) -> "EventRuntimeHints": def _create_pool(self) -> "OracleSyncConnectionPool": """Create the actual connection pool.""" + _initialize_oracle_client(self.connection_config) config = dict(self.connection_config) - thick_mode = config.pop("thick_mode", False) - lib_dir = config.pop("lib_dir", None) - if (thick_mode or lib_dir is not None or config.get("soda_metadata_cache")) and oracledb.is_thin_mode(): - client_config = {} - if lib_dir is not None: - client_config["lib_dir"] = lib_dir - if config.get("config_dir") is not None: - client_config["config_dir"] = config["config_dir"] - oracledb.init_oracle_client(**client_config) + config.pop("thick_mode", None) + config.pop("lib_dir", None) apply_wallet_ssl_context(config) @@ -669,6 +752,9 @@ async def __aexit__( self, exc_type: "type[BaseException] | None", exc_val: "BaseException | None", exc_tb: Any ) -> "bool | None": conn = self._connection + if conn is not None and exc_type is not None: + with contextlib.suppress(Exception): + await conn.rollback() session_hook = self._config.driver_features.get("session_context_hook") if conn is not None else None try: if conn is not None and session_hook is not None: @@ -752,20 +838,8 @@ def __init__( extension_config: Extension-specific configuration. **kwargs: Additional keyword arguments. """ - normalized_conn = dict(normalize_connection_config(connection_config)) + connection_config, pending_wallet_dir = _process_oracle_wallet_and_config(connection_config) self._extracted_wallet_dir: Path | None = None - pending_wallet_dir: Path | None = None - if normalized_conn.get("wallet_zip") is not None: - raw_wallet_source = normalized_conn.pop("wallet_zip") - source_is_existing_dir = isinstance(raw_wallet_source, (str, Path)) and Path(raw_wallet_source).is_dir() - wallet_dir = extract_oracle_wallet_zip( - raw_wallet_source, wallet_password=normalized_conn.get("wallet_password") - ) - if not source_is_existing_dir or wallet_dir.resolve() != Path(raw_wallet_source).resolve(): - pending_wallet_dir = wallet_dir - normalized_conn.setdefault("config_dir", str(wallet_dir)) - normalized_conn.setdefault("wallet_location", str(wallet_dir)) - connection_config = build_connection_config(normalized_conn) self._oracle_version_cache = OracleVersionCache() self._pool_session_callback = cast( "Callable[[OracleAsyncConnection, str], Any] | None", connection_config.pop("session_callback", None) @@ -798,9 +872,16 @@ async def create_connection(self) -> OracleAsyncConnection: Returns: An Oracle AsyncConnection instance. """ - if self.connection_instance is None: - self.connection_instance = await self.create_pool() - return cast("OracleAsyncConnection", await self.connection_instance.acquire()) + _validate_async_oracle_mode(self.connection_config) + config = _build_standalone_connection_params(self.connection_config) + conn = await oracledb.connect_async(**config) + try: + await self._init_connection(conn, "") + except BaseException: + with contextlib.suppress(BaseException): + await conn.close() + raise + return cast("OracleAsyncConnection", conn) async def provide_pool(self) -> "OracleAsyncConnectionPool": """Provide async pool instance. @@ -825,35 +906,18 @@ def derive_pool( **overrides: Any, ) -> Self: """Clone this OracleAsyncConfig with overridden connection/pool or feature settings and ``connection_instance=None``.""" - new_conn_cfg: dict[str, Any] = dict(self.connection_config) - new_driver_features: dict[str, Any] = dict(self.driver_features) - if self._pool_session_callback is not None and "session_callback" not in new_conn_cfg: - new_conn_cfg["session_callback"] = self._pool_session_callback - if self._user_connection_hook is not None and "on_connection_create" not in new_driver_features: - new_driver_features["on_connection_create"] = self._user_connection_hook - if connection_config_overrides: - new_conn_cfg.update(connection_config_overrides) - if driver_features_overrides: - new_driver_features.update(driver_features_overrides) - for key, val in overrides.items(): - if key in VALID_ORACLE_DRIVER_FEATURE_KEYS: - new_driver_features[key] = val - else: - new_conn_cfg[key] = val - derived = type(self)( - connection_config=new_conn_cfg, - connection_instance=None, - migration_config=cast("dict[str, Any] | None", dict(self.migration_config)) - if migration_config is None - else cast("dict[str, Any] | None", migration_config), - statement_config=self.statement_config if statement_config is None else statement_config, - driver_features=new_driver_features, - bind_key=self.bind_key if bind_key is None else bind_key, - extension_config=cast("ExtensionConfigs", dict(self.extension_config)) - if extension_config is None - else extension_config, - observability_config=self.observability_config if observability_config is None else observability_config, + derived_kwargs = _derive_oracle_pool_kwargs( + self, + connection_config_overrides, + driver_features_overrides, + statement_config, + migration_config, + extension_config, + observability_config, + bind_key, + overrides, ) + derived = type(self)(**derived_kwargs) if ( self._extracted_wallet_dir is not None and derived._extracted_wallet_dir is None @@ -900,12 +964,10 @@ def get_event_runtime_hints(self) -> "EventRuntimeHints": async def _create_pool(self) -> "OracleAsyncConnectionPool": """Create the actual async connection pool.""" + _validate_async_oracle_mode(self.connection_config) config = dict(self.connection_config) - thick_mode = config.pop("thick_mode", False) - lib_dir = config.pop("lib_dir", None) - if thick_mode or lib_dir is not None: - msg = "OracleAsyncConfig only supports Thin mode; use OracleSyncConfig for Thick mode." - raise ImproperConfigurationError(msg) + config.pop("thick_mode", None) + config.pop("lib_dir", None) apply_wallet_ssl_context(config) diff --git a/sqlspec/adapters/oracledb/core.py b/sqlspec/adapters/oracledb/core.py index fa803ab4c..6f8c245be 100644 --- a/sqlspec/adapters/oracledb/core.py +++ b/sqlspec/adapters/oracledb/core.py @@ -13,17 +13,12 @@ import threading import zipfile from collections.abc import Callable, Iterable, Mapping, Sequence, Sized +from functools import partial from pathlib import Path from typing import TYPE_CHECKING, Any, Final, Protocol, cast -from sqlspec.adapters.oracledb._json_handlers import is_json_payload -from sqlspec.adapters.oracledb._param_types import OracleBlob, OracleClob, OracleJson +from sqlspec.adapters.oracledb._param_types import OracleBlob, OracleClob, OracleJson, is_json_payload from sqlspec.adapters.oracledb._typing import DB_TYPE_BLOB, DB_TYPE_CLOB, oracledb_module -from sqlspec.adapters.oracledb.data_dictionary import ( - resolve_oracle_connection_compatible_major, - resolve_oracle_connection_major, -) -from sqlspec.adapters.oracledb.type_converter import OracleOutputConverter from sqlspec.core import ( DriverParameterProfile, ParameterStyle, @@ -33,7 +28,11 @@ create_sql_result, ) from sqlspec.core.sqlcommenter import SQLCommenterContext -from sqlspec.data_dictionary.dialects.oracle import ORACLE_JSON_STORAGE_NATIVE, resolve_oracle_json_storage +from sqlspec.data_dictionary.dialects.oracle import ( + ORACLE_JSON_STORAGE_NATIVE, + parse_oracle_version_components, + resolve_oracle_json_storage, +) from sqlspec.driver import rows_to_dicts from sqlspec.exceptions import ( CheckViolationError, @@ -48,6 +47,8 @@ OperationalError, OperationCancelledError, PermissionDeniedError, + QueryTimeoutError, + SerializationConflictError, SQLParsingError, SQLSpecError, TransactionError, @@ -56,12 +57,115 @@ from sqlspec.observability import get_trace_context from sqlspec.typing import NUMPY_INSTALLED from sqlspec.utils.correlation import CorrelationContext -from sqlspec.utils.module_loader import ensure_pyarrow, import_optional +from sqlspec.utils.module_loader import import_optional from sqlspec.utils.serializers import from_json, to_json from sqlspec.utils.text import quote_identifier from sqlspec.utils.type_converters import build_uuid_coercions from sqlspec.utils.type_guards import has_rowcount, is_async_readable, is_readable + +class OracleVersionCache: + """Pool-scoped cache for the resolved Oracle server version. + + One instance is created per config and shared with every driver spawned from + that config's pool, so the server version is resolved once per pool lifetime + rather than once per acquired session. The cache lives and dies with the + pool; a server upgrade under a live pool is out of scope and picked up on the + next pool restart. + """ + + __slots__ = ("resolved", "version") + + def __init__(self) -> None: + self.resolved: bool = False + self.version: Any | None = None + + def reset(self) -> None: + """Clear the cached version so the next resolution re-queries the server.""" + self.resolved = False + self.version = None + + +def resolve_oracle_connection_major(connection: Any, version_cache: "OracleVersionCache | None" = None) -> "int | None": + """Resolve an Oracle server major without issuing a metadata query. + + Pool-scoped data-dictionary metadata is authoritative when it has already + been resolved. Otherwise, use the connection callback cache and finally + python-oracledb's connection ``version`` string. The final fallback matters + for reacquired pool wrappers that do not retain dynamic attributes. + """ + if version_cache is not None and version_cache.resolved and version_cache.version is not None: + return int(version_cache.version.major) + cached_major = getattr(connection, "_sqlspec_oracle_major", None) + if isinstance(cached_major, int) and not isinstance(cached_major, bool): + return cached_major + version = getattr(connection, "version", None) + if not version: + return None + components = parse_oracle_version_components(str(version)) + return components[0] if components is not None else None + + +def resolve_oracle_connection_compatible_major( + connection: Any, version_cache: "OracleVersionCache | None" = None +) -> "int | None": + """Resolve an Oracle server compatible major version without issuing a metadata query.""" + if version_cache is not None and version_cache.resolved and version_cache.version is not None: + comp = version_cache.version.compatible_major + return int(comp) if comp is not None else None + cached_compatible = getattr(connection, "_sqlspec_oracle_compatible_major", None) + if isinstance(cached_compatible, int) and not isinstance(cached_compatible, bool): + return cached_compatible + return None + + +def chain_input_handler(inner: Any, fallback: "Any | None") -> Any: + """Build an input type handler that chains ``inner`` to ``fallback``.""" + if fallback is not None and _has_handler(fallback, inner, _chained_input_handler): + return fallback + return partial(_chained_input_handler, inner, fallback) + + +def chain_output_handler(inner: Any, fallback: "Any | None") -> Any: + """Build an output type handler that chains ``inner`` to ``fallback``.""" + if fallback is not None and _has_handler(fallback, inner, _chained_output_handler): + return fallback + return partial(_chained_output_handler, inner, fallback) + + +def _has_handler(handler: Any, target_inner: Any, chain_function: Any) -> bool: + current = handler + while current is not None: + if current is target_inner: + return True + if not isinstance(current, partial) or current.func is not chain_function: + return False + inner, current = current.args + if inner is target_inner: + return True + return False + + +def _chained_input_handler(inner: Any, fallback: "Any | None", cursor: Any, value: Any, arraysize: int) -> Any: + """Run ``inner`` input handler, falling back to ``fallback`` when it abstains.""" + result = inner(cursor, value, arraysize) + if result is not None: + return result + if fallback is not None: + return fallback(cursor, value, arraysize) + return None + + +def _chained_output_handler(inner: Any, fallback: "Any | None", cursor: Any, metadata: Any) -> Any: + """Run ``inner`` output handler, falling back to ``fallback`` when it abstains.""" + result = inner(cursor, metadata) + if result is not None: + return result + if fallback is not None: + return fallback(cursor, metadata) + return None + + if TYPE_CHECKING: from collections.abc import Awaitable @@ -73,7 +177,6 @@ ) from sqlspec.core import SQL from sqlspec.driver import AsyncExceptionHandler, SyncExceptionHandler - from sqlspec.typing import ArrowRecordBatch, ArrowRecordBatchReader, ArrowTable __all__ = ( @@ -95,6 +198,7 @@ "OracleSessionContextConfig", "OracleSyncStreamSource", "OracleTopologyConfig", + "OracleVersionCache", "apply_driver_features", "apply_oracle_session_trace_attributes", "apply_wallet_ssl_context", @@ -115,12 +219,12 @@ "build_tns_description", "build_truncate_statement", "build_wallet_ssl_context_if_needed", + "chain_input_handler", + "chain_output_handler", "coerce_large_parameters_async", "coerce_large_parameters_sync", "coerce_many_parameters_async", "coerce_many_parameters_sync", - "coerce_oracle_df_to_arrow_batches", - "coerce_oracle_df_to_arrow_table", "collect_async_rows", "collect_sync_rows", "connection_is_thin", @@ -138,12 +242,14 @@ "normalize_column_names", "normalize_execute_many_parameters_async", "normalize_execute_many_parameters_sync", - "oracle_df_batches_to_arrow_reader", - "oracle_df_to_arrow_table", "parse_sqlnet_ora", "parse_tnsnames_ora", + "read_oracle_lob_async", + "read_oracle_lob_sync", "release_extracted_wallet_dir", "resolve_dbms_cloud_source_uri", + "resolve_oracle_connection_compatible_major", + "resolve_oracle_connection_major", "resolve_row_metadata", "resolve_rowcount", "retain_extracted_wallet_dir", @@ -267,8 +373,6 @@ "WITH", }) _VERSION_COMPONENTS: int = 3 -TYPE_CONVERTER = OracleOutputConverter() -_pa: Any = import_optional("pyarrow") _LOB_TYPE_NAME_MARKERS: "tuple[str, ...]" = ("LOB", "BFILE") _SCALAR_PASSTHROUGH_TYPES: "tuple[type[Any], ...]" = (bool, int, float, str, bytes, bytearray, type(None)) _BIND_PASSTHROUGH_TYPES: "tuple[type[Any], ...]" = (bool, int, float, type(None)) @@ -293,6 +397,7 @@ 1400: (NotNullViolationError, "not-null constraint violation"), 1407: (NotNullViolationError, "not-null constraint violation"), 1438: (DataError, "numeric precision overflow"), + 1476: (DataError, "divisor is equal to zero"), 1652: (OperationalError, "tablespace full"), 1722: (DataError, "invalid number"), 1830: (DataError, "date format picture ends before converting entire input string"), @@ -306,9 +411,10 @@ 3113: (DatabaseConnectionError, "end-of-file on communication channel"), 3114: (DatabaseConnectionError, "not connected to ORACLE"), 3135: (DatabaseConnectionError, "connection lost contact"), + 3136: (ConnectionTimeoutError, "inbound connection timed out"), 4068: (OperationalError, "existing state of packages has been discarded"), 8176: (TransactionError, "consistent read failure"), - 8177: (TransactionError, "cannot serialize access"), + 8177: (SerializationConflictError, "cannot serialize access"), 12154: (DatabaseConnectionError, "TNS resolution failed"), 12170: (ConnectionTimeoutError, "connect timeout"), 12505: (DatabaseConnectionError, "listener rejected"), @@ -319,11 +425,15 @@ 28000: (PermissionDeniedError, "the account is locked"), 28001: (PermissionDeniedError, "the password has expired"), 30006: (OperationalError, "resource busy; acquire with WAIT timeout expired"), + 40441: (DataError, "JSON syntax error"), + 42699: (SerializationConflictError, "transaction isolation violation"), + 51803: (DataError, "vector operation error"), } _DPY_ERROR_CODE_MAPPING: "dict[str, tuple[type[SQLSpecError], str]]" = { "DPY-1001": (DatabaseConnectionError, "not connected to database"), "DPY-4011": (DatabaseConnectionError, "connection was closed"), + "DPY-4024": (QueryTimeoutError, "query execution timeout expired"), "DPY-6000": (DatabaseConnectionError, "cannot connect to database"), "DPY-6001": (DatabaseConnectionError, "cannot connect to database"), "DPY-6005": (ConnectionTimeoutError, "cannot connect to database"), @@ -787,15 +897,24 @@ def supports_df_batches(connection: object) -> bool: return hasattr(connection, "fetch_df_batches") +_FETCH_KWARGS_CACHE: dict[tuple[object, object], dict[str, object]] = {} + + def build_fetch_kwargs(driver_features: "dict[str, Any]") -> "dict[str, object]": """Build per-statement Oracle fetch keyword arguments.""" - fetch_kwargs: dict[str, object] = {} fetch_lobs = driver_features.get("fetch_lobs") + fetch_decimals = driver_features.get("fetch_decimals") + cache_key = (fetch_lobs, fetch_decimals) + cached = _FETCH_KWARGS_CACHE.get(cache_key) + if cached is not None: + return cached + + fetch_kwargs: dict[str, object] = {} if fetch_lobs is not None: fetch_kwargs["fetch_lobs"] = fetch_lobs - fetch_decimals = driver_features.get("fetch_decimals") if fetch_decimals is not None: fetch_kwargs["fetch_decimals"] = fetch_decimals + _FETCH_KWARGS_CACHE[cache_key] = fetch_kwargs return fetch_kwargs @@ -834,71 +953,6 @@ def build_arrow_fetch_kwargs( return fetch_kwargs -def oracle_df_to_arrow_table(oracle_df: Any, pa_module: Any = None, arrow_schema: Any = None) -> "ArrowTable": - """Convert an Oracle DataFrame to a PyArrow Table using zero-copy PyCapsule when supported.""" - if pa_module is None: - ensure_pyarrow() - pa_module = _pa - - if hasattr(oracle_df, "__arrow_c_stream__"): - reader = ( - pa_module.RecordBatchReader.from_stream(oracle_df, schema=arrow_schema) - if arrow_schema is not None - else pa_module.RecordBatchReader.from_stream(oracle_df) - ) - return cast("ArrowTable", reader.read_all()) - table = pa_module.table(oracle_df) - if arrow_schema is not None: - table = table.cast(arrow_schema) - return cast("ArrowTable", table) - - -def coerce_oracle_df_to_arrow_table(oracle_df: Any, pa_module: Any = None, arrow_schema: Any = None) -> "ArrowTable": - """Alias for :func:`oracle_df_to_arrow_table`.""" - return oracle_df_to_arrow_table(oracle_df, pa_module=pa_module, arrow_schema=arrow_schema) - - -def coerce_oracle_df_to_arrow_batches( - oracle_df: Any, pa_module: Any = None, arrow_schema: Any = None -) -> "list[ArrowRecordBatch]": - """Convert an Oracle DataFrame batch into PyArrow RecordBatches via PyCapsule when supported.""" - if pa_module is None: - ensure_pyarrow() - pa_module = _pa - - if hasattr(oracle_df, "__arrow_c_stream__"): - reader = ( - pa_module.RecordBatchReader.from_stream(oracle_df, schema=arrow_schema) - if arrow_schema is not None - else pa_module.RecordBatchReader.from_stream(oracle_df) - ) - return cast("list[ArrowRecordBatch]", list(reader)) - table = pa_module.table(oracle_df) - if arrow_schema is not None: - table = table.cast(arrow_schema) - return cast("list[ArrowRecordBatch]", table.to_batches()) - - -def oracle_df_batches_to_arrow_reader( - batches: "Iterable[Any]", schema: Any = None, pa_module: Any = None -) -> "ArrowRecordBatchReader": - """Convert an iterable of Oracle DataFrame batches into a PyArrow RecordBatchReader.""" - if pa_module is None: - ensure_pyarrow() - pa_module = _pa - - collected_batches: list[Any] = [] - for df_batch in batches: - collected_batches.extend(coerce_oracle_df_to_arrow_batches(df_batch, pa_module=pa_module, arrow_schema=schema)) - if collected_batches: - resolved_schema = schema if schema is not None else collected_batches[0].schema - return cast( - "ArrowRecordBatchReader", pa_module.RecordBatchReader.from_batches(resolved_schema, collected_batches) - ) - resolved_schema = schema if schema is not None else pa_module.schema([]) - return cast("ArrowRecordBatchReader", pa_module.RecordBatchReader.from_batches(resolved_schema, [])) - - def normalize_column_names(column_names: "list[str]", driver_features: "dict[str, Any]") -> "list[str]": should_lowercase = driver_features.get("enable_lowercase_column_names", False) if not should_lowercase: @@ -1196,7 +1250,7 @@ def resolve_rowcount(cursor: Any) -> int: if not has_rowcount(cursor): return 0 rowcount = cursor.rowcount - if isinstance(rowcount, int): + if isinstance(rowcount, int) and rowcount >= 0: return rowcount return 0 @@ -1317,6 +1371,20 @@ def apply_driver_features( return statement_config, features +def _extract_column_cache_key(description: "list[Any]") -> "tuple[tuple[object, object], ...]": + key_items: list[tuple[object, object]] = [] + for column in description: + try: + col_name = column[0] + type_code = column[1] + except (TypeError, IndexError, KeyError): + col_name = getattr(column, "name", None) + type_code = getattr(column, "type_code", None) + type_identifier = getattr(type_code, "name", None) or type_code + key_items.append((col_name, type_identifier)) + return tuple(key_items) + + def resolve_row_metadata( description: "list[Any] | None", driver_features: "dict[str, Any]", cache: "dict[Any, tuple[Any, list[str], bool]]" ) -> "tuple[list[str], bool]": @@ -1333,9 +1401,9 @@ def resolve_row_metadata( if not description: return [], False - cache_key = id(description) + cache_key = _extract_column_cache_key(description) cached = cache.get(cache_key) - if cached is not None and cached[0] is description: + if cached is not None: return cached[1], cached[2] column_names: list[str] = [] @@ -1408,41 +1476,44 @@ def __init__( def start(self) -> None: handler = self._driver.handle_database_exceptions() - with handler: - try: - cursor = self._driver.connection.cursor() - self._cursor = cursor - cursor.arraysize = self._chunk_size - cursor.prefetchrows = self._chunk_size - fetch_kwargs = build_fetch_kwargs(self._driver.driver_features) - if self._fetch_lobs is not None: - fetch_kwargs["fetch_lobs"] = self._fetch_lobs - parameters = coerce_large_parameters_sync( - self._driver.connection, - self._parameters, - clob_type=DB_TYPE_CLOB, - blob_type=DB_TYPE_BLOB, - varchar2_byte_limit=self._driver.driver_features.get("oracle_varchar2_byte_limit", 4000), - raw_byte_limit=self._driver.driver_features.get("oracle_raw_byte_limit", 2000), - version_cache=getattr(self._driver, "_oracle_version_cache", None), - ) - cast("Any", cursor).execute(self._sql, parameters or {}, **fetch_kwargs) - except BaseException: - self.close(error=True) - raise + self._driver._run_with_exception_handler(handler, self._start) self._driver._check_pending_exception(handler) + def _start(self) -> None: + try: + cursor = self._driver.connection.cursor() + self._cursor = cursor + cursor.arraysize = self._chunk_size + cursor.prefetchrows = self._chunk_size + fetch_kwargs = getattr(self._driver, "_fetch_kwargs", None) or build_fetch_kwargs( + self._driver.driver_features + ) + if self._fetch_lobs is not None: + fetch_kwargs = {**fetch_kwargs, "fetch_lobs": self._fetch_lobs} + parameters = coerce_large_parameters_sync( + self._driver.connection, + self._parameters, + clob_type=DB_TYPE_CLOB, + blob_type=DB_TYPE_BLOB, + varchar2_byte_limit=self._driver.driver_features.get("oracle_varchar2_byte_limit", 4000), + raw_byte_limit=self._driver.driver_features.get("oracle_raw_byte_limit", 2000), + version_cache=getattr(self._driver, "_oracle_version_cache", None), + ) + cast("Any", cursor).execute(self._sql, parameters or {}, **fetch_kwargs) + except BaseException: + self.close(error=True) + raise + def fetch_chunk(self) -> "list[dict[str, Any]]": handler = self._driver.handle_database_exceptions() cursor = self._cursor if cursor is None: return [] - rows: list[Any] = [] - with handler: - rows = cursor.fetchmany(self._chunk_size) + fetched = self._driver._run_with_exception_handler(handler, cursor.fetchmany, self._chunk_size) self._driver._check_pending_exception(handler) - if not rows: + if not fetched: return [] + rows: list[Any] = list(fetched) column_names = self._column_names if column_names is None: column_names, requires_lob_coercion = self._driver._resolve_row_metadata(cursor.description) @@ -1553,9 +1624,11 @@ async def _start(self) -> None: self._cursor = cursor cursor.arraysize = self._chunk_size cursor.prefetchrows = self._chunk_size - fetch_kwargs = build_fetch_kwargs(self._driver.driver_features) + fetch_kwargs = getattr(self._driver, "_fetch_kwargs", None) or build_fetch_kwargs( + self._driver.driver_features + ) if self._fetch_lobs is not None: - fetch_kwargs["fetch_lobs"] = self._fetch_lobs + fetch_kwargs = {**fetch_kwargs, "fetch_lobs": self._fetch_lobs} parameters = await coerce_large_parameters_async( self._driver.connection, self._parameters, @@ -1619,6 +1692,9 @@ def collect_sync_rows( if not fetched_data: return [], resolved_column_names + if driver_features.get("fetch_lobs"): + return [row if isinstance(row, tuple) else tuple(row) for row in fetched_data], resolved_column_names + if requires_lob_coercion is None: requires_lob_coercion = _description_requires_lob_coercion(description) if not requires_lob_coercion: @@ -1670,6 +1746,9 @@ async def collect_async_rows( if not fetched_data: return [], resolved_column_names + if driver_features.get("fetch_lobs"): + return [row if isinstance(row, tuple) else tuple(row) for row in fetched_data], resolved_column_names + if requires_lob_coercion is None: requires_lob_coercion = _description_requires_lob_coercion(description) if not requires_lob_coercion: @@ -1700,6 +1779,8 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpec A SQLSpec exception that wraps the original error """ del logger + if not isinstance(error, BaseException): + error = oracledb_module.DatabaseError(error) error_obj = error.args[0] if getattr(error, "args", None) else None if not error_obj: return _create_oracle_error(error, None, SQLSpecError, "database error") @@ -1725,9 +1806,13 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpec if not isinstance(error_code, int) or isinstance(error_code, bool) or not error_code: error_text = str(error_obj) for dpy_code, (error_class, error_desc) in _DPY_ERROR_CODE_MAPPING.items(): - if error_text.startswith(dpy_code): + if dpy_code in error_text: return _create_oracle_error(error, None, error_class, f"{error_desc} [{dpy_code}]") - return _create_oracle_error(error, None, SQLSpecError, "database error") + ora_match = re.search(r"ORA-(\d+)", error_text) + if ora_match: + error_code = int(ora_match.group(1)) + else: + return _create_oracle_error(error, None, SQLSpecError, "database error") mapping = _ERROR_CODE_MAPPING.get(error_code) if mapping: @@ -3129,6 +3214,22 @@ def _resolve_callable_or_str(value: "str | Callable[[], str | None] | None") -> return value +async def read_oracle_lob_async(value: Any) -> Any: + """Read a LOB value asynchronously from an async connection or return raw data.""" + if is_async_readable(value): + return await value.read() + if is_readable(value): + return value.read() + return value + + +def read_oracle_lob_sync(value: Any) -> Any: + """Read a LOB value synchronously from a sync connection or return raw data.""" + if is_readable(value): + return value.read() + return value + + class _OracleSyncStreamDriver(Protocol): connection: "OracleSyncConnection" driver_features: dict[str, Any] @@ -3137,6 +3238,10 @@ def handle_database_exceptions(self) -> "SyncExceptionHandler": ... def _check_pending_exception(self, exc_handler: "SyncExceptionHandler") -> None: ... + def _run_with_exception_handler( + self, exc_handler: "SyncExceptionHandler", operation: "Callable[..., Any]", *args: Any, **kwargs: Any + ) -> "Any | None": ... + def _resolve_row_metadata(self, description: object) -> tuple[list[str], bool]: ... @@ -3195,6 +3300,16 @@ def _parse_version_tuple(version: str) -> "tuple[int, int, int]": ORACLEDB_VERSION: "tuple[int, int, int]" = _parse_version_tuple(oracledb_module.__version__) +def _string_exceeds_byte_limit(value: str, limit: int) -> bool: + """Return whether a string's UTF-8 byte length exceeds the given limit.""" + char_len = len(value) + if char_len > limit: + return True + if value.isascii(): + return False + return len(value.encode("utf-8")) > limit + + def _parameter_values_need_coercion( values: "Iterable[Any]", *, @@ -3208,7 +3323,7 @@ def _parameter_values_need_coercion( if value_type in _BIND_PASSTHROUGH_TYPES: continue if value_type is str: - if len(value.encode("utf-8")) > varchar2_byte_limit: + if _string_exceeds_byte_limit(value, varchar2_byte_limit): return True continue if value_type is bytes or value_type is bytearray: @@ -3226,7 +3341,7 @@ def _parameter_values_need_coercion( return True continue if isinstance(value, str): - if len(value.encode("utf-8")) > varchar2_byte_limit: + if _string_exceeds_byte_limit(value, varchar2_byte_limit): return True continue if isinstance(value, (bytes, bytearray)): @@ -3379,7 +3494,7 @@ def _coerce_value_sync( if value_type in _BIND_PASSTHROUGH_TYPES: return value if value_type is str: - if len(value.encode("utf-8")) > varchar2_byte_limit: + if _string_exceeds_byte_limit(value, varchar2_byte_limit): return connection.createlob(clob_type, value) return value if value_type is bytes: @@ -3408,7 +3523,7 @@ def _coerce_value_sync( if json_binding_state.uses_blob(): return connection.createlob(blob_type, to_json(value.value, as_bytes=True)) return value - if isinstance(value, str) and len(value.encode("utf-8")) > varchar2_byte_limit: + if isinstance(value, str) and _string_exceeds_byte_limit(value, varchar2_byte_limit): return connection.createlob(clob_type, value) if isinstance(value, (bytes, bytearray)) and len(value) > raw_byte_limit: return connection.createlob(blob_type, bytes(value)) @@ -3432,7 +3547,7 @@ async def _coerce_value_async( if value_type in _BIND_PASSTHROUGH_TYPES: return value if value_type is str: - if len(value.encode("utf-8")) > varchar2_byte_limit: + if _string_exceeds_byte_limit(value, varchar2_byte_limit): return await connection.createlob(clob_type, value) return value if value_type is bytes: @@ -3461,7 +3576,7 @@ async def _coerce_value_async( if json_binding_state.uses_blob(): return await connection.createlob(blob_type, to_json(value.value, as_bytes=True)) return value - if isinstance(value, str) and len(value.encode("utf-8")) > varchar2_byte_limit: + if isinstance(value, str) and _string_exceeds_byte_limit(value, varchar2_byte_limit): return await connection.createlob(clob_type, value) if isinstance(value, (bytes, bytearray)) and len(value) > raw_byte_limit: return await connection.createlob(blob_type, bytes(value)) @@ -3570,7 +3685,7 @@ async def _coerce_async_row_values(row: "tuple[Any, ...]") -> "tuple[Any, ...]": continue if is_readable(value) or is_async_readable(value): - processed_value = await TYPE_CONVERTER.process_lob(value) + processed_value = await read_oracle_lob_async(value) if coerced_values is None: if processed_value is value: continue diff --git a/sqlspec/adapters/oracledb/data_dictionary.py b/sqlspec/adapters/oracledb/data_dictionary.py index 4c4a0f8fb..28e96af0e 100644 --- a/sqlspec/adapters/oracledb/data_dictionary.py +++ b/sqlspec/adapters/oracledb/data_dictionary.py @@ -1,10 +1,17 @@ """Oracle-specific data dictionary for metadata queries.""" +from collections import defaultdict from enum import StrEnum +from inspect import isawaitable from typing import TYPE_CHECKING, Any, ClassVar, cast from mypy_extensions import mypyc_attr +from sqlspec.adapters.oracledb.core import ( + OracleVersionCache, + resolve_oracle_connection_compatible_major, + resolve_oracle_connection_major, +) from sqlspec.data_dictionary import ( ColumnMetadata, DDLResult, @@ -31,7 +38,6 @@ from sqlspec.data_dictionary.dialects.oracle import ( extract_oracle_version_value, list_oracle_available_features, - merge_oracle_table_lists, oracle_supports_json_blob, oracle_supports_native_json, oracle_supports_oson_blob, @@ -49,7 +55,7 @@ from collections.abc import Sequence from sqlspec.adapters.oracledb.driver import OracleAsyncDriver, OracleSyncDriver - from sqlspec.data_dictionary._types import DialectConfig + from sqlspec.data_dictionary import DialectConfig __all__ = ( "JSONStorageType", @@ -169,63 +175,36 @@ class JSONStorageType(StrEnum): BLOB_PLAIN = "blob_plain" -class OracleVersionCache: - """Pool-scoped cache for the resolved Oracle server version. - - One instance is created per config and shared with every driver spawned from - that config's pool, so the server version is resolved once per pool lifetime - rather than once per acquired session. The cache lives and dies with the - pool; a server upgrade under a live pool is out of scope and picked up on the - next pool restart. - """ - - __slots__ = ("resolved", "version") - - def __init__(self) -> None: - self.resolved: bool = False - self.version: OracleVersionInfo | None = None - - def reset(self) -> None: - """Clear the cached version so the next resolution re-queries the server.""" - self.resolved = False - self.version = None - - def storage_type_from_version(version_info: "OracleVersionInfo | None") -> JSONStorageType: """Public alias for :func:`_storage_type_from_version`.""" return _storage_type_from_version(version_info) -def resolve_oracle_connection_major(connection: Any, version_cache: "OracleVersionCache | None" = None) -> "int | None": - """Resolve an Oracle server major without issuing a metadata query. +def _merge_index_expressions( + indexes: "list[IndexMetadata]", expr_rows: "list[dict[str, Any]]", schema_name: "str | None" = None +) -> "list[IndexMetadata]": + """Merge expression-based index columns into index metadata records.""" + if not expr_rows: + return indexes - Pool-scoped data-dictionary metadata is authoritative when it has already - been resolved. Otherwise, use the connection callback cache and finally - python-oracledb's connection ``version`` string. The final fallback matters - for reacquired pool wrappers that do not retain dynamic attributes. - """ - if version_cache is not None and version_cache.resolved and version_cache.version is not None: - return version_cache.version.major - cached_major = getattr(connection, "_sqlspec_oracle_major", None) - if isinstance(cached_major, int) and not isinstance(cached_major, bool): - return cached_major - version = getattr(connection, "version", None) - if not version: - return None - components = parse_oracle_version_components(str(version)) - return components[0] if components is not None else None + exprs_by_index: dict[tuple[str, str], list[str]] = defaultdict(list) + for row in expr_rows: + owner = str(row.get("schema_name") or row.get("index_owner") or schema_name or "").upper() + idx_name = str(row.get("index_name") or "").upper() + expr = row.get("column_expression") + if expr is not None: + exprs_by_index[(owner, idx_name)].append(str(expr)) + for idx in indexes: + owner = str(idx.get("schema_name") or schema_name or "").upper() + idx_name = str(idx.get("index_name") or "").upper() + matching_exprs = exprs_by_index.get((owner, idx_name)) + if matching_exprs: + expr_str = ",".join(matching_exprs) + idx["expressions"] = expr_str + idx["expression"] = expr_str -def resolve_oracle_connection_compatible_major( - connection: Any, version_cache: "OracleVersionCache | None" = None -) -> "int | None": - """Resolve an Oracle server compatible major version without issuing a metadata query.""" - if version_cache is not None and version_cache.resolved and version_cache.version is not None: - return version_cache.version.compatible_major - cached_compatible = getattr(connection, "_sqlspec_oracle_compatible_major", None) - if isinstance(cached_compatible, int) and not isinstance(cached_compatible, bool): - return cached_compatible - return None + return indexes @mypyc_attr(allow_interpreted_subclasses=True, native_class=False) @@ -439,13 +418,9 @@ def get_tables(self, driver: "OracleSyncDriver", schema: "str | None" = None) -> """Get tables sorted by dependency order with full coverage.""" schema_name = self.resolve_schema(schema) self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="tables") - ordered_rows = driver.select( + return driver.select( self.get_query("tables", "by_schema"), schema_name=schema_name, table_name=None, schema_type=TableMetadata ) - all_rows = driver.select( - self.get_query("tables", "names_by_schema"), schema_name=schema_name, schema_type=TableMetadata - ) - return merge_oracle_table_lists(ordered_rows, all_rows) def get_columns( self, driver: "OracleSyncDriver", table: "str | None" = None, schema: "str | None" = None @@ -473,17 +448,25 @@ def get_indexes( schema_name = self.resolve_schema(schema) if table is None: self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="indexes") - return driver.select( + indexes: list[IndexMetadata] = driver.select( self.get_query("indexes", "by_schema"), schema_name=schema_name, schema_type=IndexMetadata ) - table_name = self.resolve_identifier(table) - self._log_table_describe(driver, schema_name=schema_name, table_name=table_name, operation="indexes") - return driver.select( - self.get_query("indexes", "by_table"), - schema_name=schema_name, - table_name=table_name, - schema_type=IndexMetadata, - ) + expr_rows: list[dict[str, Any]] = driver.select( + self.get_query("indexes", "expressions"), schema_name=schema_name, table_name=None + ) + else: + table_name = self.resolve_identifier(table) + self._log_table_describe(driver, schema_name=schema_name, table_name=table_name, operation="indexes") + indexes = driver.select( + self.get_query("indexes", "by_table"), + schema_name=schema_name, + table_name=table_name, + schema_type=IndexMetadata, + ) + expr_rows = driver.select( + self.get_query("indexes", "expressions"), schema_name=schema_name, table_name=table_name + ) + return _merge_index_expressions(indexes, expr_rows, schema_name=schema_name) def get_foreign_keys( self, driver: "OracleSyncDriver", table: "str | None" = None, schema: "str | None" = None @@ -856,13 +839,9 @@ async def get_tables(self, driver: "OracleAsyncDriver", schema: "str | None" = N """Get tables sorted by dependency order with full coverage.""" schema_name = self.resolve_schema(schema) self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="tables") - ordered_rows = await driver.select( + return await driver.select( self.get_query("tables", "by_schema"), schema_name=schema_name, table_name=None, schema_type=TableMetadata ) - all_rows = await driver.select( - self.get_query("tables", "names_by_schema"), schema_name=schema_name, schema_type=TableMetadata - ) - return merge_oracle_table_lists(ordered_rows, all_rows) async def get_columns( self, driver: "OracleAsyncDriver", table: "str | None" = None, schema: "str | None" = None @@ -890,17 +869,25 @@ async def get_indexes( schema_name = self.resolve_schema(schema) if table is None: self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="indexes") - return await driver.select( + indexes: list[IndexMetadata] = await driver.select( self.get_query("indexes", "by_schema"), schema_name=schema_name, schema_type=IndexMetadata ) - table_name = self.resolve_identifier(table) - self._log_table_describe(driver, schema_name=schema_name, table_name=table_name, operation="indexes") - return await driver.select( - self.get_query("indexes", "by_table"), - schema_name=schema_name, - table_name=table_name, - schema_type=IndexMetadata, - ) + expr_rows: list[dict[str, Any]] = await driver.select( + self.get_query("indexes", "expressions"), schema_name=schema_name, table_name=None + ) + else: + table_name = self.resolve_identifier(table) + self._log_table_describe(driver, schema_name=schema_name, table_name=table_name, operation="indexes") + indexes = await driver.select( + self.get_query("indexes", "by_table"), + schema_name=schema_name, + table_name=table_name, + schema_type=IndexMetadata, + ) + expr_rows = await driver.select( + self.get_query("indexes", "expressions"), schema_name=schema_name, table_name=table_name + ) + return _merge_index_expressions(indexes, expr_rows, schema_name=schema_name) async def get_foreign_keys( self, driver: "OracleAsyncDriver", table: "str | None" = None, schema: "str | None" = None @@ -993,11 +980,19 @@ async def _get_native_object_ddl( ddl_value = await driver.select_value( query.sql, object_type=normalized_type, object_name=object_name, owner=schema ) + if ddl_value is not None: + read = getattr(ddl_value, "read", None) + if callable(read): + res = read() + if isawaitable(res): + ddl_value = await res + else: + ddl_value = res return _oracle_ddl_metadata_result( object_type=normalized_type, object_name=object_name, owner=schema, - ddl_text=_coerce_oracle_ddl_text(ddl_value), + ddl_text=str(ddl_value) if ddl_value is not None else None, ) def _build_version_info( diff --git a/sqlspec/adapters/oracledb/driver.py b/sqlspec/adapters/oracledb/driver.py index 174e6595f..53034807a 100644 --- a/sqlspec/adapters/oracledb/driver.py +++ b/sqlspec/adapters/oracledb/driver.py @@ -1,8 +1,10 @@ """Oracle Database driver implementation.""" import logging +from itertools import chain from pathlib import Path -from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, Protocol, cast, overload +from time import perf_counter +from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, NamedTuple, Protocol, cast, overload from sqlspec.adapters.oracledb._typing import ( DB_TYPE_BLOB, @@ -32,7 +34,6 @@ coerce_large_parameters_sync, coerce_many_parameters_async, coerce_many_parameters_sync, - coerce_oracle_df_to_arrow_table, collect_async_rows, collect_sync_rows, connection_is_thin, @@ -54,6 +55,8 @@ ) from sqlspec.core import ( SQL, + DMLResult, + SQLResult, StackResult, StatementConfig, StatementStack, @@ -81,18 +84,22 @@ hash_stack_operations, validate_savepoint_name, ) +from sqlspec.driver._common import CommonDriverAttributesMixin from sqlspec.exceptions import ImproperConfigurationError, SQLSpecError, StackExecutionError +from sqlspec.storage._utils import import_pyarrow +from sqlspec.utils.arrow_helpers import coerce_oracle_df_to_arrow_table from sqlspec.utils.logging import get_logger, log_with_context -from sqlspec.utils.module_loader import ensure_pyarrow, import_optional +from sqlspec.utils.module_loader import ensure_pyarrow from sqlspec.utils.text import normalize_identifier, quote_identifier, split_qualified_identifier from sqlspec.utils.type_guards import has_pipeline_capability, is_async_readable from sqlspec.utils.uuids import uuid4 if TYPE_CHECKING: - from collections.abc import Mapping, Sequence + from collections.abc import Iterator, Mapping, Sequence from sqlspec.builder import QueryBuilder from sqlspec.core import ArrowResult, Statement, StatementFilter + from sqlspec.core.cache import CachedQuery from sqlspec.core.stack import StackOperation from sqlspec.driver import ExecutionResult from sqlspec.storage import StorageBridgeJob, StorageDestination, StorageFormat, StorageTelemetry @@ -109,7 +116,7 @@ logger = get_logger(__name__) -pa: Any = import_optional("pyarrow") + PLAN_STATEMENT_ID_MAX_LENGTH: Final[int] = 30 """Width of the PLAN_TABLE STATEMENT_ID column.""" @@ -151,19 +158,13 @@ def _compiled_sql(self, statement: "SQL", statement_config: "StatementConfig") - """ -class OraclePipelineMixin: +class OraclePipelineMixin(CommonDriverAttributesMixin): """Shared helpers for Oracle pipeline execution.""" __slots__ = () - def _session_tracing_enabled(self) -> bool: - driver = cast("OraclePipelineDriver", self) - return bool( - driver.driver_features.get("enable_session_tracing", False) or driver.statement_config.enable_sqlcommenter - ) - def _apply_session_tracing(self, connection: Any = None) -> None: - if not self._session_tracing_enabled(): + if not getattr(self, "_session_tracing_enabled", False): return driver = cast("OraclePipelineDriver", self) target_conn = connection if connection is not None else getattr(self, "connection", None) @@ -186,7 +187,7 @@ def _apply_session_tracing(self, connection: Any = None) -> None: def _format_wire_sql(self, sql: str, config: "StatementConfig | None" = None) -> str: driver = cast("OraclePipelineDriver", self) effective_config = config or driver.statement_config - if self._session_tracing_enabled(): + if getattr(self, "_session_tracing_enabled", False): return strip_sqlcomment(sql) if not effective_config.enable_sqlcommenter: return sql @@ -195,10 +196,10 @@ def _format_wire_sql(self, sql: str, config: "StatementConfig | None" = None) -> def _compiled_sql( self, statement: "SQL", statement_config: "StatementConfig", flatten_single_parameters: bool = False ) -> "tuple[str, Any]": - compiled_sql, prepared_parameters = super()._compiled_sql( # type: ignore[misc] + compiled_sql, prepared_parameters = super()._compiled_sql( statement, statement_config, flatten_single_parameters=flatten_single_parameters ) - if self._session_tracing_enabled(): + if getattr(self, "_session_tracing_enabled", False): return strip_sqlcomment(compiled_sql), prepared_parameters return compiled_sql, prepared_parameters @@ -250,9 +251,7 @@ def _prepare_pipeline_operation(self, operation: "StackOperation") -> "_Compiled msg = f"Unsupported stack operation method: {operation.method}" raise ValueError(msg) - compiled_sql, prepared_parameters = driver._compiled_sql( # pyright: ignore[reportPrivateUsage] - sql_statement, config - ) + compiled_sql, prepared_parameters = self._compiled_sql(sql_statement, config) summary = describe_stack_statement(operation.statement) return _CompiledStackOperation( statement=sql_statement, @@ -295,10 +294,11 @@ def _build_stack_results_from_pipeline( except AttributeError: error = None if error is not None: + mapped_error = create_mapped_exception(error) if not isinstance(error, SQLSpecError) else error stack_error = StackExecutionError( index, compiled.summary, - error, + mapped_error, adapter=type(self).__name__, mode="continue-on-error" if continue_on_error else "fail-fast", ) @@ -324,8 +324,9 @@ def _wrap_pipeline_error( self, error: Exception, stack: "StatementStack", continue_on_error: bool ) -> StackExecutionError: mode = "continue-on-error" if continue_on_error else "fail-fast" + mapped_error = create_mapped_exception(error) if not isinstance(error, SQLSpecError) else error return StackExecutionError( - -1, "Oracle pipeline execution failed", error, adapter=type(self).__name__, mode=mode + -1, "Oracle pipeline execution failed", mapped_error, adapter=type(self).__name__, mode=mode ) @@ -376,13 +377,19 @@ class OracleSyncDriver(OraclePipelineMixin, SyncDriverAdapterBase): __slots__ = ( "_data_dictionary", + "_fetch_kwargs", "_oracle_version_cache", - "_pipeline_support", - "_pipeline_support_reason", "_row_metadata_cache", + "_session_tracing_enabled", "_transaction_active", ) dialect = "oracle" + _pipeline_support: ClassVar[bool] = False + _pipeline_support_reason: ClassVar[str] = "asyncio_thin_required" + + def _pipeline_native_supported(self) -> bool: + """Sync Oracle connections do not support native pipelining.""" + return False def __init__( self, @@ -397,11 +404,13 @@ def __init__( super().__init__(connection=connection, statement_config=statement_config, driver_features=driver_features) self._data_dictionary: OracledbSyncDataDictionary | None = None - self._pipeline_support: bool | None = None - self._pipeline_support_reason: str | None = None self._oracle_version_cache: OracleVersionCache | None = None self._row_metadata_cache: dict[int, tuple[Any, list[str], bool]] = {} self._transaction_active = False + self._fetch_kwargs = build_fetch_kwargs(self.driver_features) + self._session_tracing_enabled = bool( + self.driver_features.get("enable_session_tracing", False) or self.statement_config.enable_sqlcommenter + ) @property def data_dictionary(self) -> "OracledbSyncDataDictionary": @@ -449,7 +458,7 @@ def dispatch_execute(self, cursor: Any, statement: "SQL") -> "ExecutionResult": cursor, explained_sql, prepared_parameters, display_cursor=display_cursor, xplan_format=xplan_format ) - cursor.execute(sql, prepared_parameters or {}, **build_fetch_kwargs(self.driver_features)) + cursor.execute(sql, prepared_parameters or {}, **self._fetch_kwargs) is_select_like = statement.returns_rows() or self._should_force_select(statement, cursor) @@ -476,6 +485,47 @@ def dispatch_execute(self, cursor: Any, statement: "SQL") -> "ExecutionResult": affected_rows = resolve_rowcount(cursor) return self.create_execution_result(cursor, rowcount_override=affected_rows) + def _execute_cache_hit_with_cursor( + self, sql: str, params: "tuple[Any, ...] | list[Any] | dict[str, Any]", cached: "CachedQuery" + ) -> "SQLResult": + """Execute a cached query while owning only the cursor context.""" + with self.with_cursor(self.connection) as cursor: + wire_sql = self._format_wire_sql(cached.compiled_sql) + returns_rows = cached.operation_profile.returns_rows + fetch_kwargs = self._fetch_kwargs + parameters = coerce_large_parameters_sync( + self.connection, + params, + clob_type=DB_TYPE_CLOB, + blob_type=DB_TYPE_BLOB, + varchar2_byte_limit=self.driver_features.get("oracle_varchar2_byte_limit", 4000), + raw_byte_limit=self.driver_features.get("oracle_raw_byte_limit", 2000), + version_cache=self._oracle_version_cache, + ) + cast("Any", cursor).execute(wire_sql, parameters or {}, **fetch_kwargs) + if returns_rows and cursor.description is not None: + fetched_data = cursor.fetchall() + column_names, requires_lob_coercion = self._resolve_row_metadata(cursor.description) + data, column_names = collect_sync_rows( + cast("list[Any] | None", fetched_data), + cursor.description, + self.driver_features, + column_names=column_names, + requires_lob_coercion=requires_lob_coercion, + ) + return self.build_cached_select_result( + sql, + params, + cached, + selected_data=data, + column_names=column_names, + data_row_count=len(data), + row_format="tuple", + compiled_sql=wire_sql, + ) + affected_rows = resolve_rowcount(cursor) + return cast("SQLResult", DMLResult(cached.operation_type, affected_rows)) + def dispatch_execute_many(self, cursor: Any, statement: "SQL") -> "ExecutionResult": """Execute SQL with multiple parameter sets using Oracle batch processing. @@ -525,17 +575,7 @@ def dispatch_execute_many(self, cursor: Any, statement: "SQL") -> "ExecutionResu ) def dispatch_execute_script(self, cursor: Any, statement: "SQL") -> "ExecutionResult": - """Execute SQL script with statement splitting and parameter handling. - - Parameters are embedded as static values for script execution compatibility. - - Args: - cursor: Oracle cursor object - statement: SQL script statement to execute - - Returns: - Execution result containing statement count and success information - """ + """Execute SQL script with statement splitting and parameter handling.""" sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) prepared_parameters = coerce_large_parameters_sync( self.connection, @@ -546,18 +586,42 @@ def dispatch_execute_script(self, cursor: Any, statement: "SQL") -> "ExecutionRe raw_byte_limit=self.driver_features.get("oracle_raw_byte_limit", 2000), version_cache=self._oracle_version_cache, ) - prepared_parameters = cast("list[Any] | tuple[Any, ...] | dict[Any, Any] | None", prepared_parameters) statements = self.split_script_statements(sql, statement.statement_config, strip_trailing_semicolon=True) successful_count = 0 last_cursor = cursor - - for stmt in statements: - cursor.execute(stmt, _filter_script_parameters(stmt, prepared_parameters)) + pos_offset = 0 + total_statements = len(statements) + runtime = self._observability + has_observers = runtime is not None and runtime.has_statement_observers + + for idx, stmt in enumerate(statements): + stmt_params, pos_offset = _slice_script_parameters(stmt, prepared_parameters, pos_offset) + stmt_started = perf_counter() if has_observers else 0.0 + cursor.execute(stmt, stmt_params) successful_count += 1 + if has_observers: + stmt_duration = perf_counter() - stmt_started + single_stmt = self._sub_statement(stmt, stmt_params) + try: + stmt_operation = single_stmt.operation_type + finally: + self._release_pooled_statement(single_stmt) + self._emit_script_statement_event( + sql=stmt, + parameters=stmt_params, + operation=stmt_operation, + rows_affected=resolve_rowcount(cursor), + duration_s=stmt_duration, + started_at=stmt_started, + profile=None, + script_ordinal=idx, + script_statement_count=total_statements, + ) + self._pending_statement_profile = None return self.create_execution_result( - last_cursor, statement_count=len(statements), successful_statements=successful_count, is_script_result=True + last_cursor, statement_count=total_statements, successful_statements=successful_count, is_script_result=True ) def begin(self) -> None: @@ -579,7 +643,8 @@ def commit(self) -> None: self.connection.commit() except OracleError as e: raise create_mapped_exception(e) from e - self._transaction_active = False + finally: + self._transaction_active = False def rollback(self) -> None: """Rollback the current transaction. @@ -626,7 +691,7 @@ def with_cursor(self, connection: OracleSyncConnection) -> OracleSyncCursor: connection, arraysize=self.driver_features.get("arraysize"), prefetchrows=self.driver_features.get("prefetchrows"), - on_enter=self._apply_session_tracing if self._session_tracing_enabled() else None, + on_enter=self._apply_session_tracing if self._session_tracing_enabled else None, ) @overload @@ -747,23 +812,43 @@ def select_to_arrow( ) prepared_parameters = cast("list[Any] | tuple[Any, ...] | dict[Any, Any] | None", prepared_parameters) + exc_handler = self.handle_database_exceptions() try: self._apply_session_tracing(self.connection) - if return_format in {"batches", "reader"} and supports_df_batches(self.connection): - record_batches, batch_schema = self._fetch_arrow_record_batches( - sql, prepared_parameters, batch_size, arrow_schema + if return_format == "reader" and supports_df_batches(self.connection): + batch_generator = self._iter_arrow_record_batches(sql, prepared_parameters, batch_size, arrow_schema) + try: + first_batch = next(batch_generator) + except StopIteration: + pa = import_pyarrow() + empty_schema = arrow_schema if arrow_schema is not None else pa.schema([]) + reader = pa.RecordBatchReader.from_batches(empty_schema, ()) + return create_arrow_result(prepared_statement, reader, rows_affected=0) + pa = import_pyarrow() + schema = arrow_schema if arrow_schema is not None else first_batch.schema + lazy_reader = pa.RecordBatchReader.from_batches(schema, chain((first_batch,), batch_generator)) + return create_arrow_result(prepared_statement, lazy_reader, rows_affected=-1) + + if return_format == "batches" and supports_df_batches(self.connection): + batches_result = cast( + "tuple[list[ArrowRecordBatch], ArrowSchema | None]", + self._run_with_exception_handler( + exc_handler, + self._fetch_arrow_record_batches, + sql, + prepared_parameters, + batch_size, + arrow_schema, + ), ) + self._check_pending_exception(exc_handler) + record_batches, _ = batches_result rows_affected = sum(batch.num_rows for batch in record_batches) - if batch_schema is None: - batch_schema = pa.schema([]) - arrow_data: object - if return_format == "reader": - arrow_data = pa.RecordBatchReader.from_batches(batch_schema, record_batches) - else: - arrow_data = record_batches - return create_arrow_result(prepared_statement, arrow_data, rows_affected=rows_affected) + return create_arrow_result(prepared_statement, record_batches, rows_affected=rows_affected) - oracle_df = self._execute_arrow_dataframe(sql, prepared_parameters, batch_size) + oracle_df = self._run_with_exception_handler( + exc_handler, self._execute_arrow_dataframe, sql, prepared_parameters, batch_size + ) except AttributeError as exc: if native_only: msg = "Oracle native Arrow support is not available for this connection." @@ -776,8 +861,9 @@ def select_to_arrow( batch_size=batch_size, arrow_schema=arrow_schema, ) + self._check_pending_exception(exc_handler) - arrow_table = coerce_oracle_df_to_arrow_table(oracle_df, pa) + arrow_table = coerce_oracle_df_to_arrow_table(oracle_df, import_pyarrow()) column_names = normalize_column_names(arrow_table.column_names, self.driver_features) if column_names != arrow_table.column_names: arrow_table = arrow_table.rename_columns(column_names) @@ -791,17 +877,8 @@ def select_to_arrow( ) def execute_stack(self, stack: "StatementStack", *, continue_on_error: bool = False) -> "tuple[StackResult, ...]": - """Execute a StatementStack using Oracle's pipeline when available.""" - if not isinstance(stack, StatementStack) or not stack: - return super().execute_stack(stack, continue_on_error=continue_on_error) - - blocker = self._stack_native_blocker(stack) - if blocker is not None: - self._log_pipeline_skip(blocker, stack) - return super().execute_stack(stack, continue_on_error=continue_on_error) - - if not self._pipeline_native_supported(): - self._log_pipeline_skip(self._pipeline_support_reason or "database_version", stack) + """Execute a StatementStack sequentially on the sync driver.""" + self._log_pipeline_skip("asyncio_thin_required", stack) return super().execute_stack(stack, continue_on_error=continue_on_error) def select_to_storage( @@ -883,14 +960,16 @@ def load_from_arrow( if overwrite: statement = build_truncate_statement(table) exc_handler = self.handle_database_exceptions() - with exc_handler: + + def _truncate() -> None: cursor = self.connection.cursor() try: cursor.execute(statement) finally: cursor.close() - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + + self._run_with_exception_handler(exc_handler, _truncate) + self._check_pending_exception(exc_handler) use_direct_path = ( self.driver_features.get("enable_direct_path_load", True) is not False and supports_direct_path_load(self.connection) @@ -900,34 +979,37 @@ def load_from_arrow( if use_direct_path: schema_name, table_name = _resolve_direct_path_target(self.connection, table) exc_handler = self.handle_database_exceptions() - with exc_handler: - self.connection.direct_path_load( - schema_name=schema_name, - table_name=table_name, - column_names=list(arrow_table.column_names), - data=arrow_table, - ) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + self._run_with_exception_handler( + exc_handler, + self.connection.direct_path_load, + schema_name=schema_name, + table_name=table_name, + column_names=list(arrow_table.column_names), + data=arrow_table, + ) + self._check_pending_exception(exc_handler) else: columns = list(arrow_table.column_names) statement = build_insert_statement(table, columns) exc_handler = self.handle_database_exceptions() - with self.with_cursor(self.connection) as cursor, exc_handler: - for batch in arrow_table.to_batches(max_chunksize=5000): - batch_rows = [tuple(row[col] for col in columns) for row in batch.to_pylist()] - coerced_rows = coerce_many_parameters_sync( - self.connection, - batch_rows, - clob_type=DB_TYPE_CLOB, - blob_type=DB_TYPE_BLOB, - varchar2_byte_limit=self.driver_features.get("oracle_varchar2_byte_limit", 4000), - raw_byte_limit=self.driver_features.get("oracle_raw_byte_limit", 2000), - version_cache=self._oracle_version_cache, - ) - cursor.executemany(statement, coerced_rows) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + + def _insert_batches() -> None: + with self.with_cursor(self.connection) as cursor: + for batch in arrow_table.to_batches(max_chunksize=5000): + batch_rows = self._arrow_batch_to_rows(batch, columns) + coerced_rows = coerce_many_parameters_sync( + self.connection, + batch_rows, + clob_type=DB_TYPE_CLOB, + blob_type=DB_TYPE_BLOB, + varchar2_byte_limit=self.driver_features.get("oracle_varchar2_byte_limit", 4000), + raw_byte_limit=self.driver_features.get("oracle_raw_byte_limit", 2000), + version_cache=self._oracle_version_cache, + ) + cursor.executemany(statement, coerced_rows) + + self._run_with_exception_handler(exc_handler, _insert_batches) + self._check_pending_exception(exc_handler) telemetry_payload = self._ingest_telemetry(arrow_table) telemetry_payload["destination"] = table self._attach_partition_telemetry(telemetry_payload, partitioner) @@ -974,13 +1056,17 @@ def load_from_storage( truncate_sql = build_truncate_statement(table) if overwrite else None rows_processed = 0 exc_handler = self.handle_database_exceptions() - with self.with_cursor(self.connection) as cursor, exc_handler: - if truncate_sql is not None: - cursor.execute(truncate_sql) - cursor.execute(plsql, binds) - rows_processed = resolve_rowcount(cursor) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + + def _run_dbms_cloud() -> None: + nonlocal rows_processed + with self.with_cursor(self.connection) as cursor: + if truncate_sql is not None: + cursor.execute(truncate_sql) + cursor.execute(plsql, binds) + rows_processed = resolve_rowcount(cursor) + + self._run_with_exception_handler(exc_handler, _run_dbms_cloud) + self._check_pending_exception(exc_handler) telemetry_payload: StorageTelemetry = { "destination": table, "backend": "dbms_cloud", @@ -1058,10 +1144,10 @@ def _dispatch_explain_plan( Execution result carrying the formatted plan rows. """ if display_cursor: - cursor.execute(explained_sql, parameters or {}, **build_fetch_kwargs(self.driver_features)) + cursor.execute(explained_sql, parameters or {}, **self._fetch_kwargs) if cursor.description is not None: cursor.fetchall() - cursor.execute(_format_display_cursor_sql(xplan_format), {}, **build_fetch_kwargs(self.driver_features)) + cursor.execute(_format_display_cursor_sql(xplan_format), {}, **self._fetch_kwargs) fetched_data = cursor.fetchall() column_names, requires_lob_coercion = self._resolve_row_metadata(cursor.description) data, column_names = collect_sync_rows( @@ -1085,7 +1171,7 @@ def _dispatch_explain_plan( display_sql = PLAN_DISPLAY_SQL if xplan_format == "TYPICAL" else _format_plan_display_sql(xplan_format) try: - cursor.execute(display_sql, {"statement_id": statement_id}, **build_fetch_kwargs(self.driver_features)) + cursor.execute(display_sql, {"statement_id": statement_id}, **self._fetch_kwargs) fetched_data = cursor.fetchall() column_names, requires_lob_coercion = self._resolve_row_metadata(cursor.description) data, column_names = collect_sync_rows( @@ -1138,6 +1224,40 @@ def _execute_arrow_dataframe(self, sql: str, parameters: "Any", batch_size: int statement=sql, parameters=params, arraysize=batch_size or 1000, **retry_kwargs ) + def _iter_arrow_record_batches( + self, + sql: str, + parameters: "StatementParameters | None", + batch_size: int | None, + arrow_schema: "ArrowSchema | None", + ) -> "Iterator[ArrowRecordBatch]": + params = parameters if parameters is not None else [] + fetch_kwargs = build_arrow_fetch_kwargs(self.driver_features, arrow_schema=arrow_schema) + try: + oracle_dfs = self.connection.fetch_df_batches( + statement=sql, parameters=params, size=batch_size or 1000, **fetch_kwargs + ) + except TypeError as exc: + retry_kwargs = _retry_arrow_without_fetch_lobs(exc, fetch_kwargs) + if retry_kwargs is None: + raise + oracle_dfs = self.connection.fetch_df_batches( + statement=sql, parameters=params, size=batch_size or 1000, **retry_kwargs + ) + except OracleError as exc: + raise create_mapped_exception(exc) from exc + try: + for oracle_df in oracle_dfs: + batch_table = coerce_oracle_df_to_arrow_table(oracle_df, import_pyarrow()) + column_names = normalize_column_names(batch_table.column_names, self.driver_features) + if column_names != batch_table.column_names: + batch_table = batch_table.rename_columns(column_names) + if arrow_schema is not None: + batch_table = batch_table.cast(arrow_schema) + yield from batch_table.to_batches() + except OracleError as exc: + raise create_mapped_exception(exc) from exc + def _fetch_arrow_record_batches( self, sql: str, @@ -1148,7 +1268,7 @@ def _fetch_arrow_record_batches( params = parameters if parameters is not None else [] record_batches: list[ArrowRecordBatch] = [] batch_schema: ArrowSchema | None = None - fetch_kwargs = build_arrow_fetch_kwargs(self.driver_features) + fetch_kwargs = build_arrow_fetch_kwargs(self.driver_features, arrow_schema=arrow_schema) try: oracle_dfs = self.connection.fetch_df_batches( statement=sql, parameters=params, size=batch_size or 1000, **fetch_kwargs @@ -1161,7 +1281,7 @@ def _fetch_arrow_record_batches( statement=sql, parameters=params, size=batch_size or 1000, **retry_kwargs ) for oracle_df in oracle_dfs: - batch_table = coerce_oracle_df_to_arrow_table(oracle_df, pa) + batch_table = coerce_oracle_df_to_arrow_table(oracle_df, import_pyarrow()) column_names = normalize_column_names(batch_table.column_names, self.driver_features) if column_names != batch_table.column_names: batch_table = batch_table.rename_columns(column_names) @@ -1172,19 +1292,6 @@ def _fetch_arrow_record_batches( record_batches.extend(batch_table.to_batches()) return record_batches, batch_schema - def _pipeline_native_supported(self) -> bool: - if self._pipeline_support is not None: - return self._pipeline_support - - if self.stack_native_disabled: - self._pipeline_support = False - self._pipeline_support_reason = "env_override" - return False - - self._pipeline_support = False - self._pipeline_support_reason = "asyncio_thin_required" - return False - class OracleAsyncDriver(OraclePipelineMixin, AsyncDriverAdapterBase): """Asynchronous Oracle Database driver. @@ -1195,10 +1302,12 @@ class OracleAsyncDriver(OraclePipelineMixin, AsyncDriverAdapterBase): __slots__ = ( "_data_dictionary", + "_fetch_kwargs", "_oracle_version_cache", "_pipeline_support", "_pipeline_support_reason", "_row_metadata_cache", + "_session_tracing_enabled", "_transaction_active", ) dialect = "oracle" @@ -1221,6 +1330,10 @@ def __init__( self._oracle_version_cache: OracleVersionCache | None = None self._row_metadata_cache: dict[int, tuple[Any, list[str], bool]] = {} self._transaction_active = False + self._fetch_kwargs = build_fetch_kwargs(self.driver_features) + self._session_tracing_enabled = bool( + self.driver_features.get("enable_session_tracing", False) or self.statement_config.enable_sqlcommenter + ) @property def data_dictionary(self) -> "OracledbAsyncDataDictionary": @@ -1268,7 +1381,7 @@ async def dispatch_execute(self, cursor: Any, statement: "SQL") -> "ExecutionRes cursor, explained_sql, prepared_parameters, display_cursor=display_cursor, xplan_format=xplan_format ) - await cursor.execute(sql, prepared_parameters or {}, **build_fetch_kwargs(self.driver_features)) + await cursor.execute(sql, prepared_parameters or {}, **self._fetch_kwargs) is_select_like = statement.returns_rows() or self._should_force_select(statement, cursor) @@ -1295,6 +1408,47 @@ async def dispatch_execute(self, cursor: Any, statement: "SQL") -> "ExecutionRes affected_rows = resolve_rowcount(cursor) return self.create_execution_result(cursor, rowcount_override=affected_rows) + async def _execute_cache_hit_with_cursor( + self, sql: str, params: "tuple[Any, ...] | list[Any] | dict[str, Any]", cached: "CachedQuery" + ) -> "SQLResult": + """Execute an async cached query while owning only the cursor context.""" + async with self.with_cursor(self.connection) as cursor: + wire_sql = self._format_wire_sql(cached.compiled_sql) + returns_rows = cached.operation_profile.returns_rows + fetch_kwargs = self._fetch_kwargs + parameters = await coerce_large_parameters_async( + self.connection, + params, + clob_type=DB_TYPE_CLOB, + blob_type=DB_TYPE_BLOB, + varchar2_byte_limit=self.driver_features.get("oracle_varchar2_byte_limit", 4000), + raw_byte_limit=self.driver_features.get("oracle_raw_byte_limit", 2000), + version_cache=self._oracle_version_cache, + ) + await cast("Any", cursor).execute(wire_sql, parameters or {}, **fetch_kwargs) + if returns_rows and cursor.description is not None: + fetched_data = await cursor.fetchall() + column_names, requires_lob_coercion = self._resolve_row_metadata(cursor.description) + data, column_names = await collect_async_rows( + cast("list[Any] | None", fetched_data), + cursor.description, + self.driver_features, + column_names=column_names, + requires_lob_coercion=requires_lob_coercion, + ) + return self.build_cached_select_result( + sql, + params, + cached, + selected_data=data, + column_names=column_names, + data_row_count=len(data), + row_format="tuple", + compiled_sql=wire_sql, + ) + affected_rows = resolve_rowcount(cursor) + return cast("SQLResult", DMLResult(cached.operation_type, affected_rows)) + async def dispatch_execute_many(self, cursor: Any, statement: "SQL") -> "ExecutionResult": """Execute SQL with multiple parameter sets using Oracle batch processing. @@ -1346,17 +1500,7 @@ async def dispatch_execute_many(self, cursor: Any, statement: "SQL") -> "Executi ) async def dispatch_execute_script(self, cursor: Any, statement: "SQL") -> "ExecutionResult": - """Execute SQL script with statement splitting and parameter handling. - - Parameters are embedded as static values for script execution compatibility. - - Args: - cursor: Oracle cursor object - statement: SQL script statement to execute - - Returns: - Execution result containing statement count and success information - """ + """Execute async SQL script with statement splitting and parameter handling.""" sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) prepared_parameters = await coerce_large_parameters_async( self.connection, @@ -1367,18 +1511,42 @@ async def dispatch_execute_script(self, cursor: Any, statement: "SQL") -> "Execu raw_byte_limit=self.driver_features.get("oracle_raw_byte_limit", 2000), version_cache=self._oracle_version_cache, ) - prepared_parameters = cast("list[Any] | tuple[Any, ...] | dict[Any, Any] | None", prepared_parameters) statements = self.split_script_statements(sql, statement.statement_config, strip_trailing_semicolon=True) successful_count = 0 last_cursor = cursor - - for stmt in statements: - await cursor.execute(stmt, _filter_script_parameters(stmt, prepared_parameters)) + pos_offset = 0 + total_statements = len(statements) + runtime = self._observability + has_observers = runtime is not None and runtime.has_statement_observers + + for idx, stmt in enumerate(statements): + stmt_params, pos_offset = _slice_script_parameters(stmt, prepared_parameters, pos_offset) + stmt_started = perf_counter() if has_observers else 0.0 + await cursor.execute(stmt, stmt_params) successful_count += 1 + if has_observers: + stmt_duration = perf_counter() - stmt_started + single_stmt = self._sub_statement(stmt, stmt_params) + try: + stmt_operation = single_stmt.operation_type + finally: + self._release_pooled_statement(single_stmt) + self._emit_script_statement_event( + sql=stmt, + parameters=stmt_params, + operation=stmt_operation, + rows_affected=resolve_rowcount(cursor), + duration_s=stmt_duration, + started_at=stmt_started, + profile=None, + script_ordinal=idx, + script_statement_count=total_statements, + ) + self._pending_statement_profile = None return self.create_execution_result( - last_cursor, statement_count=len(statements), successful_statements=successful_count, is_script_result=True + last_cursor, statement_count=total_statements, successful_statements=successful_count, is_script_result=True ) async def begin(self) -> None: @@ -1400,7 +1568,8 @@ async def commit(self) -> None: await self.connection.commit() except OracleError as e: raise create_mapped_exception(e) from e - self._transaction_active = False + finally: + self._transaction_active = False async def rollback(self) -> None: """Rollback the current transaction. @@ -1450,7 +1619,7 @@ def with_cursor(self, connection: OracleAsyncConnection) -> OracleAsyncCursor: connection, arraysize=self.driver_features.get("arraysize"), prefetchrows=self.driver_features.get("prefetchrows"), - on_enter=self._apply_session_tracing if self._session_tracing_enabled() else None, + on_enter=self._apply_session_tracing if self._session_tracing_enabled else None, ) @overload @@ -1571,13 +1740,25 @@ async def select_to_arrow( ) prepared_parameters = cast("list[Any] | tuple[Any, ...] | dict[Any, Any] | None", prepared_parameters) + exc_handler = self.handle_database_exceptions() try: self._apply_session_tracing(self.connection) if return_format in {"batches", "reader"} and supports_df_batches(self.connection): - record_batches, batch_schema = await self._fetch_arrow_record_batches( - sql, prepared_parameters, batch_size, arrow_schema + batches_result = cast( + "tuple[list[ArrowRecordBatch], ArrowSchema | None]", + await self._run_with_exception_handler( + exc_handler, + self._fetch_arrow_record_batches, + sql, + prepared_parameters, + batch_size, + arrow_schema, + ), ) + self._check_pending_exception(exc_handler) + record_batches, batch_schema = batches_result rows_affected = sum(batch.num_rows for batch in record_batches) + pa = import_pyarrow() if batch_schema is None: batch_schema = pa.schema([]) arrow_data: object @@ -1587,7 +1768,9 @@ async def select_to_arrow( arrow_data = record_batches return create_arrow_result(prepared_statement, arrow_data, rows_affected=rows_affected) - oracle_df = await self._execute_arrow_dataframe(sql, prepared_parameters, batch_size) + oracle_df = await self._run_with_exception_handler( + exc_handler, self._execute_arrow_dataframe, sql, prepared_parameters, batch_size + ) except AttributeError as exc: if native_only: msg = "Oracle native Arrow support is not available for this connection." @@ -1600,8 +1783,9 @@ async def select_to_arrow( batch_size=batch_size, arrow_schema=arrow_schema, ) + self._check_pending_exception(exc_handler) - arrow_table = coerce_oracle_df_to_arrow_table(oracle_df, pa) + arrow_table = coerce_oracle_df_to_arrow_table(oracle_df, import_pyarrow()) column_names = normalize_column_names(arrow_table.column_names, self.driver_features) if column_names != arrow_table.column_names: arrow_table = arrow_table.rename_columns(column_names) @@ -1711,10 +1895,8 @@ async def load_from_arrow( if overwrite: statement = build_truncate_statement(table) exc_handler = self.handle_database_exceptions() - async with exc_handler: - await self.connection.execute(statement) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + await self._run_with_exception_handler(exc_handler, self.connection.execute, statement) + self._check_pending_exception(exc_handler) use_direct_path = ( self.driver_features.get("enable_direct_path_load", True) is not False and supports_direct_path_load(self.connection) @@ -1724,34 +1906,37 @@ async def load_from_arrow( if use_direct_path: schema_name, table_name = _resolve_direct_path_target(self.connection, table) exc_handler = self.handle_database_exceptions() - async with exc_handler: - await self.connection.direct_path_load( - schema_name=schema_name, - table_name=table_name, - column_names=list(arrow_table.column_names), - data=arrow_table, - ) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + await self._run_with_exception_handler( + exc_handler, + self.connection.direct_path_load, + schema_name=schema_name, + table_name=table_name, + column_names=list(arrow_table.column_names), + data=arrow_table, + ) + self._check_pending_exception(exc_handler) else: columns = list(arrow_table.column_names) statement = build_insert_statement(table, columns) exc_handler = self.handle_database_exceptions() - async with self.with_cursor(self.connection) as cursor, exc_handler: - for batch in arrow_table.to_batches(max_chunksize=5000): - batch_rows = [tuple(row[col] for col in columns) for row in batch.to_pylist()] - coerced_rows = await coerce_many_parameters_async( - self.connection, - batch_rows, - clob_type=DB_TYPE_CLOB, - blob_type=DB_TYPE_BLOB, - varchar2_byte_limit=self.driver_features.get("oracle_varchar2_byte_limit", 4000), - raw_byte_limit=self.driver_features.get("oracle_raw_byte_limit", 2000), - version_cache=self._oracle_version_cache, - ) - await cursor.executemany(statement, coerced_rows) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + + async def _insert_batches() -> None: + async with self.with_cursor(self.connection) as cursor: + for batch in arrow_table.to_batches(max_chunksize=5000): + batch_rows = self._arrow_batch_to_rows(batch, columns) + coerced_rows = await coerce_many_parameters_async( + self.connection, + batch_rows, + clob_type=DB_TYPE_CLOB, + blob_type=DB_TYPE_BLOB, + varchar2_byte_limit=self.driver_features.get("oracle_varchar2_byte_limit", 4000), + raw_byte_limit=self.driver_features.get("oracle_raw_byte_limit", 2000), + version_cache=self._oracle_version_cache, + ) + await cursor.executemany(statement, coerced_rows) + + await self._run_with_exception_handler(exc_handler, _insert_batches) + self._check_pending_exception(exc_handler) telemetry_payload = self._ingest_telemetry(arrow_table) telemetry_payload["destination"] = table self._attach_partition_telemetry(telemetry_payload, partitioner) @@ -1798,13 +1983,17 @@ async def load_from_storage( truncate_sql = build_truncate_statement(table) if overwrite else None rows_processed = 0 exc_handler = self.handle_database_exceptions() - async with self.with_cursor(self.connection) as cursor, exc_handler: - if truncate_sql is not None: - await cursor.execute(truncate_sql) - await cursor.execute(plsql, binds) - rows_processed = resolve_rowcount(cursor) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + + async def _run_dbms_cloud_async() -> None: + nonlocal rows_processed + async with self.with_cursor(self.connection) as cursor: + if truncate_sql is not None: + await cursor.execute(truncate_sql) + await cursor.execute(plsql, binds) + rows_processed = resolve_rowcount(cursor) + + await self._run_with_exception_handler(exc_handler, _run_dbms_cloud_async) + self._check_pending_exception(exc_handler) telemetry_payload: StorageTelemetry = { "destination": table, "backend": "dbms_cloud", @@ -1891,12 +2080,10 @@ async def _dispatch_explain_plan( Execution result carrying the formatted plan rows. """ if display_cursor: - await cursor.execute(explained_sql, parameters or {}, **build_fetch_kwargs(self.driver_features)) + await cursor.execute(explained_sql, parameters or {}, **self._fetch_kwargs) if cursor.description is not None: await cursor.fetchall() - await cursor.execute( - _format_display_cursor_sql(xplan_format), {}, **build_fetch_kwargs(self.driver_features) - ) + await cursor.execute(_format_display_cursor_sql(xplan_format), {}, **self._fetch_kwargs) fetched_data = await cursor.fetchall() column_names, requires_lob_coercion = self._resolve_row_metadata(cursor.description) data, column_names = await collect_async_rows( @@ -1920,9 +2107,7 @@ async def _dispatch_explain_plan( display_sql = PLAN_DISPLAY_SQL if xplan_format == "TYPICAL" else _format_plan_display_sql(xplan_format) try: - await cursor.execute( - display_sql, {"statement_id": statement_id}, **build_fetch_kwargs(self.driver_features) - ) + await cursor.execute(display_sql, {"statement_id": statement_id}, **self._fetch_kwargs) fetched_data = await cursor.fetchall() column_names, requires_lob_coercion = self._resolve_row_metadata(cursor.description) data, column_names = await collect_async_rows( @@ -1985,7 +2170,7 @@ async def _fetch_arrow_record_batches( params = parameters if parameters is not None else [] record_batches: list[ArrowRecordBatch] = [] batch_schema: ArrowSchema | None = None - fetch_kwargs = build_arrow_fetch_kwargs(self.driver_features) + fetch_kwargs = build_arrow_fetch_kwargs(self.driver_features, arrow_schema=arrow_schema) try: oracle_dfs = self.connection.fetch_df_batches( statement=sql, parameters=params, size=batch_size or 1000, **fetch_kwargs @@ -1998,7 +2183,7 @@ async def _fetch_arrow_record_batches( statement=sql, parameters=params, size=batch_size or 1000, **retry_kwargs ) async for oracle_df in oracle_dfs: - batch_table = coerce_oracle_df_to_arrow_table(oracle_df, pa) + batch_table = coerce_oracle_df_to_arrow_table(oracle_df, import_pyarrow()) column_names = normalize_column_names(batch_table.column_names, self.driver_features) if column_names != batch_table.column_names: batch_table = batch_table.rename_columns(column_names) @@ -2063,7 +2248,7 @@ async def _execute_stack_native( if owns_transaction: try: await self.rollback() - except Exception as rollback_error: # pragma: no cover + except Exception as rollback_error: logger.debug("Rollback after pipeline failure failed: %s", rollback_error) if isinstance(exc, StackExecutionError): raise @@ -2142,6 +2327,7 @@ def _arrow_schema_supports_direct_path(arrow_table: Any) -> bool: True when every column type is known to be convertible. """ ensure_pyarrow() + pa = import_pyarrow() for field in arrow_table.schema: data_type = field.type @@ -2267,28 +2453,75 @@ def _retry_arrow_without_fetch_lobs(exc: TypeError, fetch_kwargs: "dict[str, obj return retry_kwargs -def _statement_contains_bind(statement: str, name: str) -> bool: - """Return whether a SQL statement references a named Oracle bind placeholder.""" - lowered_sql = statement.lower() - target = f":{name.lower()}" - start = 0 - while True: - idx = lowered_sql.find(target, start) - if idx == -1: - return False - prev_ok = idx == 0 or lowered_sql[idx - 1] != ":" - end_idx = idx + len(target) - next_ch = lowered_sql[end_idx] if end_idx < len(lowered_sql) else "" - next_ok = not (next_ch.isalnum() or next_ch in {"_", "$", "#"}) - if prev_ok and next_ok: - return True - start = idx + 1 +def _extract_statement_binds(statement: str) -> "tuple[set[str], int]": + """Extract named bind names (lowercased) and count of positional binds outside quotes and comments.""" + named_binds: set[str] = set() + positional_count = 0 + n = len(statement) + i = 0 + while i < n: + ch = statement[i] + if ch == "-" and i + 1 < n and statement[i + 1] == "-": + end = statement.find("\n", i + 2) + i = n if end == -1 else end + 1 + continue + if ch == "/" and i + 1 < n and statement[i + 1] == "*": + end = statement.find("*/", i + 2) + i = n if end == -1 else end + 2 + continue + if ch == "'": + end = i + 1 + while end < n: + if statement[end] == "'": + if end + 1 < n and statement[end + 1] == "'": + end += 2 + continue + end += 1 + break + end += 1 + i = end + continue + if ch == '"': + end = i + 1 + while end < n: + if statement[end] == '"': + if end + 1 < n and statement[end + 1] == '"': + end += 2 + continue + end += 1 + break + end += 1 + i = end + continue + if ch == ":" and (i == 0 or statement[i - 1] != ":"): + start = i + 1 + if start < n and statement[start] != ":": + end = start + while end < n and (statement[end].isalnum() or statement[end] in {"_", "$", "#"}): + end += 1 + if end > start: + token = statement[start:end] + if token.isdigit(): + positional_count += 1 + else: + named_binds.add(token.lower()) + i = end + continue + i += 1 + return named_binds, positional_count -def _filter_script_parameters(statement: str, parameters: Any) -> Any: - """Filter dictionary script parameters to the bind placeholders present in statement.""" +def _slice_script_parameters(statement: str, parameters: Any, pos_offset: int) -> "tuple[Any, int]": + """Filter dictionary binds or slice positional binds for a single statement.""" if not parameters: - return {} + return {}, pos_offset + named_binds, pos_count = _extract_statement_binds(statement) if isinstance(parameters, dict): - return {key: val for key, val in parameters.items() if _statement_contains_bind(statement, str(key))} - return parameters + filtered = {k: v for k, v in parameters.items() if str(k).lower() in named_binds} + return filtered, pos_offset + if isinstance(parameters, (list, tuple)): + if pos_count > 0: + sliced = list(parameters[pos_offset : pos_offset + pos_count]) + return sliced, pos_offset + pos_count + return [], pos_offset + return parameters, pos_offset diff --git a/sqlspec/adapters/oracledb/events/_hub.py b/sqlspec/adapters/oracledb/events/_hub.py index 50f748568..14226f667 100644 --- a/sqlspec/adapters/oracledb/events/_hub.py +++ b/sqlspec/adapters/oracledb/events/_hub.py @@ -6,8 +6,6 @@ caller's poll_interval (capped at the configured aq_wait_seconds ceiling). """ -# pyright: reportPrivateUsage=false - import asyncio import contextlib import logging @@ -84,7 +82,7 @@ def dequeue(self, channel: str, poll_interval: float) -> "Any | None": _apply_deq_options(queue, self._visibility, self._default_visibility, wait_seconds) try: message = queue.deqone() - except Exception as error: # pragma: no cover + except Exception as error: if not isinstance(error, OracleDatabaseError): raise log_with_context( @@ -218,7 +216,7 @@ async def dequeue(self, channel: str, poll_interval: float) -> "Any | None": _apply_deq_options(queue, self._visibility, self._default_visibility, wait_seconds) try: message = await queue.deqone() - except Exception as error: # pragma: no cover + except Exception as error: if not isinstance(error, OracleDatabaseError): raise log_with_context( diff --git a/sqlspec/adapters/oracledb/events/backend.py b/sqlspec/adapters/oracledb/events/backend.py index fc2501ec3..c8bb32077 100644 --- a/sqlspec/adapters/oracledb/events/backend.py +++ b/sqlspec/adapters/oracledb/events/backend.py @@ -3,13 +3,14 @@ import contextlib import logging from datetime import UTC, datetime -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, cast from sqlspec.adapters.oracledb._typing import DEQ_IMMEDIATE, DEQ_ON_COMMIT from sqlspec.adapters.oracledb.events._hub import OracleAsyncAQHub, OracleSyncAQHub from sqlspec.exceptions import ImproperConfigurationError, MissingDependencyError from sqlspec.extensions.events import EventMessage, parse_event_timestamp from sqlspec.utils.logging import get_logger, log_with_context +from sqlspec.utils.serializers import from_json from sqlspec.utils.uuids import uuid4 if TYPE_CHECKING: @@ -212,22 +213,22 @@ def create_event_backend( match (backend_name, is_async): case ("aq", False): try: - return OracleSyncAQEventBackend(config, extension_settings) # type: ignore[arg-type] + return OracleSyncAQEventBackend(cast("OracleSyncConfig", config), extension_settings) except (ImproperConfigurationError, MissingDependencyError): return None case ("aq", True): try: - return OracleAsyncAQEventBackend(config, extension_settings) # type: ignore[arg-type] + return OracleAsyncAQEventBackend(cast("OracleAsyncConfig", config), extension_settings) except (ImproperConfigurationError, MissingDependencyError): return None case ("txeventq", False): try: - return OracleSyncTxEventQEventBackend(config, extension_settings) # type: ignore[arg-type] + return OracleSyncTxEventQEventBackend(cast("OracleSyncConfig", config), extension_settings) except (ImproperConfigurationError, MissingDependencyError): return None case ("txeventq", True): try: - return OracleAsyncTxEventQEventBackend(config, extension_settings) # type: ignore[arg-type] + return OracleAsyncTxEventQEventBackend(cast("OracleAsyncConfig", config), extension_settings) except (ImproperConfigurationError, MissingDependencyError): return None case _: @@ -272,6 +273,11 @@ def _build_envelope( def _parse_message(channel: str, payload: Any) -> EventMessage: """Parse Oracle AQ message payload into EventMessage.""" + if isinstance(payload, (str, bytes)): + with contextlib.suppress(Exception): + parsed = from_json(payload) + if isinstance(parsed, dict): + payload = parsed if not isinstance(payload, dict): payload = {"payload": payload} payload_channel = payload.get("channel") diff --git a/sqlspec/adapters/oracledb/events/store.py b/sqlspec/adapters/oracledb/events/store.py index 36e5f4156..48fc64dbb 100644 --- a/sqlspec/adapters/oracledb/events/store.py +++ b/sqlspec/adapters/oracledb/events/store.py @@ -22,7 +22,7 @@ from typing_extensions import TypedDict -from sqlspec.adapters.oracledb._storage import oracle_table_feature_report +from sqlspec.adapters.oracledb._storage import oracle_table_feature_report, validate_oracle_identifier from sqlspec.adapters.oracledb.data_dictionary import JSONStorageType, storage_type_from_version from sqlspec.config import EventsConfig from sqlspec.extensions.events import BaseEventQueueStore @@ -133,6 +133,8 @@ class OracleSyncEventQueueStore(BaseEventQueueStore["OracleSyncConfig"]): def __init__(self, config: "OracleSyncConfig") -> None: """Initialize Oracle sync event queue store.""" super().__init__(config) + for part in self.table_name.split("."): + validate_oracle_identifier(part, "table name") self._in_memory, self._json_storage_override = _init_oracle_settings(self._extension_settings) def create_statements(self) -> "list[str]": @@ -218,7 +220,7 @@ def _table_feature_clause(self) -> str: range_partition_key="available_at", table_options_key="table_options", ) - return report["clause"] + return report["plsql_clause"] class OracleAsyncEventQueueStore(BaseEventQueueStore["OracleAsyncConfig"]): @@ -248,6 +250,8 @@ class OracleAsyncEventQueueStore(BaseEventQueueStore["OracleAsyncConfig"]): def __init__(self, config: "OracleAsyncConfig") -> None: """Initialize Oracle async event queue store.""" super().__init__(config) + for part in self.table_name.split("."): + validate_oracle_identifier(part, "table name") self._in_memory, self._json_storage_override = _init_oracle_settings(self._extension_settings) def create_statements(self) -> "list[str]": @@ -333,7 +337,7 @@ def _table_feature_clause(self) -> str: range_partition_key="available_at", table_options_key="table_options", ) - return report["clause"] + return report["plsql_clause"] def _init_oracle_settings(extension_settings: "dict[str, Any]") -> "tuple[bool, JSONStorageType | None]": diff --git a/sqlspec/adapters/oracledb/litestar/store.py b/sqlspec/adapters/oracledb/litestar/store.py index 080146d6a..ecd14e4da 100644 --- a/sqlspec/adapters/oracledb/litestar/store.py +++ b/sqlspec/adapters/oracledb/litestar/store.py @@ -1,14 +1,14 @@ """Oracle session store for Litestar integration.""" +import contextlib from datetime import UTC, timedelta from typing import TYPE_CHECKING, Any, NotRequired, TypedDict, cast -from sqlspec.adapters.oracledb._storage import oracle_table_feature_report -from sqlspec.adapters.oracledb.core import DB_TYPE_BLOB +from sqlspec.adapters.oracledb._storage import oracle_table_feature_report, validate_oracle_identifier +from sqlspec.adapters.oracledb.core import DB_TYPE_BLOB, read_oracle_lob_async, read_oracle_lob_sync, resolve_rowcount from sqlspec.config import LitestarConfig from sqlspec.extensions.litestar.store import BaseSQLSpecStore from sqlspec.utils.sync_tools import async_ -from sqlspec.utils.type_guards import is_async_readable, is_readable if TYPE_CHECKING: from sqlspec.adapters.oracledb.config import OracleAsyncConfig, OracleSyncConfig @@ -109,6 +109,8 @@ def __init__(self, config: "OracleAsyncConfig") -> None: config: OracleAsyncConfig instance. """ super().__init__(config) + for part in self._table_name.split("."): + validate_oracle_identifier(part, "table name") litestar_config = config.extension_config.get("litestar", {}) self._in_memory = bool(litestar_config.get("in_memory", False)) @@ -137,14 +139,14 @@ async def get(self, key: str, renew_for: "int | timedelta | None" = None) -> "by """ sql = f""" SELECT data, expires_at FROM {self._table_name} - WHERE session_id = :session_id + WHERE session_id = :1 AND (expires_at IS NULL OR expires_at > SYSTIMESTAMP) """ conn_context = self._config.provide_connection() async with conn_context as conn: with conn.cursor() as cursor: - await cursor.execute(sql, {"session_id": key}) + await cursor.execute(sql, (key,)) row = await cursor.fetchone() if row is None: @@ -158,14 +160,19 @@ async def get(self, key: str, renew_for: "int | timedelta | None" = None) -> "by update_sql = f""" UPDATE {self._table_name} SET expires_at = CASE - WHEN :expires_in_seconds IS NULL THEN NULL - ELSE SYSTIMESTAMP + NUMTODSINTERVAL(:expires_in_seconds, 'SECOND') + WHEN :1 IS NULL THEN NULL + ELSE SYSTIMESTAMP + NUMTODSINTERVAL(:1, 'SECOND') END, updated_at = SYSTIMESTAMP - WHERE session_id = :session_id + WHERE session_id = :2 """ - await cursor.execute(update_sql, {"expires_in_seconds": expires_in_seconds, "session_id": key}) - await conn.commit() + try: + await cursor.execute(update_sql, (expires_in_seconds, expires_in_seconds, key)) + await conn.commit() + except BaseException: + with contextlib.suppress(Exception): + await conn.rollback() + raise return await _read_blob_async(data_blob) @@ -208,12 +215,19 @@ async def set(self, key: str, value: "str | bytes", expires_in: "int | timedelta conn_context = self._config.provide_connection() async with conn_context as conn: - bind_data: Any = await conn.createlob(DB_TYPE_BLOB, data) if len(data) > ORACLE_SMALL_BLOB_LIMIT else data - with conn.cursor() as cursor: - await cursor.execute( - sql, {"session_id": key, "data": bind_data, "expires_in_seconds": expires_in_seconds} + try: + bind_data: Any = ( + await conn.createlob(DB_TYPE_BLOB, data) if len(data) > ORACLE_SMALL_BLOB_LIMIT else data ) - await conn.commit() + with conn.cursor() as cursor: + await cursor.execute( + sql, {"session_id": key, "data": bind_data, "expires_in_seconds": expires_in_seconds} + ) + await conn.commit() + except BaseException: + with contextlib.suppress(Exception): + await conn.rollback() + raise async def delete(self, key: str) -> None: """Delete a session by key. @@ -221,13 +235,18 @@ async def delete(self, key: str) -> None: Args: key: Session ID to delete. """ - sql = f"DELETE FROM {self._table_name} WHERE session_id = :session_id" + sql = f"DELETE FROM {self._table_name} WHERE session_id = :1" conn_context = self._config.provide_connection() async with conn_context as conn: - with conn.cursor() as cursor: - await cursor.execute(sql, {"session_id": key}) - await conn.commit() + try: + with conn.cursor() as cursor: + await cursor.execute(sql, (key,)) + await conn.commit() + except BaseException: + with contextlib.suppress(Exception): + await conn.rollback() + raise async def delete_all(self) -> None: """Delete all sessions from the store.""" @@ -235,9 +254,14 @@ async def delete_all(self) -> None: conn_context = self._config.provide_connection() async with conn_context as conn: - with conn.cursor() as cursor: - await cursor.execute(sql) - await conn.commit() + try: + with conn.cursor() as cursor: + await cursor.execute(sql) + await conn.commit() + except BaseException: + with contextlib.suppress(Exception): + await conn.rollback() + raise self._log_delete_all() async def exists(self, key: str) -> bool: @@ -251,14 +275,14 @@ async def exists(self, key: str) -> bool: """ sql = f""" SELECT 1 FROM {self._table_name} - WHERE session_id = :session_id + WHERE session_id = :1 AND (expires_at IS NULL OR expires_at > SYSTIMESTAMP) """ conn_context = self._config.provide_connection() async with conn_context as conn: with conn.cursor() as cursor: - await cursor.execute(sql, {"session_id": key}) + await cursor.execute(sql, (key,)) result = await cursor.fetchone() return result is not None @@ -273,13 +297,13 @@ async def expires_in(self, key: str) -> "int | None": """ sql = f""" SELECT expires_at, SYSTIMESTAMP FROM {self._table_name} - WHERE session_id = :session_id + WHERE session_id = :1 """ conn_context = self._config.provide_connection() async with conn_context as conn: with conn.cursor() as cursor: - await cursor.execute(sql, {"session_id": key}) + await cursor.execute(sql, (key,)) row = await cursor.fetchone() if row is None or row[0] is None: @@ -308,78 +332,26 @@ async def delete_expired(self) -> int: conn_context = self._config.provide_connection() async with conn_context as conn: - with conn.cursor() as cursor: - await cursor.execute(sql) - count = cursor.rowcount if cursor.rowcount is not None else 0 - await conn.commit() - if count > 0: - self._log_delete_expired(count) - return count + try: + with conn.cursor() as cursor: + await cursor.execute(sql) + count = resolve_rowcount(cursor) + await conn.commit() + if count > 0: + self._log_delete_expired(count) + return count + except BaseException: + with contextlib.suppress(Exception): + await conn.rollback() + raise def _table_ddl(self) -> str: - """Get Oracle CREATE TABLE SQL with optimized schema. - - Returns: - SQL statement to create the sessions table with proper indexes. - """ - table_clause = _litestar_table_feature_clause(self._config, self._in_memory) - index_name = f"idx_{self._table_name.replace('.', '_')}_expires_at"[:30] - return f""" - BEGIN - EXECUTE IMMEDIATE 'CREATE TABLE {self._table_name} ( - session_id VARCHAR2(255) PRIMARY KEY, - data BLOB NOT NULL, - expires_at TIMESTAMP WITH TIME ZONE, - created_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL, - updated_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL - ){table_clause}'; - EXCEPTION - WHEN OTHERS THEN - IF SQLCODE != -955 THEN - RAISE; - END IF; - END; - - BEGIN - EXECUTE IMMEDIATE 'CREATE INDEX {index_name} - ON {self._table_name}(expires_at)'; - EXCEPTION - WHEN OTHERS THEN - IF SQLCODE != -955 THEN - RAISE; - END IF; - END; - """ + """Get Oracle CREATE TABLE SQL with optimized schema.""" + return _build_session_table_ddl(self._table_name, self._config, self._in_memory) def _drop_table_sql(self) -> "list[str]": - """Get Oracle DROP TABLE SQL with PL/SQL error handling. - - Returns: - List of SQL statements with exception handling for non-existent objects. - """ - index_name = f"idx_{self._table_name.replace('.', '_')}_expires_at"[:30] - return [ - f""" - BEGIN - EXECUTE IMMEDIATE 'DROP INDEX {index_name}'; - EXCEPTION - WHEN OTHERS THEN - IF SQLCODE != -1418 THEN - RAISE; - END IF; - END; - """, - f""" - BEGIN - EXECUTE IMMEDIATE 'DROP TABLE {self._table_name}'; - EXCEPTION - WHEN OTHERS THEN - IF SQLCODE != -942 THEN - RAISE; - END IF; - END; - """, - ] + """Get Oracle DROP TABLE SQL with PL/SQL error handling.""" + return _build_drop_session_table_sql(self._table_name) class OracleSyncStore(BaseSQLSpecStore["OracleSyncConfig"]): @@ -415,6 +387,8 @@ def __init__(self, config: "OracleSyncConfig") -> None: config: OracleSyncConfig instance. """ super().__init__(config) + for part in self._table_name.split("."): + validate_oracle_identifier(part, "table name") litestar_config = config.extension_config.get("litestar", {}) self._in_memory = bool(litestar_config.get("in_memory", False)) @@ -492,69 +466,12 @@ async def delete_expired(self) -> int: return await async_(self._delete_expired)() def _table_ddl(self) -> str: - """Get Oracle CREATE TABLE SQL with optimized schema. - - Returns: - SQL statement to create the sessions table with proper indexes. - """ - table_clause = _litestar_table_feature_clause(self._config, self._in_memory) - index_name = f"idx_{self._table_name.replace('.', '_')}_expires_at"[:30] - return f""" - BEGIN - EXECUTE IMMEDIATE 'CREATE TABLE {self._table_name} ( - session_id VARCHAR2(255) PRIMARY KEY, - data BLOB NOT NULL, - expires_at TIMESTAMP WITH TIME ZONE, - created_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL, - updated_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL - ){table_clause}'; - EXCEPTION - WHEN OTHERS THEN - IF SQLCODE != -955 THEN - RAISE; - END IF; - END; - - BEGIN - EXECUTE IMMEDIATE 'CREATE INDEX {index_name} - ON {self._table_name}(expires_at)'; - EXCEPTION - WHEN OTHERS THEN - IF SQLCODE != -955 THEN - RAISE; - END IF; - END; - """ + """Get Oracle CREATE TABLE SQL with optimized schema.""" + return _build_session_table_ddl(self._table_name, self._config, self._in_memory) def _drop_table_sql(self) -> "list[str]": - """Get Oracle DROP TABLE SQL with PL/SQL error handling. - - Returns: - List of SQL statements with exception handling for non-existent objects. - """ - index_name = f"idx_{self._table_name.replace('.', '_')}_expires_at"[:30] - return [ - f""" - BEGIN - EXECUTE IMMEDIATE 'DROP INDEX {index_name}'; - EXCEPTION - WHEN OTHERS THEN - IF SQLCODE != -1418 THEN - RAISE; - END IF; - END; - """, - f""" - BEGIN - EXECUTE IMMEDIATE 'DROP TABLE {self._table_name}'; - EXCEPTION - WHEN OTHERS THEN - IF SQLCODE != -942 THEN - RAISE; - END IF; - END; - """, - ] + """Get Oracle DROP TABLE SQL with PL/SQL error handling.""" + return _build_drop_session_table_sql(self._table_name) def _create_table(self) -> None: """Synchronous implementation of create_table.""" @@ -568,12 +485,12 @@ def _get(self, key: str, renew_for: "int | timedelta | None" = None) -> "bytes | """Synchronous implementation of get.""" sql = f""" SELECT data, expires_at FROM {self._table_name} - WHERE session_id = :session_id + WHERE session_id = :1 AND (expires_at IS NULL OR expires_at > SYSTIMESTAMP) """ with self._config.provide_connection() as conn, conn.cursor() as cursor: - cursor.execute(sql, {"session_id": key}) + cursor.execute(sql, (key,)) row = cursor.fetchone() if row is None: @@ -587,14 +504,19 @@ def _get(self, key: str, renew_for: "int | timedelta | None" = None) -> "bytes | update_sql = f""" UPDATE {self._table_name} SET expires_at = CASE - WHEN :expires_in_seconds IS NULL THEN NULL - ELSE SYSTIMESTAMP + NUMTODSINTERVAL(:expires_in_seconds, 'SECOND') + WHEN :1 IS NULL THEN NULL + ELSE SYSTIMESTAMP + NUMTODSINTERVAL(:1, 'SECOND') END, updated_at = SYSTIMESTAMP - WHERE session_id = :session_id + WHERE session_id = :2 """ - cursor.execute(update_sql, {"expires_in_seconds": expires_in_seconds, "session_id": key}) - conn.commit() + try: + cursor.execute(update_sql, (expires_in_seconds, expires_in_seconds, key)) + conn.commit() + except BaseException: + with contextlib.suppress(Exception): + conn.rollback() + raise return _read_blob_sync(data_blob) @@ -630,38 +552,57 @@ def _set(self, key: str, value: "str | bytes", expires_in: "int | timedelta | No """ with self._config.provide_connection() as conn: - bind_data: Any = conn.createlob(DB_TYPE_BLOB, data) if len(data) > ORACLE_SMALL_BLOB_LIMIT else data - with conn.cursor() as cursor: - cursor.execute(sql, {"session_id": key, "data": bind_data, "expires_in_seconds": expires_in_seconds}) - conn.commit() + try: + bind_data: Any = conn.createlob(DB_TYPE_BLOB, data) if len(data) > ORACLE_SMALL_BLOB_LIMIT else data + with conn.cursor() as cursor: + cursor.execute( + sql, {"session_id": key, "data": bind_data, "expires_in_seconds": expires_in_seconds} + ) + conn.commit() + except BaseException: + with contextlib.suppress(Exception): + conn.rollback() + raise def _delete(self, key: str) -> None: """Synchronous implementation of delete.""" - sql = f"DELETE FROM {self._table_name} WHERE session_id = :session_id" + sql = f"DELETE FROM {self._table_name} WHERE session_id = :1" - with self._config.provide_connection() as conn, conn.cursor() as cursor: - cursor.execute(sql, {"session_id": key}) - conn.commit() + with self._config.provide_connection() as conn: + try: + with conn.cursor() as cursor: + cursor.execute(sql, (key,)) + conn.commit() + except BaseException: + with contextlib.suppress(Exception): + conn.rollback() + raise def _delete_all(self) -> None: """Synchronous implementation of delete_all.""" sql = f"DELETE FROM {self._table_name}" - with self._config.provide_connection() as conn, conn.cursor() as cursor: - cursor.execute(sql) - conn.commit() + with self._config.provide_connection() as conn: + try: + with conn.cursor() as cursor: + cursor.execute(sql) + conn.commit() + except BaseException: + with contextlib.suppress(Exception): + conn.rollback() + raise self._log_delete_all() def _exists(self, key: str) -> bool: """Synchronous implementation of exists.""" sql = f""" SELECT 1 FROM {self._table_name} - WHERE session_id = :session_id + WHERE session_id = :1 AND (expires_at IS NULL OR expires_at > SYSTIMESTAMP) """ with self._config.provide_connection() as conn, conn.cursor() as cursor: - cursor.execute(sql, {"session_id": key}) + cursor.execute(sql, (key,)) result = cursor.fetchone() return result is not None @@ -669,11 +610,11 @@ def _expires_in(self, key: str) -> "int | None": """Synchronous implementation of expires_in.""" sql = f""" SELECT expires_at, SYSTIMESTAMP FROM {self._table_name} - WHERE session_id = :session_id + WHERE session_id = :1 """ with self._config.provide_connection() as conn, conn.cursor() as cursor: - cursor.execute(sql, {"session_id": key}) + cursor.execute(sql, (key,)) row = cursor.fetchone() if row is None or row[0] is None: @@ -696,13 +637,19 @@ def _delete_expired(self) -> int: """Synchronous implementation of delete_expired.""" sql = f"DELETE FROM {self._table_name} WHERE expires_at <= SYSTIMESTAMP" - with self._config.provide_connection() as conn, conn.cursor() as cursor: - cursor.execute(sql) - count = cursor.rowcount if cursor.rowcount is not None else 0 - conn.commit() - if count > 0: - self._log_delete_expired(count) - return count + with self._config.provide_connection() as conn: + try: + with conn.cursor() as cursor: + cursor.execute(sql) + count = resolve_rowcount(cursor) + conn.commit() + if count > 0: + self._log_delete_expired(count) + return count + except BaseException: + with contextlib.suppress(Exception): + conn.rollback() + raise def _oracle_expiry_seconds(expires_in: "int | timedelta | None") -> "int | None": @@ -729,18 +676,73 @@ def _coerce_bytes_payload(value: object) -> bytes: async def _read_blob_async(value: object) -> bytes: """Read LOB values from async connections into bytes.""" - if is_async_readable(value): - return _coerce_bytes_payload(await value.read()) - if is_readable(value): - return _coerce_bytes_payload(value.read()) - return _coerce_bytes_payload(value) + data = await read_oracle_lob_async(value) + return _coerce_bytes_payload(data) def _read_blob_sync(value: object) -> bytes: """Read LOB values from sync connections into bytes.""" - if is_readable(value): - return _coerce_bytes_payload(value.read()) - return _coerce_bytes_payload(value) + data = read_oracle_lob_sync(value) + return _coerce_bytes_payload(data) + + +def _build_session_table_ddl(table_name: str, config: Any, in_memory: bool) -> str: + """Build Oracle CREATE TABLE and INDEX SQL for the session store.""" + table_clause = _litestar_table_feature_clause(config, in_memory) + index_name = f"idx_{table_name.replace('.', '_')}_expires_at"[:30] + return f""" + BEGIN + EXECUTE IMMEDIATE 'CREATE TABLE {table_name} ( + session_id VARCHAR2(255) PRIMARY KEY, + data BLOB NOT NULL, + expires_at TIMESTAMP WITH TIME ZONE, + created_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL, + updated_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL + ){table_clause}'; + EXCEPTION + WHEN OTHERS THEN + IF SQLCODE != -955 THEN + RAISE; + END IF; + END; + + BEGIN + EXECUTE IMMEDIATE 'CREATE INDEX {index_name} + ON {table_name}(expires_at)'; + EXCEPTION + WHEN OTHERS THEN + IF SQLCODE != -955 THEN + RAISE; + END IF; + END; + """ + + +def _build_drop_session_table_sql(table_name: str) -> "list[str]": + """Build Oracle DROP TABLE and DROP INDEX statements for the session store.""" + index_name = f"idx_{table_name.replace('.', '_')}_expires_at"[:30] + return [ + f""" + BEGIN + EXECUTE IMMEDIATE 'DROP INDEX {index_name}'; + EXCEPTION + WHEN OTHERS THEN + IF SQLCODE != -1418 THEN + RAISE; + END IF; + END; + """, + f""" + BEGIN + EXECUTE IMMEDIATE 'DROP TABLE {table_name}'; + EXCEPTION + WHEN OTHERS THEN + IF SQLCODE != -942 THEN + RAISE; + END IF; + END; + """, + ] def _litestar_table_feature_clause(config: Any, in_memory: bool) -> str: @@ -756,4 +758,4 @@ def _litestar_table_feature_clause(config: Any, in_memory: bool) -> str: range_partition_key="expires_at", table_options_key="table_options", ) - return report["clause"] + return report["plsql_clause"] diff --git a/sqlspec/adapters/oracledb/migrations.py b/sqlspec/adapters/oracledb/migrations.py index 7e24f2bff..37dce84fb 100644 --- a/sqlspec/adapters/oracledb/migrations.py +++ b/sqlspec/adapters/oracledb/migrations.py @@ -5,9 +5,8 @@ """ import getpass -from typing import TYPE_CHECKING, cast - -from rich.console import Console +import importlib +from typing import TYPE_CHECKING, Any, cast from sqlspec.builder import CreateTable, Select, sql from sqlspec.migrations.tracker import AsyncMigrationTracker, SyncMigrationTracker @@ -24,7 +23,14 @@ __all__ = ("OracleAsyncMigrationTracker", "OracleSyncMigrationTracker") logger = get_logger("sqlspec.migrations.oracle") -console = Console() +_console: Any = None + + +def _get_console() -> Any: + global _console + if _console is None: + _console = importlib.import_module("rich.console").Console() + return _console class OracleMigrationTrackerMixin: @@ -367,7 +373,7 @@ def _migrate_schema_if_needed(self, driver: "SyncDriverAdapterBase") -> None: return if self._should_echo(): - console.print( + _get_console().print( f"[cyan]Migrating tracking table schema, adding columns: {', '.join(sorted(missing_columns))}[/]" ) @@ -376,7 +382,7 @@ def _migrate_schema_if_needed(self, driver: "SyncDriverAdapterBase") -> None: driver.commit() if self._should_echo(): - console.print("[green]Migration tracking table schema updated successfully[/]") + _get_console().print("[green]Migration tracking table schema updated successfully[/]") except Exception as e: logger.warning("Could not check or migrate tracking table schema: %s", e) @@ -593,7 +599,7 @@ async def _migrate_schema_if_needed(self, driver: "AsyncDriverAdapterBase") -> N return if self._should_echo(): - console.print( + _get_console().print( f"[cyan]Migrating tracking table schema, adding columns: {', '.join(sorted(missing_columns))}[/]" ) @@ -602,7 +608,7 @@ async def _migrate_schema_if_needed(self, driver: "AsyncDriverAdapterBase") -> N await driver.commit() if self._should_echo(): - console.print("[green]Migration tracking table schema updated successfully[/]") + _get_console().print("[green]Migration tracking table schema updated successfully[/]") except Exception as e: logger.warning("Could not check or migrate tracking table schema: %s", e) diff --git a/sqlspec/adapters/oracledb/type_converter.py b/sqlspec/adapters/oracledb/type_converter.py index 1011a3c60..25e3b773c 100644 --- a/sqlspec/adapters/oracledb/type_converter.py +++ b/sqlspec/adapters/oracledb/type_converter.py @@ -7,22 +7,20 @@ import array from typing import Any -from sqlspec.adapters.oracledb._vector_handlers import ( - numpy_converter_in, # pyright: ignore[reportPrivateUsage] - numpy_converter_out, # pyright: ignore[reportPrivateUsage] -) +from sqlspec.adapters.oracledb._vector_handlers import numpy_converter_in, numpy_converter_out from sqlspec.typing import NUMPY_INSTALLED +from sqlspec.utils.module_loader import import_optional from sqlspec.utils.sync_tools import ensure_async_ from sqlspec.utils.type_guards import is_async_readable, is_readable -try: - import numpy as np -except ImportError: - np = None # type: ignore[assignment] - __all__ = ("OracleOutputConverter",) +def _numpy_module() -> Any: + """Return the ``numpy`` module when installed, otherwise ``None``.""" + return import_optional("numpy") + + class OracleOutputConverter: """Oracle-specific output conversion with LOB optimization. @@ -82,7 +80,11 @@ def convert_numpy_to_vector(self, value: Any) -> Any: array.array compatible with Oracle VECTOR if value is ndarray, otherwise original value. """ - if not NUMPY_INSTALLED or np is None: + if not NUMPY_INSTALLED: + return value + + np = _numpy_module() + if np is None: return value if isinstance(value, np.ndarray): diff --git a/sqlspec/adapters/psqlpy/_typing.py b/sqlspec/adapters/psqlpy/_typing.py index 7be4219bf..914679c8f 100644 --- a/sqlspec/adapters/psqlpy/_typing.py +++ b/sqlspec/adapters/psqlpy/_typing.py @@ -8,60 +8,27 @@ import importlib from typing import TYPE_CHECKING, Any -from sqlspec.typing import import_optional_attr - - -class _PsqlpyUnavailableError(Exception): - """Fallback psqlpy exception base when optional exception classes are unavailable.""" - +from psqlpy import Connection as PsqlpyConnection +from psqlpy import ConnectionPool as PsqlpyConnectionPool +from psqlpy import Listener as PsqlpyListener +from psqlpy.exceptions import ConnectionExecuteError as PsqlpyConnectionExecuteError +from psqlpy.exceptions import DatabaseError as PsqlpyDatabaseError +from psqlpy.exceptions import DataError as PsqlpyDataError +from psqlpy.exceptions import Error as PsqlpyError +from psqlpy.exceptions import IntegrityError as PsqlpyIntegrityError +from psqlpy.exceptions import NotSupportedError as PsqlpyNotSupportedError +from psqlpy.exceptions import OperationalError as PsqlpyOperationalError +from psqlpy.extra_types import JSONB as PSQLPY_JSONB + +RustPSQLDriverPyBaseError = PsqlpyError if TYPE_CHECKING: from collections.abc import Callable from types import TracebackType - from typing import TypeAlias - - from psqlpy import Connection as _PsqlpyConnection - from psqlpy import ConnectionPool as PsqlpyConnectionPool - from psqlpy import Listener as _PsqlpyListener - from psqlpy.exceptions import ConnectionExecuteError as _PsqlpyConnectionExecuteError - from psqlpy.exceptions import DatabaseError as _PsqlpyDatabaseError - from psqlpy.exceptions import DataError as _PsqlpyDataError - from psqlpy.exceptions import Error as _PsqlpyError - from psqlpy.exceptions import IntegrityError as _PsqlpyIntegrityError - from psqlpy.exceptions import NotSupportedError as _PsqlpyNotSupportedError - from psqlpy.exceptions import OperationalError as _PsqlpyOperationalError - from psqlpy.extra_types import JSONB as PSQLPY_JSONB from sqlspec.adapters.psqlpy.driver import PsqlpyDriver from sqlspec.core import StatementConfig - PsqlpyConnection: TypeAlias = _PsqlpyConnection - PsqlpyDataError: TypeAlias = _PsqlpyDataError - PsqlpyDatabaseError: TypeAlias = _PsqlpyDatabaseError - PsqlpyConnectionExecuteError: TypeAlias = _PsqlpyConnectionExecuteError - PsqlpyError: TypeAlias = _PsqlpyError - PsqlpyIntegrityError: TypeAlias = _PsqlpyIntegrityError - PsqlpyListener: TypeAlias = _PsqlpyListener - PsqlpyNotSupportedError: TypeAlias = _PsqlpyNotSupportedError - PsqlpyOperationalError: TypeAlias = _PsqlpyOperationalError - - -else: - PSQLPY_JSONB = import_optional_attr("psqlpy.extra_types", "JSONB") - PsqlpyConnection = import_optional_attr("psqlpy", "Connection") or Any - PsqlpyConnectionPool = import_optional_attr("psqlpy", "ConnectionPool") or Any - PsqlpyDataError = import_optional_attr("psqlpy.exceptions", "DataError") or _PsqlpyUnavailableError - PsqlpyDatabaseError = import_optional_attr("psqlpy.exceptions", "DatabaseError") or _PsqlpyUnavailableError - PsqlpyConnectionExecuteError = ( - import_optional_attr("psqlpy.exceptions", "ConnectionExecuteError") or _PsqlpyUnavailableError - ) - PsqlpyError = import_optional_attr("psqlpy.exceptions", "Error") or _PsqlpyUnavailableError - PsqlpyIntegrityError = import_optional_attr("psqlpy.exceptions", "IntegrityError") or _PsqlpyUnavailableError - PsqlpyListener = import_optional_attr("psqlpy", "Listener") or Any - PsqlpyNotSupportedError = import_optional_attr("psqlpy.exceptions", "NotSupportedError") or _PsqlpyUnavailableError - PsqlpyOperationalError = import_optional_attr("psqlpy.exceptions", "OperationalError") or _PsqlpyUnavailableError - - __all__ = ( "PSQLPY_JSONB", "PsqlpyConnection", @@ -76,6 +43,7 @@ class _PsqlpyUnavailableError(Exception): "PsqlpyNotSupportedError", "PsqlpyOperationalError", "PsqlpySessionContext", + "RustPSQLDriverPyBaseError", ) @@ -168,6 +136,13 @@ async def __aexit__( return None +_CACHED_PSQLPY_DRIVER_CLS: "type[PsqlpyDriver] | None" = None + + def _resolve_psqlpy_driver() -> "type[PsqlpyDriver]": - module = importlib.import_module("sqlspec.adapters.psqlpy.driver") - return module.PsqlpyDriver # type: ignore[no-any-return] + global _CACHED_PSQLPY_DRIVER_CLS + if _CACHED_PSQLPY_DRIVER_CLS is None: + module = importlib.import_module("sqlspec.adapters.psqlpy.driver") + _CACHED_PSQLPY_DRIVER_CLS = module.PsqlpyDriver + assert _CACHED_PSQLPY_DRIVER_CLS is not None + return _CACHED_PSQLPY_DRIVER_CLS diff --git a/sqlspec/adapters/psqlpy/adk/store.py b/sqlspec/adapters/psqlpy/adk/store.py index de64aa015..6ca32d356 100644 --- a/sqlspec/adapters/psqlpy/adk/store.py +++ b/sqlspec/adapters/psqlpy/adk/store.py @@ -5,10 +5,12 @@ from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, NotRequired, cast from sqlspec.adapters.psqlpy._typing import PsqlpyConnectionExecuteError, PsqlpyDatabaseError +from sqlspec.adapters.psqlpy.core import build_postgres_storage_params_sql from sqlspec.config import ADKConfig from sqlspec.extensions.adk import BaseAsyncADKStore, StoredEvent, StoredSession, normalize_session_list_options from sqlspec.extensions.adk.memory.store import BaseAsyncADKMemoryStore from sqlspec.utils.logging import get_logger +from sqlspec.utils.serializers import to_json from sqlspec.utils.type_guards import has_query_result_metadata if TYPE_CHECKING: @@ -23,7 +25,6 @@ logger = get_logger("sqlspec.adapters.psqlpy.adk.store") -POSTGRES_TABLE_NOT_FOUND_SQLSTATE: Final = "42P01" PSQLPY_STATUS_REGEX: Final[re.Pattern[str]] = re.compile(r"^([A-Z]+)(?:\s+(\d+))?\s+(\d+)$", re.IGNORECASE) @@ -553,7 +554,9 @@ async def _sessions_table_ddl(self) -> str: if self._owner_id_column_ddl: owner_id_line = f",\n {self._owner_id_column_ddl}" - table_options = _postgres_table_options(_adk_config(self._config)) + table_options = build_postgres_storage_params_sql( + _adk_config(self._config), default_fillfactor=80, section="adk" + ) return f""" CREATE TABLE IF NOT EXISTS {self._session_table} ( id VARCHAR(128) PRIMARY KEY, @@ -577,7 +580,9 @@ async def _sessions_table_ddl(self) -> str: async def _events_table_ddl(self) -> str: adk_config = _adk_config(self._config) - table_options = _postgres_table_options(adk_config, include_autovacuum=True) + table_options = build_postgres_storage_params_sql( + adk_config, include_autovacuum=True, default_fillfactor=80, section="adk" + ) generated_columns, generated_indexes, covering_columns = _postgres_event_ddl_options( adk_config, self._events_table ) @@ -603,7 +608,9 @@ async def _events_table_ddl(self) -> str: """ async def _app_states_table_ddl(self) -> str: - table_options = _postgres_table_options(_adk_config(self._config)) + table_options = build_postgres_storage_params_sql( + _adk_config(self._config), default_fillfactor=80, section="adk" + ) return f""" CREATE TABLE IF NOT EXISTS {self._app_state_table} ( app_name VARCHAR(128) PRIMARY KEY, @@ -613,7 +620,9 @@ async def _app_states_table_ddl(self) -> str: """ async def _user_states_table_ddl(self) -> str: - table_options = _postgres_table_options(_adk_config(self._config)) + table_options = build_postgres_storage_params_sql( + _adk_config(self._config), default_fillfactor=80, section="adk" + ) return f""" CREATE TABLE IF NOT EXISTS {self._user_state_table} ( app_name VARCHAR(128) NOT NULL, @@ -685,73 +694,96 @@ async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: " if not entries: return 0 - inserted_count = 0 + ids = [entry["id"] for entry in entries] + session_ids = [entry["session_id"] for entry in entries] + app_names = [entry["app_name"] for entry in entries] + user_ids = [entry["user_id"] for entry in entries] + scopes = [entry.get("scope", "user") for entry in entries] + event_ids = [entry["event_id"] for entry in entries] + authors = [entry.get("author") for entry in entries] + timestamps = [entry["timestamp"] for entry in entries] + content_jsons = [entry["content_json"] for entry in entries] + content_texts = [entry["content_text"] for entry in entries] + metadata_jsons = [entry.get("metadata_json") for entry in entries] + inserted_ats = [entry["inserted_at"] for entry in entries] + embeddings = [ + to_json([float(value) for value in embedding]) + if (embedding := entry.get("embedding")) is not None + else None + for entry in entries + ] + if self._owner_id_column_name: sql = f""" INSERT INTO {self._memory_table} ( id, session_id, app_name, user_id, scope, event_id, author, - {self._owner_id_column_name}, timestamp, embedding, content_json, - content_text, metadata_json, inserted_at - ) VALUES ( - $1, $2, $3, $4, $5, $6, $7, $8, $9, $10::float8[]::vector, $11, $12, $13, $14 + {self._owner_id_column_name}, timestamp, content_json, + content_text, metadata_json, inserted_at, embedding ) + SELECT id, session_id, app_name, user_id, scope, event_id, author, + $14, timestamp, content_json, content_text, metadata_json, inserted_at, embedding::vector + FROM unnest( + $1::varchar[], $2::varchar[], $3::varchar[], $4::varchar[], + $5::varchar[], $6::varchar[], $7::varchar[], + $8::timestamptz[], $9::jsonb[], $10::text[], $11::jsonb[], $12::timestamptz[], $13::text[] + ) AS batch(id, session_id, app_name, user_id, scope, event_id, author, + timestamp, content_json, content_text, metadata_json, inserted_at, embedding) ON CONFLICT (event_id) DO NOTHING + RETURNING 1 """ + params = [ + ids, + session_ids, + app_names, + user_ids, + scopes, + event_ids, + authors, + timestamps, + content_jsons, + content_texts, + metadata_jsons, + inserted_ats, + embeddings, + owner_id, + ] else: sql = f""" INSERT INTO {self._memory_table} ( id, session_id, app_name, user_id, scope, event_id, author, - timestamp, embedding, content_json, content_text, metadata_json, inserted_at - ) VALUES ( - $1, $2, $3, $4, $5, $6, $7, $8, $9::float8[]::vector, $10, $11, $12, $13 + timestamp, content_json, content_text, metadata_json, inserted_at, embedding ) + SELECT id, session_id, app_name, user_id, scope, event_id, author, + timestamp, content_json, content_text, metadata_json, inserted_at, embedding::vector + FROM unnest( + $1::varchar[], $2::varchar[], $3::varchar[], $4::varchar[], + $5::varchar[], $6::varchar[], $7::varchar[], + $8::timestamptz[], $9::jsonb[], $10::text[], $11::jsonb[], $12::timestamptz[], $13::text[] + ) AS batch(id, session_id, app_name, user_id, scope, event_id, author, + timestamp, content_json, content_text, metadata_json, inserted_at, embedding) ON CONFLICT (event_id) DO NOTHING + RETURNING 1 """ + params = [ + ids, + session_ids, + app_names, + user_ids, + scopes, + event_ids, + authors, + timestamps, + content_jsons, + content_texts, + metadata_jsons, + inserted_ats, + embeddings, + ] async with self._config.provide_connection() as conn: - for entry in entries: - raw_embedding = entry.get("embedding") - embedding_val = [float(v) for v in raw_embedding] if raw_embedding is not None else None - if self._owner_id_column_name: - params = [ - entry["id"], - entry["session_id"], - entry["app_name"], - entry["user_id"], - entry.get("scope", "user"), - entry["event_id"], - entry.get("author"), - owner_id, - entry["timestamp"], - embedding_val, - entry["content_json"], - entry["content_text"], - entry.get("metadata_json"), - entry["inserted_at"], - ] - else: - params = [ - entry["id"], - entry["session_id"], - entry["app_name"], - entry["user_id"], - entry.get("scope", "user"), - entry["event_id"], - entry.get("author"), - entry["timestamp"], - embedding_val, - entry["content_json"], - entry["content_text"], - entry.get("metadata_json"), - entry["inserted_at"], - ] - - result = await conn.execute(sql, params) - rows_affected = self._extract_rows_affected(result) - if rows_affected > 0: - inserted_count += rows_affected - - return inserted_count + result = await conn.fetch(sql, params) + rows = result.result() if result else [] + return len(rows) async def search_entries( self, @@ -1156,33 +1188,6 @@ def _adk_config(config: Any) -> PsqlpyADKConfig: return cast("PsqlpyADKConfig", adk_config) -def _postgres_table_options(adk_config: PsqlpyADKConfig, *, include_autovacuum: bool = False) -> str: - options = [_postgres_fillfactor_option(adk_config)] - if include_autovacuum: - options.extend(_postgres_autovacuum_options(adk_config)) - return f" WITH ({', '.join(options)})" - - -def _postgres_fillfactor_option(adk_config: PsqlpyADKConfig) -> str: - value = adk_config.get("fillfactor", 80) - if not isinstance(value, int) or isinstance(value, bool) or value not in range(10, 101): - msg = "extension_config['adk']['fillfactor'] must be an integer from 10 to 100" - raise ValueError(msg) - return f"fillfactor = {value}" - - -def _postgres_autovacuum_options(adk_config: PsqlpyADKConfig) -> "list[str]": - options: list[str] = [] - for key in ("autovacuum_vacuum_scale_factor", "autovacuum_analyze_scale_factor"): - value = adk_config.get(key) - if value is not None: - if not isinstance(value, (int, float)) or isinstance(value, bool) or not 0 <= float(value) <= 1: - msg = f"extension_config['adk']['{key}'] must be a number from 0 to 1" - raise ValueError(msg) - options.append(f"{key} = {float(value):g}") - return options - - def _postgres_memory_indexes( memory_table: str, use_fts: bool, diff --git a/sqlspec/adapters/psqlpy/config.py b/sqlspec/adapters/psqlpy/config.py index 540f093bb..06b09eaba 100644 --- a/sqlspec/adapters/psqlpy/config.py +++ b/sqlspec/adapters/psqlpy/config.py @@ -148,10 +148,6 @@ class PsqlpyDriverFeatures(TypedDict): """Enable ``google_columnar_engine.enable_columnar_scan=on`` at session startup or per transaction.""" pooler_mode: NotRequired[Literal["session", "transaction", "statement", "auto"] | None] """External pooler compatibility mode; executes queries with ``prepared=False`` when not ``session``.""" - pgbouncer: NotRequired[bool] - """Alias for ``pooler_mode='transaction'`` disabling prepared execution.""" - pgbouncer_mode: NotRequired[bool] - """Alias for ``pooler_mode='transaction'`` disabling prepared execution.""" direct_dsn: NotRequired["str | Mapping[str, Any] | None"] """Direct PostgreSQL DSN or connection mapping bypassing the pooler for migrations and ``LISTEN/NOTIFY``.""" json_serializer: NotRequired["Callable[[Any], str]"] @@ -162,6 +158,8 @@ class PsqlpyDriverFeatures(TypedDict): """Async callback invoked once per physical ``psqlpy`` connection acquired from the pool.""" enable_events: NotRequired[bool] """Enable SQLSpec event channel pub/sub support (defaults to ``True`` when ``events`` is configured).""" + enable_pgcopy_binary: NotRequired[bool] + """Enable binary wire format for PostgreSQL COPY operations.""" events_backend: NotRequired[Literal["notify", "notify_queue", "poll_queue"]] """Event transport backend (``notify`` for LISTEN/NOTIFY, ``notify_queue``, or ``poll_queue``).""" @@ -240,10 +238,10 @@ class PsqlpyConfig(AsyncDatabaseConfig[PsqlpyConnection, "ConnectionPool", Psqlp connection_type: "ClassVar[type[PsqlpyConnection]]" = PsqlpyConnection supports_transactional_ddl: "ClassVar[bool]" = True supports_migration_schemas: "ClassVar[bool]" = True - supports_native_arrow_export: ClassVar[bool] = True + supports_native_arrow_export: ClassVar[bool] = False supports_native_arrow_import: ClassVar[bool] = True - supports_native_parquet_export: ClassVar[bool] = True - supports_native_parquet_import: ClassVar[bool] = True + supports_native_parquet_export: ClassVar[bool] = False + supports_native_parquet_import: ClassVar[bool] = False supports_native_row_streaming: ClassVar[bool] = True type_coercion_capabilities: "ClassVar[TypeCoercionCapabilities]" = TypeCoercionCapabilities( datetime_binding="native", timestamp_precision="microsecond", json_columns_decoded=True, uuid_binding="native" @@ -286,9 +284,7 @@ def __init__( statement_config, driver_features = apply_driver_features(statement_config, driver_features) features_dict = dict(driver_features) if driver_features else {} - is_transaction_pooler = features_dict.get("pooler_mode") in {"transaction", "statement"} or bool( - features_dict.get("pgbouncer") or features_dict.get("pgbouncer_mode") - ) + is_transaction_pooler = features_dict.get("pooler_mode") in {"transaction", "statement"} if bool(features_dict.get("enable_alloydb_columnar", False)) and not is_transaction_pooler: columnar_opt = "-c google_columnar_engine.enable_columnar_scan=on" existing_options = connection_config.get("options") @@ -411,8 +407,6 @@ def derive_direct_config(self) -> Self: conn_cfg.update(build_connection_config(direct_dsn)) features: dict[str, Any] = dict(self.driver_features) features["pooler_mode"] = "session" - features["pgbouncer"] = False - features["pgbouncer_mode"] = False features["direct_dsn"] = None if self._user_connection_hook is not None: features["on_connection_create"] = self._user_connection_hook @@ -456,7 +450,9 @@ async def _ensure_connection(self, connection: "PsqlpyConnection") -> None: conn_id = id(connection) if conn_id not in self._initialized_connection_ids: if self._user_connection_hook is not None: - await self._user_connection_hook(connection) + hook_res = self._user_connection_hook(connection) + if inspect.isawaitable(hook_res): + await hook_res if len(self._initialized_connection_ids) >= _MAX_INITIALIZED_CONNECTION_IDS: self._initialized_connection_ids.clear() self._initialized_connection_ids.add(conn_id) diff --git a/sqlspec/adapters/psqlpy/core.py b/sqlspec/adapters/psqlpy/core.py index 0101439a4..97cdb66ee 100644 --- a/sqlspec/adapters/psqlpy/core.py +++ b/sqlspec/adapters/psqlpy/core.py @@ -3,10 +3,11 @@ import contextlib import datetime as dt import decimal -import inspect import re import uuid +from collections.abc import Mapping from functools import lru_cache +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, cast import sqlglot @@ -14,7 +15,12 @@ from sqlglot.errors import ParseError from sqlspec.adapters.psqlpy._typing import PSQLPY_JSONB as _JSONB_TYPE -from sqlspec.adapters.psqlpy._typing import PsqlpyDataError, PsqlpyIntegrityError, PsqlpyOperationalError +from sqlspec.adapters.psqlpy._typing import ( + PsqlpyConnectionExecuteError, + PsqlpyDataError, + PsqlpyIntegrityError, + PsqlpyOperationalError, +) from sqlspec.adapters.psqlpy.type_converter import coerce_pgvector from sqlspec.core import ( DriverParameterProfile, @@ -29,7 +35,7 @@ resolve_postgres_extension_state, resolve_runtime_statement_config, ) -from sqlspec.core.parameters import type_coercion_dispatcher +from sqlspec.core.parameters import type_coercion_dispatcher, type_coercion_fallbacks from sqlspec.driver import CachedQuery from sqlspec.exceptions import ( CheckViolationError, @@ -58,7 +64,7 @@ from sqlspec.utils.uuids import uuid_from_string if TYPE_CHECKING: - from collections.abc import Callable, Mapping + from collections.abc import Callable from sqlspec.core import SQL @@ -67,6 +73,7 @@ "apply_driver_features", "build_connection_config", "build_postgres_extension_probe_names", + "build_postgres_storage_params_sql", "build_profile", "build_statement_config", "coerce_json_columns", @@ -74,7 +81,6 @@ "collect_rows", "create_mapped_exception", "default_statement_config", - "detect_postgres_extensions", "driver_profile", "extract_rows_affected", "format_execute_many_parameters", @@ -106,6 +112,39 @@ _NUMERIC_COERCE_TYPES: "tuple[type[Any], ...]" = (float, decimal.Decimal, list, tuple) +def build_postgres_storage_params_sql( + settings: "dict[str, object] | Mapping[str, object] | None", + *, + include_autovacuum: bool = False, + default_fillfactor: int | None = None, + default_vacuum_scale: float | None = None, + default_analyze_scale: float | None = None, + section: str = "storage", +) -> str: + """Format PostgreSQL table storage parameters (WITH fillfactor = ...) clause.""" + options: list[str] = [] + config_dict = dict(settings) if settings else {} + fillfactor = config_dict.get("fillfactor", default_fillfactor) + if fillfactor is not None: + if not isinstance(fillfactor, int) or isinstance(fillfactor, bool) or fillfactor not in range(10, 101): + msg = f"extension_config['{section}']['fillfactor'] must be an integer from 10 to 100" + raise ValueError(msg) + options.append(f"fillfactor = {fillfactor}") + if include_autovacuum: + scales = ( + ("autovacuum_vacuum_scale_factor", default_vacuum_scale), + ("autovacuum_analyze_scale_factor", default_analyze_scale), + ) + for key, default_val in scales: + value = config_dict.get(key, default_val) + if value is not None: + if not isinstance(value, (int, float)) or isinstance(value, bool) or not 0 <= float(value) <= 1: + msg = f"extension_config['{section}']['{key}'] must be a number from 0 to 1" + raise ValueError(msg) + options.append(f"{key} = {float(value):g}") + return f" WITH ({', '.join(options)})" if options else "" + + def build_profile() -> "DriverParameterProfile": """Create the psqlpy driver parameter profile.""" coercions: dict[type, Callable[[Any], Any]] = build_uuid_coercions(native=True) @@ -169,13 +208,7 @@ def build_connection_config(connection_config: "Mapping[str, Any]") -> "dict[str username = config.pop("username", None) or config.pop("user", None) if username is not None: config["username"] = username - auth_token_provider = config.get("auth_token_provider") - if callable(auth_token_provider) and not inspect.iscoroutinefunction(auth_token_provider): - config.pop("auth_token_provider", None) - token_result = auth_token_provider() - if not inspect.isawaitable(token_result): - config["password"] = str(token_result) - elif auth_token_provider is None: + if "auth_token_provider" in config and config["auth_token_provider"] is None: config.pop("auth_token_provider", None) return config @@ -193,12 +226,9 @@ def apply_driver_features( features.setdefault("enable_pgvector", PGVECTOR_INSTALLED) features.setdefault("enable_paradedb", True) features.setdefault("enable_pg_textsearch", True) - pgbouncer_enabled = bool(features.get("pgbouncer") or features.get("pgbouncer_mode")) - features["pgbouncer"] = pgbouncer_enabled - features["pgbouncer_mode"] = pgbouncer_enabled raw_pooler_mode = features.get("pooler_mode") if raw_pooler_mode is None: - features["pooler_mode"] = "transaction" if pgbouncer_enabled else "session" + features["pooler_mode"] = "session" else: normalized_mode = str(raw_pooler_mode).lower() if normalized_mode not in {"session", "transaction", "statement", "auto"}: @@ -206,10 +236,7 @@ def apply_driver_features( f"Invalid pooler_mode: {raw_pooler_mode!r}. Expected 'session', 'transaction', 'statement', or 'auto'." ) raise ImproperConfigurationError(msg) - features["pooler_mode"] = "transaction" if normalized_mode == "auto" and pgbouncer_enabled else normalized_mode - if features["pooler_mode"] in {"transaction", "statement"}: - features["pgbouncer"] = True - features["pgbouncer_mode"] = True + features["pooler_mode"] = "session" if normalized_mode == "auto" else normalized_mode features.setdefault("direct_dsn", None) features.setdefault("enable_alloydb_columnar", False) @@ -417,23 +444,28 @@ def extract_rows_affected(result: Any) -> int: return 0 -def get_parameter_casts(statement: "SQL | CachedQuery") -> "dict[int, str]": +_EMPTY_PARAMETER_CASTS: Final[Mapping[int, str]] = MappingProxyType({}) + + +def get_parameter_casts(statement: "SQL | CachedQuery") -> "Mapping[int, str]": """Get parameter cast metadata from a SQLSpec statement.""" if isinstance(statement, CachedQuery): return statement.parameter_casts state = statement.get_processed_state() - return {} if state is Empty else state.parameter_casts or {} + return _EMPTY_PARAMETER_CASTS if state is Empty or not state.parameter_casts else state.parameter_casts def prepare_parameters_with_casts( - parameters: Any, parameter_casts: "dict[int, str]", statement_config: "StatementConfig" + parameters: Any, parameter_casts: "Mapping[int, str]", statement_config: "StatementConfig" ) -> Any: """Prepare parameters with cast-aware type coercion.""" if isinstance(parameters, (list, tuple)): - result: list[Any] = [] serializer = statement_config.parameter_config.json_serializer or to_json type_map = statement_config.parameter_config.type_coercion_map - dispatcher = type_coercion_dispatcher(tuple(type_map.items())) if type_map else None + fallbacks = type_coercion_fallbacks(type_map) if type_map else () + dispatcher = type_coercion_dispatcher(fallbacks) if fallbacks else None + coerced: list[Any] | None = None + for idx, param in enumerate(parameters, start=1): cast_type = parameter_casts.get(idx, "") prepared_value = param @@ -447,8 +479,14 @@ def prepare_parameters_with_casts( prepared_value = fallback_converter(prepared_value) if cast_type: prepared_value = _coerce_parameter_for_cast(prepared_value, cast_type, serializer) - result.append(prepared_value) - return tuple(result) if isinstance(parameters, tuple) else result + if prepared_value is not param and coerced is None: + coerced = list(parameters[: idx - 1]) + if coerced is not None: + coerced.append(prepared_value) + + if coerced is not None: + return tuple(coerced) if isinstance(parameters, tuple) else coerced + return parameters return parameters @@ -478,14 +516,28 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpec Returns: A SQLSpec exception that wraps the original error """ - del logger + _ = logger error_msg = str(error).lower() - if "unique" in error_msg or "duplicate key" in error_msg: + if isinstance(error, PsqlpyDataError): + return _create_postgres_error(error, DataError, "data error") + + if "syntax error" in error_msg or "parse error" in error_msg: + return _create_postgres_error(error, SQLParsingError, "SQL syntax error") + + if ( + "unique constraint" in error_msg + or "duplicate key value" in error_msg + or "violates unique constraint" in error_msg + ): return _create_postgres_error(error, UniqueViolationError, "unique constraint violation") - if "foreign key" in error_msg or "violates foreign key" in error_msg: + if "foreign key constraint" in error_msg or "violates foreign key constraint" in error_msg: return _create_postgres_error(error, ForeignKeyViolationError, "foreign key constraint violation") - if "not null" in error_msg or ("null value" in error_msg and "violates not-null" in error_msg): + if ( + "not-null constraint" in error_msg + or "violates not-null constraint" in error_msg + or ("null value in column" in error_msg and "violates not-null" in error_msg) + ): return _create_postgres_error(error, NotNullViolationError, "not-null constraint violation") if "check constraint" in error_msg or "violates check constraint" in error_msg: return _create_postgres_error(error, CheckViolationError, "check constraint violation") @@ -497,13 +549,18 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpec if "serialization failure" in error_msg or "could not serialize" in error_msg: return _create_postgres_error(error, SerializationConflictError, "serialization failure") + if "connection" in error_msg: + if "timeout" in error_msg or "timed out" in error_msg: + return _create_postgres_error(error, ConnectionTimeoutError, "connection timeout") + return _create_postgres_error(error, DatabaseConnectionError, "connection error") + if ( - "connection" in error_msg - or "could not connect" in error_msg - or "connection refused" in error_msg + "connection refused" in error_msg or "connection reset" in error_msg - or "broken pipe" in error_msg + or "could not connect to server" in error_msg or "server closed the connection" in error_msg + or "broken pipe" in error_msg + or "failed to connect" in error_msg ): if "timeout" in error_msg or "timed out" in error_msg: return _create_postgres_error(error, ConnectionTimeoutError, "connection timeout") @@ -521,16 +578,17 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpec ): return _create_postgres_error(error, PermissionDeniedError, "authentication error") - if isinstance(error, PsqlpyDataError): - return _create_postgres_error(error, DataError, "data error") - - if "syntax error" in error_msg or "parse" in error_msg: - return _create_postgres_error(error, SQLParsingError, "SQL syntax error") - if isinstance(error, PsqlpyOperationalError): return _create_postgres_error(error, OperationalError, "operational error") if isinstance(error, PsqlpyIntegrityError): return _create_postgres_error(error, IntegrityError, "integrity constraint violation") + if isinstance(error, (PsqlpyConnectionExecuteError,)): + return _create_postgres_error(error, DatabaseConnectionError, "connection execute error") + + if "connection" in error_msg: + if "timeout" in error_msg or "timed out" in error_msg: + return _create_postgres_error(error, ConnectionTimeoutError, "connection timeout") + return _create_postgres_error(error, DatabaseConnectionError, "connection error") return _create_postgres_error(error, SQLSpecError, "database error") @@ -572,41 +630,6 @@ def format_execute_many_parameters(parameters: Any, *, coerce_numeric: bool) -> return [_format_execute_many_param_set(parameters, coerce_numeric=coerce_numeric)] -async def detect_postgres_extensions( - connection: Any, *, enable_pgvector: bool = False, enable_paradedb: bool = False, enable_pg_textsearch: bool = False -) -> "tuple[bool, bool, bool]": - """Detect whether pgvector, paradedb, and pg_textsearch extensions are installed on PostgreSQL. - - Args: - connection: Active psqlpy connection. - enable_pgvector: Whether to probe for pgvector extension. - enable_paradedb: Whether to probe for paradedb extension. - enable_pg_textsearch: Whether to probe for pg_textsearch extension. - - Returns: - Tuple of ``(pgvector_available, paradedb_available, pg_textsearch_available)``. - """ - driver_features: dict[str, Any] = { - "enable_pgvector": enable_pgvector, - "enable_paradedb": enable_paradedb, - "enable_pg_textsearch": enable_pg_textsearch, - } - extensions_to_check = build_postgres_extension_probe_names(driver_features) - if not extensions_to_check: - return False, False, False - - try: - result = await connection.fetch( - "SELECT extname FROM pg_extension WHERE extname = ANY($1::text[])", [extensions_to_check] - ) - rows = result.result() if result else [] - detected = {cast("str", row["extname"] if isinstance(row, dict) else row[0]) for row in rows} - except Exception: - logger.debug("PostgreSQL extension detection failed for psqlpy connection") - return False, False, False - return "vector" in detected, "pg_search" in detected, "pg_textsearch" in detected - - def _coerce_json_parameter(value: Any, cast_type: str, serializer: "Callable[[Any], str]") -> Any: """Serialize JSON parameters according to the detected cast type.""" diff --git a/sqlspec/adapters/psqlpy/data_dictionary.py b/sqlspec/adapters/psqlpy/data_dictionary.py index a45f39843..a4706a434 100644 --- a/sqlspec/adapters/psqlpy/data_dictionary.py +++ b/sqlspec/adapters/psqlpy/data_dictionary.py @@ -28,7 +28,12 @@ system_metadata_gated_result, unsupported_system_metadata_capability, ) -from sqlspec.data_dictionary.dialects.postgres import resolve_postgres_json_type +from sqlspec.data_dictionary.dialects.postgres import ( + POSTGRES_METADATA_DOMAINS, + POSTGRES_SUPPORTED_DOMAINS, + POSTGRES_SYSTEM_METADATA_QUERIES, + resolve_postgres_json_type, +) from sqlspec.driver import AsyncDataDictionaryBase if TYPE_CHECKING: @@ -38,37 +43,6 @@ __all__ = ("PsqlpyDataDictionary",) -_POSTGRES_METADATA_DOMAINS: tuple[str, ...] = ( - "schemas", - "objects", - "tables", - "columns", - "constraints", - "indexes", - "views", - "materialized_views", - "sequences", - "types", - "routines", - "triggers", - "comments", - "privileges", - "dependencies", - "extensions", - "partitions", - "ddl", - "system", -) -_POSTGRES_SUPPORTED_DOMAINS = frozenset(d for d in _POSTGRES_METADATA_DOMAINS if d != "system") -_POSTGRES_SYSTEM_METADATA_QUERIES = { - "settings": "settings", - "statement_history": "pg_stat_statements", - "pg_stat_statements": "pg_stat_statements", - "table_statistics": "table_stats", - "table_stats": "table_stats", -} - - @mypyc_attr(allow_interpreted_subclasses=True, native_class=False) class PsqlpyDataDictionary(AsyncDataDictionaryBase): """PostgreSQL-specific async data dictionary via psqlpy.""" @@ -89,7 +63,7 @@ async def get_system_metadata_capabilities( ) -> tuple[SystemMetadataCapability, ...]: """Get PostgreSQL opt-in system metadata capability disclosures.""" _ = driver - requested_domains = tuple(_POSTGRES_SYSTEM_METADATA_QUERIES) if domains is None else tuple(domains) + requested_domains = tuple(POSTGRES_SYSTEM_METADATA_QUERIES) if domains is None else tuple(domains) return tuple(_postgres_system_metadata_capability(domain) for domain in requested_domains) async def get_schemas(self, driver: "PsqlpyDriver") -> MetadataResult: @@ -163,7 +137,7 @@ async def get_system_metadata( if gate_result.capability.support != MetadataSupport.SUPPORTED: return gate_result - query_name = _POSTGRES_SYSTEM_METADATA_QUERIES.get(metadata_request.domain) + query_name = POSTGRES_SYSTEM_METADATA_QUERIES.get(metadata_request.domain) if query_name is None: return gate_result parameters: dict[str, object] = {} @@ -266,10 +240,11 @@ async def get_indexes( ) table_name = self.resolve_identifier(table) - self._log_table_describe(driver, schema_name=schema_name, table_name=table_name, operation="indexes") + table_schema_name = None if schema is None else schema_name + self._log_table_describe(driver, schema_name=table_schema_name, table_name=table_name, operation="indexes") return await driver.select( _postgres_domain_sql("indexes", "by_schema"), - schema_name=schema_name, + schema_name=table_schema_name, table_name=table_name, schema_type=IndexMetadata, ) @@ -334,7 +309,7 @@ def _postgres_metadata_capability(domain: str) -> MetadataCapability: risks=(MetadataRisk.EXPENSIVE, MetadataRisk.PRIVILEGED), warnings=("System metadata is opt-in and disabled by default.",), ) - if domain in _POSTGRES_SUPPORTED_DOMAINS: + if domain in POSTGRES_SUPPORTED_DOMAINS: return MetadataCapability( domain=domain, support=MetadataSupport.SUPPORTED, @@ -345,7 +320,7 @@ def _postgres_metadata_capability(domain: str) -> MetadataCapability: def _postgres_metadata_profile(adapter: str, domains: Sequence[str] | None) -> MetadataCapabilityProfile: - requested_domains = _POSTGRES_METADATA_DOMAINS if domains is None else tuple(domains) + requested_domains = POSTGRES_METADATA_DOMAINS if domains is None else tuple(domains) return MetadataCapabilityProfile( "postgres", adapter=adapter, @@ -404,7 +379,7 @@ def _ddl_result_from_rows( def _postgres_system_metadata_capability(domain: str) -> SystemMetadataCapability: - if domain not in _POSTGRES_SYSTEM_METADATA_QUERIES: + if domain not in POSTGRES_SYSTEM_METADATA_QUERIES: return unsupported_system_metadata_capability(domain) return SystemMetadataCapability( domain, diff --git a/sqlspec/adapters/psqlpy/driver.py b/sqlspec/adapters/psqlpy/driver.py index 8abc1f355..4a4429a96 100644 --- a/sqlspec/adapters/psqlpy/driver.py +++ b/sqlspec/adapters/psqlpy/driver.py @@ -4,7 +4,9 @@ and transaction management. """ +from contextlib import suppress from itertools import chain +from time import perf_counter from typing import TYPE_CHECKING, Any, Literal, cast from sqlspec.adapters.psqlpy._typing import ( @@ -32,7 +34,7 @@ split_schema_and_table, ) from sqlspec.adapters.psqlpy.data_dictionary import PsqlpyDataDictionary -from sqlspec.core import SQL, StatementConfig, get_cache_config, register_driver_profile +from sqlspec.core import SQL, DMLResult, StatementConfig, get_cache_config, register_driver_profile from sqlspec.core.sqlcommenter import split_sqlcomment from sqlspec.driver import AsyncDriverAdapterBase, AsyncRowStream, BaseAsyncExceptionHandler from sqlspec.exceptions import SQLSpecError @@ -44,7 +46,7 @@ from sqlspec.adapters.psqlpy._typing import PsqlpyConnection from sqlspec.core import ArrowResult, SQLResult - from sqlspec.driver import ExecutionResult + from sqlspec.driver import CachedQuery, ExecutionResult from sqlspec.storage import StorageBridgeJob, StorageDestination, StorageFormat, StorageTelemetry __all__ = ("PsqlpyCursor", "PsqlpyDriver", "PsqlpyExceptionHandler", "PsqlpySessionContext") @@ -78,7 +80,7 @@ class PsqlpyDriver(AsyncDriverAdapterBase): and transaction management. """ - __slots__ = ("_data_dictionary", "_transaction_active") + __slots__ = ("_data_dictionary", "_pooler_unprepared", "_transaction_active") dialect = "postgres" def __init__( @@ -95,6 +97,8 @@ def __init__( super().__init__(connection=connection, statement_config=statement_config, driver_features=driver_features) self._data_dictionary: PsqlpyDataDictionary | None = None self._transaction_active = False + pooler_mode = self.driver_features.get("pooler_mode", "session") + self._pooler_unprepared = pooler_mode in {"transaction", "statement"} @property def data_dictionary(self) -> "PsqlpyDataDictionary": @@ -138,23 +142,8 @@ async def dispatch_execute(self, cursor: "PsqlpyConnection", statement: SQL) -> ) if statement.operation_type in {"INSERT", "UPDATE", "DELETE"}: - base_sql, comment_suffix = split_sqlcomment(sql) - count_sql = _dml_count_query(base_sql) - if count_sql is not None: - wire_count_sql = f"{count_sql} {comment_suffix}" if comment_suffix else count_sql - count_result = ( - await cursor.fetch(wire_count_sql, params, prepared=False) - if self._is_transaction_pooler_mode() - else await cursor.fetch(wire_count_sql, params) - ) - count_rows, _ = collect_rows(count_result) - if len(count_rows) != 1 or set(count_rows[0]) != {_DML_COUNT_COLUMN}: - msg = "psqlpy DML row count query returned an invalid result" - raise SQLSpecError(msg) - rows_affected = count_rows[0][_DML_COUNT_COLUMN] - if type(rows_affected) is not int or rows_affected < 0: - msg = "psqlpy DML row count query returned an invalid count" - raise SQLSpecError(msg) + rows_affected = await self._fetch_dml_rowcount(cursor, sql, params) + if rows_affected is not None: return self.create_execution_result(cursor, rowcount_override=rows_affected) result = ( @@ -166,6 +155,29 @@ async def dispatch_execute(self, cursor: "PsqlpyConnection", statement: SQL) -> return self.create_execution_result(cursor, rowcount_override=rows_affected) + async def _fetch_dml_rowcount( + self, cursor: "PsqlpyConnection", sql: str, params: "Sequence[Any] | Mapping[str, Any]" + ) -> int | None: + base_sql, comment_suffix = split_sqlcomment(sql) + count_sql = _dml_count_query(base_sql) + if count_sql is None: + return None + wire_count_sql = f"{count_sql} {comment_suffix}" if comment_suffix else count_sql + count_result = ( + await cursor.fetch(wire_count_sql, params, prepared=False) + if self._pooler_unprepared + else await cursor.fetch(wire_count_sql, params) + ) + count_rows, _ = collect_rows(count_result) + if len(count_rows) != 1 or set(count_rows[0]) != {_DML_COUNT_COLUMN}: + msg = "psqlpy DML row count query returned an invalid result" + raise SQLSpecError(msg) + rows_affected = count_rows[0][_DML_COUNT_COLUMN] + if type(rows_affected) is not int or rows_affected < 0: + msg = "psqlpy DML row count query returned an invalid count" + raise SQLSpecError(msg) + return rows_affected + async def dispatch_execute_many(self, cursor: "PsqlpyConnection", statement: SQL) -> "ExecutionResult": """Execute SQL with multiple parameter sets. @@ -207,18 +219,39 @@ async def dispatch_execute_script(self, cursor: "PsqlpyConnection", statement: S statement_config = statement.statement_config statements = self.split_script_statements(sql, statement_config, strip_trailing_semicolon=True) + statement_count = len(statements) successful_count = 0 - last_result = None params = prepared_parameters or [] is_single = len(statements) == 1 + runtime = self._observability + has_observers = runtime is not None and runtime.has_statement_observers + use_prepared = not self._is_transaction_pooler_mode() - for stmt in statements: + for idx, stmt in enumerate(statements): stmt_params = params if (is_single or "$" in stmt) else [] - last_result = await cursor.execute(stmt, stmt_params) + stmt_started = perf_counter() if has_observers else 0.0 + self._pending_statement_profile = None + await cursor.execute(stmt, stmt_params, prepared=use_prepared) + if has_observers: + stmt_duration = perf_counter() - stmt_started + stmt_profile = self._pending_statement_profile + self._pending_statement_profile = None + self._emit_script_statement_event( + sql=stmt, + parameters=stmt_params, + operation="SCRIPT", + rows_affected=None, + duration_s=stmt_duration, + started_at=stmt_started, + profile=stmt_profile, + script_ordinal=idx, + script_statement_count=statement_count, + ) successful_count += 1 + self._pending_statement_profile = None return self.create_execution_result( - last_result, statement_count=len(statements), successful_statements=successful_count, is_script_result=True + cursor, statement_count=statement_count, successful_statements=successful_count, is_script_result=True ) async def begin(self) -> None: @@ -228,32 +261,32 @@ async def begin(self) -> None: try: await self.connection.execute("BEGIN") + except (PsqlpyDatabaseError, PsqlpyConnectionExecuteError, PsqlpyError, Exception) as e: + raise create_mapped_exception(e) from e + try: if self._should_set_local_columnar(): await self.connection.execute("SET LOCAL google_columnar_engine.enable_columnar_scan = 'on'") - except (PsqlpyDatabaseError, PsqlpyConnectionExecuteError, PsqlpyError) as e: - msg = f"Failed to begin psqlpy transaction: {e}" - raise SQLSpecError(msg) from e + except Exception as e: + with suppress(Exception): + await self.connection.execute("ROLLBACK") + raise create_mapped_exception(e) from e self._transaction_active = True async def commit(self) -> None: """Commit the current transaction.""" - try: await self.connection.execute("COMMIT") - except (PsqlpyDatabaseError, PsqlpyConnectionExecuteError, PsqlpyError) as e: - msg = f"Failed to commit psqlpy transaction: {e}" - raise SQLSpecError(msg) from e + except (PsqlpyDatabaseError, PsqlpyConnectionExecuteError, PsqlpyError, Exception) as e: + raise create_mapped_exception(e) from e finally: self._transaction_active = False async def rollback(self) -> None: """Rollback the current transaction.""" - try: await self.connection.execute("ROLLBACK") - except (PsqlpyDatabaseError, PsqlpyConnectionExecuteError, PsqlpyError) as e: - msg = f"Failed to rollback psqlpy transaction: {e}" - raise SQLSpecError(msg) from e + except (PsqlpyDatabaseError, PsqlpyConnectionExecuteError, PsqlpyError, Exception) as e: + raise create_mapped_exception(e) from e finally: self._transaction_active = False @@ -321,9 +354,8 @@ async def select_to_storage( telemetry: "StorageTelemetry | None" = None, **kwargs: Any, ) -> "StorageBridgeJob": - """Execute a query and stream Arrow results to a storage backend.""" + """Execute a query and write client-converted Arrow results to storage.""" - self._require_capability("arrow_export_enabled") arrow_result = await self.select_to_arrow(statement, *parameters, statement_config=statement_config, **kwargs) async_pipeline = self._storage_pipeline() telemetry_payload = await self._write_storage_result( @@ -373,9 +405,8 @@ async def load_from_arrow( qualified = format_table_identifier(table) exc_handler = self.handle_database_exceptions() async with exc_handler, self.with_cursor(self.connection) as cursor: - await cursor.execute(f"TRUNCATE TABLE {qualified}") - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + await self._run_with_exception_handler(exc_handler, cursor.execute, f"TRUNCATE TABLE {qualified}") + self._check_pending_exception(exc_handler) columns, batch_iter = self._coerce_arrow_batches(source) first_batch = next((b for b in batch_iter if b.num_rows > 0), None) @@ -396,7 +427,10 @@ async def load_from_arrow( payload = self._encode_arrow_batch_pg_binary( batch, columns=columns, include_header=True, include_trailer=True ) - await cursor.binary_copy_to_table(payload, table_name, **copy_kwargs) + await self._run_with_exception_handler( + exc_handler, cursor.binary_copy_to_table, payload, table_name, **copy_kwargs + ) + self._check_pending_exception(exc_handler) else: json_columns = await self._resolve_json_columns(schema_name, table_name) deserializer = cast( @@ -407,13 +441,15 @@ async def load_from_arrow( metrics["bytes_processed"] += int(batch.nbytes) records = self._arrow_batch_to_rows(batch, columns) if records: - await cursor.copy_records_to_table( + await self._run_with_exception_handler( + exc_handler, + cursor.copy_records_to_table, table_name, coerce_json_columns(records, columns, json_columns, deserializer=deserializer), **copy_kwargs, ) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + self._check_pending_exception(exc_handler) + self._check_pending_exception(exc_handler) telemetry_payload = self._ingest_telemetry_from_batches(metrics["rows_processed"], metrics["bytes_processed"]) telemetry_payload["destination"] = table @@ -526,20 +562,47 @@ async def _resolve_json_columns(self, schema_name: "str | None", table_name: str data, _ = collect_rows(rows) return {str(row["column_name"]) for row in data} - async def _execute_cache_hit( - self, sql: str, params: "tuple[Any, ...] | list[Any] | dict[str, Any]", cached: Any + async def _execute_cache_hit_with_cursor( + self, sql: str, params: "tuple[Any, ...] | list[Any] | dict[str, Any]", cached: "CachedQuery" ) -> "SQLResult": """Execute cached psqlpy queries with cast-aware parameter preparation.""" prepared_params = self.prepare_driver_parameters(params, self.statement_config, prepared_statement=cached) - direct_statement = self._cached_statement( - sql, - params, - cached, - cast("tuple[Any, ...] | list[Any] | dict[str, Any]", prepared_params), - params_are_simple=True, - compiled_sql=cached.compiled_sql, + args = ( + prepared_params + if isinstance(prepared_params, (tuple, list)) + else tuple(prepared_params.values()) + if isinstance(prepared_params, dict) + else () + ) + wire_sql = self._format_wire_sql(cached.compiled_sql) + if cached.operation_profile.returns_rows: + query_result = ( + await self.connection.fetch(wire_sql, args, prepared=False) + if self._pooler_unprepared + else await self.connection.fetch(wire_sql, args) + ) + dict_rows, column_names = collect_rows(query_result) + return self.build_cached_select_result( + sql, + params, + cached, + selected_data=dict_rows, + column_names=column_names, + data_row_count=len(dict_rows), + row_format="dict", + compiled_sql=wire_sql, + ) + if cached.operation_type in {"INSERT", "UPDATE", "DELETE"}: + rows_affected = await self._fetch_dml_rowcount(self.connection, wire_sql, args) + if rows_affected is not None: + return DMLResult(cached.operation_type, rows_affected) + result = ( + await self.connection.execute(wire_sql, args, prepared=False) + if self._pooler_unprepared + else await self.connection.execute(wire_sql, args) ) - return await self._execute_cached_statement(direct_statement) + affected_rows = extract_rows_affected(result) + return DMLResult(cached.operation_type, affected_rows) def _connection_in_transaction(self) -> bool: """Check if connection is in transaction. @@ -550,10 +613,7 @@ def _connection_in_transaction(self) -> bool: return self._transaction_active def _is_transaction_pooler_mode(self) -> bool: - pooler_mode = self.driver_features.get("pooler_mode", "session") - return pooler_mode in {"transaction", "statement"} or bool( - self.driver_features.get("pgbouncer") or self.driver_features.get("pgbouncer_mode") - ) + return self._pooler_unprepared def _should_set_local_columnar(self) -> bool: return bool(self.driver_features.get("enable_alloydb_columnar", False)) and self._is_transaction_pooler_mode() diff --git a/sqlspec/adapters/psqlpy/events/_hub.py b/sqlspec/adapters/psqlpy/events/_hub.py index 20da69998..e8c767ed3 100644 --- a/sqlspec/adapters/psqlpy/events/_hub.py +++ b/sqlspec/adapters/psqlpy/events/_hub.py @@ -13,6 +13,7 @@ import asyncio import contextlib import logging +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, cast from weakref import WeakKeyDictionary @@ -237,16 +238,9 @@ def _resolve_listener_config(self) -> "PsqlpyConfig": return self._listener_config raw_features = getattr(self._config, "driver_features", None) driver_features = raw_features if isinstance(raw_features, dict) else {} - pooler_mode = str( - driver_features.get("pooler_mode") - or ( - "transaction" - if driver_features.get("pgbouncer") or driver_features.get("pgbouncer_mode") - else "session" - ) - ).lower() + pooler_mode = str(driver_features.get("pooler_mode") or "session").lower() raw_direct_dsn = driver_features.get("direct_dsn") - direct_dsn = raw_direct_dsn if isinstance(raw_direct_dsn, str) and raw_direct_dsn else None + direct_dsn = raw_direct_dsn if isinstance(raw_direct_dsn, (str, Mapping)) and raw_direct_dsn else None if pooler_mode in {"transaction", "statement"} and not direct_dsn: logger.warning( "PostgreSQL LISTEN/NOTIFY cannot run under pooler_mode=%r without configuring driver_features['direct_dsn']", diff --git a/sqlspec/adapters/psqlpy/events/backend.py b/sqlspec/adapters/psqlpy/events/backend.py index cf49e89a4..2d305bc4e 100644 --- a/sqlspec/adapters/psqlpy/events/backend.py +++ b/sqlspec/adapters/psqlpy/events/backend.py @@ -28,7 +28,7 @@ __all__ = ("PsqlpyEventsBackend", "PsqlpyHybridEventsBackend", "create_event_backend") -logger = get_logger("sqlspec.events.psqlpy") +logger = get_logger("sqlspec.adapters.psqlpy.events.backend") _MIN_LISTENER_POOL_SIZE = 2 _MAX_SEEN_MARKERS = 1_024 @@ -241,14 +241,15 @@ async def dequeue(self, channel: str, poll_interval: float) -> EventMessage | No return await self._dequeue_reconciled(channel) self._queue._reset_empty_poll_delay(channel) self._runtime.increment_metric("events.listener.wakeup") - batch_size = _extract_batch_size(payload) - is_new, should_drain = self._marker_state.register(channel, _extract_marker_id(payload), batch_size) + parsed_payload = _parse_notify_payload(payload) + batch_size = _extract_batch_size(parsed_payload) + is_new, should_drain = self._marker_state.register(channel, _extract_marker_id(parsed_payload), batch_size) if is_new and batch_size > 1: self._runtime.increment_metric("events.marker.coalesced", batch_size - 1) if not should_drain: self._runtime.increment_metric("events.marker.duplicate") continue - event_id = _extract_event_id(payload) + event_id = _extract_event_id(parsed_payload) if event_id is not None: event = await self._queue.dequeue_by_event_id(event_id) if event is not None: @@ -354,6 +355,10 @@ def create_event_backend( def _validate_listener_pool_capacity(config: "PsqlpyConfig") -> None: + raw_features = config.driver_features + driver_features = raw_features if isinstance(raw_features, dict) else {} + if driver_features.get("direct_dsn"): + return max_size = config.connection_config.get("max_db_pool_size") if max_size is not None and max_size < _MIN_LISTENER_POOL_SIZE: msg = ( @@ -362,34 +367,40 @@ def _validate_listener_pool_capacity(config: "PsqlpyConfig") -> None: raise ImproperConfigurationError(msg) -def _extract_event_id(payload: "str | None") -> "str | None": - if not payload: +def _parse_notify_payload(payload: "dict[str, Any] | str | None") -> "dict[str, Any] | None": + if isinstance(payload, dict): + return payload + if not payload or not isinstance(payload, str): + return None + try: + raw = from_json(payload) + return raw if isinstance(raw, dict) else None + except (TypeError, ValueError): + return None + + +def _extract_event_id(payload: "dict[str, Any] | str | None") -> "str | None": + parsed = payload if isinstance(payload, dict) else _parse_notify_payload(payload) + if parsed is None: return None - raw = from_json(payload) - if isinstance(raw, dict): - event_id = raw.get("event_id") - return event_id if isinstance(event_id, str) else None - return None + event_id = parsed.get("event_id") + return event_id if isinstance(event_id, str) else None -def _extract_batch_size(payload: "str | None") -> int: - if not payload: +def _extract_batch_size(payload: "dict[str, Any] | str | None") -> int: + parsed = payload if isinstance(payload, dict) else _parse_notify_payload(payload) + if parsed is None: return 0 - raw = from_json(payload) - if isinstance(raw, dict): - batch_size = raw.get("batch_size") - return batch_size if isinstance(batch_size, int) else 0 - return 0 + batch_size = parsed.get("batch_size") + return batch_size if isinstance(batch_size, int) else 0 -def _extract_marker_id(payload: "str | None") -> "str | None": - if not payload: +def _extract_marker_id(payload: "dict[str, Any] | str | None") -> "str | None": + parsed = payload if isinstance(payload, dict) else _parse_notify_payload(payload) + if parsed is None: return None - raw = from_json(payload) - if isinstance(raw, dict): - marker_id = raw.get("marker_id", raw.get("event_id")) - return marker_id if isinstance(marker_id, str) else None - return None + marker_id = parsed.get("marker_id", parsed.get("event_id")) + return marker_id if isinstance(marker_id, str) else None def _record_dequeue_result(runtime: Any, event: "EventMessage | None") -> None: diff --git a/sqlspec/adapters/psqlpy/events/store.py b/sqlspec/adapters/psqlpy/events/store.py index db99cc248..5bdf5494f 100644 --- a/sqlspec/adapters/psqlpy/events/store.py +++ b/sqlspec/adapters/psqlpy/events/store.py @@ -3,6 +3,7 @@ from typing import NotRequired from sqlspec.adapters.psqlpy.config import PsqlpyConfig +from sqlspec.adapters.psqlpy.core import build_postgres_storage_params_sql from sqlspec.config import EventsConfig from sqlspec.extensions.events import BaseEventQueueStore @@ -42,27 +43,4 @@ def _column_types(self) -> "tuple[str, str, str]": def _table_clause(self) -> str: """Return explicitly configured PostgreSQL table tuning options.""" - return _postgres_table_options(self.settings) - - -def _postgres_table_options(settings: "dict[str, object]") -> str: - options = _postgres_option_values(settings) - return f" WITH ({', '.join(options)})" if options else "" - - -def _postgres_option_values(settings: "dict[str, object]") -> "list[str]": - options: list[str] = [] - fillfactor = settings.get("fillfactor") - if fillfactor is not None: - if not isinstance(fillfactor, int) or isinstance(fillfactor, bool) or fillfactor not in range(10, 101): - msg = "extension_config['events']['fillfactor'] must be an integer from 10 to 100" - raise ValueError(msg) - options.append(f"fillfactor = {fillfactor}") - for key in ("autovacuum_vacuum_scale_factor", "autovacuum_analyze_scale_factor"): - value = settings.get(key) - if value is not None: - if not isinstance(value, (int, float)) or isinstance(value, bool) or not 0 <= float(value) <= 1: - msg = f"extension_config['events']['{key}'] must be a number from 0 to 1" - raise ValueError(msg) - options.append(f"{key} = {float(value):g}") - return options + return build_postgres_storage_params_sql(self.settings, include_autovacuum=True, section="events") diff --git a/sqlspec/adapters/psqlpy/litestar/store.py b/sqlspec/adapters/psqlpy/litestar/store.py index c4b764fa0..8a1df7a26 100644 --- a/sqlspec/adapters/psqlpy/litestar/store.py +++ b/sqlspec/adapters/psqlpy/litestar/store.py @@ -1,8 +1,9 @@ """Psqlpy session store for Litestar integration.""" from datetime import UTC, datetime, timedelta -from typing import TYPE_CHECKING, Any, NotRequired, cast +from typing import TYPE_CHECKING, NotRequired +from sqlspec.adapters.psqlpy.core import build_postgres_storage_params_sql from sqlspec.config import LitestarConfig from sqlspec.extensions.litestar.store import BaseSQLSpecStore @@ -169,7 +170,7 @@ async def exists(self, key: str) -> bool: async with self._config.provide_connection() as conn: query_result = await conn.fetch(sql, [key]) - rows = query_result.result() + rows = query_result.result() if query_result else [] return len(rows) > 0 async def expires_in(self, key: str) -> "int | None": @@ -188,7 +189,7 @@ async def expires_in(self, key: str) -> "int | None": async with self._config.provide_connection() as conn: query_result = await conn.fetch(sql, [key]) - rows = query_result.result() + rows = query_result.result() if query_result else [] if not rows: return None @@ -219,7 +220,7 @@ async def delete_expired(self) -> int: async with self._config.provide_connection() as conn: query_result = await conn.fetch(sql, []) - rows = query_result.result() + rows = query_result.result() if query_result else [] count = len(rows) if count > 0: self._log_delete_expired(count) @@ -231,7 +232,14 @@ def _table_ddl(self) -> str: Returns: SQL statement to create the sessions table with proper indexes. """ - fillfactor, vacuum_scale, analyze_scale = _postgres_litestar_tuning(self._config) + tuning_clause = build_postgres_storage_params_sql( + self._config.extension_config.get("litestar", {}), + default_fillfactor=80, + default_vacuum_scale=0.05, + default_analyze_scale=0.02, + include_autovacuum=True, + section="litestar", + ) return f""" CREATE TABLE IF NOT EXISTS {self._table_name} ( session_id TEXT PRIMARY KEY, @@ -239,15 +247,10 @@ def _table_ddl(self) -> str: expires_at TIMESTAMPTZ, created_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP, updated_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP - ) WITH (fillfactor = {fillfactor}); + ){tuning_clause}; CREATE INDEX IF NOT EXISTS idx_{self._table_name}_expires_at ON {self._table_name}(expires_at) WHERE expires_at IS NOT NULL; - - ALTER TABLE {self._table_name} SET ( - autovacuum_vacuum_scale_factor = {vacuum_scale:g}, - autovacuum_analyze_scale_factor = {analyze_scale:g} - ); """ def _drop_table_sql(self) -> "list[str]": @@ -257,19 +260,3 @@ def _drop_table_sql(self) -> "list[str]": List of SQL statements to drop indexes and table. """ return [f"DROP INDEX IF EXISTS idx_{self._table_name}_expires_at", f"DROP TABLE IF EXISTS {self._table_name}"] - - -def _postgres_litestar_tuning(config: Any) -> "tuple[int, float, float]": - settings = cast("dict[str, Any]", config.extension_config.get("litestar", {})) - fillfactor = settings.get("fillfactor", 80) - if not isinstance(fillfactor, int) or isinstance(fillfactor, bool) or fillfactor not in range(10, 101): - msg = "extension_config['litestar']['fillfactor'] must be an integer from 10 to 100" - raise ValueError(msg) - scales: list[float] = [] - for key, default in (("autovacuum_vacuum_scale_factor", 0.05), ("autovacuum_analyze_scale_factor", 0.02)): - value = settings.get(key, default) - if not isinstance(value, (int, float)) or isinstance(value, bool) or not 0 <= float(value) <= 1: - msg = f"extension_config['litestar']['{key}'] must be a number from 0 to 1" - raise ValueError(msg) - scales.append(float(value)) - return fillfactor, scales[0], scales[1] diff --git a/sqlspec/adapters/psqlpy/type_converter.py b/sqlspec/adapters/psqlpy/type_converter.py index 7d3941837..1aa253efc 100644 --- a/sqlspec/adapters/psqlpy/type_converter.py +++ b/sqlspec/adapters/psqlpy/type_converter.py @@ -1,42 +1,21 @@ -"""PostgreSQL-specific helpers for the psqlpy adapter. +"""PostgreSQL-specific helpers for the psqlpy adapter.""" -This module preserves the ``register_pgvector`` placeholder used by the -driver configuration layer. -""" +from typing import Any -from typing import TYPE_CHECKING, Any +from psqlpy.extra_types import PgVector -from sqlspec.typing import PGVECTOR_INSTALLED, import_optional_attr - -if TYPE_CHECKING: - from sqlspec.adapters.psqlpy._typing import PsqlpyConnection as Connection - -__all__ = ("coerce_pgvector", "register_pgvector") - -_PGVECTOR_TYPE = import_optional_attr("psqlpy.extra_types", "PgVector") +__all__ = ("coerce_pgvector",) def coerce_pgvector(value: Any) -> Any: """Wrap supported dense vector values using psqlpy's native vector encoder.""" - if value is None or _PGVECTOR_TYPE is None or isinstance(value, _PGVECTOR_TYPE): + if value is None or isinstance(value, PgVector): return value if isinstance(value, (list, tuple)): - return _PGVECTOR_TYPE(list(value)) + return PgVector(list(value)) tolist = getattr(value, "tolist", None) if callable(tolist): - return _PGVECTOR_TYPE(tolist()) + raw_list = tolist() + if isinstance(raw_list, list): + return PgVector(raw_list) return value - - -def register_pgvector(connection: "Connection") -> None: - """Register pgvector type handlers on psqlpy connection. - - Currently a placeholder for future implementation. The psqlpy library - does not yet expose a type handler registration API compatible with - pgvector's automatic conversion system. - - Args: - connection: Psqlpy connection instance. - """ - if not PGVECTOR_INSTALLED: - return diff --git a/sqlspec/adapters/psycopg/_typing.py b/sqlspec/adapters/psycopg/_typing.py index 9a74fe490..63d3fe484 100644 --- a/sqlspec/adapters/psycopg/_typing.py +++ b/sqlspec/adapters/psycopg/_typing.py @@ -46,10 +46,10 @@ from sqlspec.builder import QueryBuilder from sqlspec.core import SQL, Statement, StatementConfig - PsycopgSyncConnection: TypeAlias = PsycopgConnection[PsycopgDictRow] - PsycopgAsyncConnection: TypeAlias = PsycopgNativeAsyncConnection[PsycopgDictRow] - PsycopgSyncRawCursor: TypeAlias = PsycopgCursor[PsycopgDictRow] - PsycopgAsyncRawCursor: TypeAlias = PsycopgNativeAsyncCursor[PsycopgDictRow] + PsycopgSyncConnection: TypeAlias = PsycopgConnection[Any] + PsycopgAsyncConnection: TypeAlias = PsycopgNativeAsyncConnection[Any] + PsycopgSyncRawCursor: TypeAlias = PsycopgCursor[Any] + PsycopgAsyncRawCursor: TypeAlias = PsycopgNativeAsyncCursor[Any] else: PsycopgSyncConnection = PsycopgConnection PsycopgAsyncConnection = PsycopgNativeAsyncConnection @@ -215,7 +215,10 @@ def __exit__( self, exc_type: "type[BaseException] | None", exc_val: "BaseException | None", exc_tb: "TracebackType | None" ) -> "bool | None": try: - if self._driver is not None and getattr(self._driver, "_transaction_active", False): + if self._driver is not None and ( + getattr(self._driver, "_transaction_active", False) + or getattr(self._driver, "_connection_in_transaction", lambda: False)() + ): with contextlib.suppress(Exception): self._driver.rollback() finally: @@ -276,7 +279,10 @@ async def __aexit__( self, exc_type: "type[BaseException] | None", exc_val: "BaseException | None", exc_tb: "TracebackType | None" ) -> "bool | None": try: - if self._driver is not None and getattr(self._driver, "_transaction_active", False): + if self._driver is not None and ( + getattr(self._driver, "_transaction_active", False) + or getattr(self._driver, "_connection_in_transaction", lambda: False)() + ): with contextlib.suppress(Exception): await self._driver.rollback() finally: diff --git a/sqlspec/adapters/psycopg/adk/store.py b/sqlspec/adapters/psycopg/adk/store.py index d2cdcba31..ee47a037e 100644 --- a/sqlspec/adapters/psycopg/adk/store.py +++ b/sqlspec/adapters/psycopg/adk/store.py @@ -7,6 +7,7 @@ from sqlspec.adapters.psycopg._typing import psycopg_dict_row as dict_row from sqlspec.adapters.psycopg._typing import psycopg_errors as errors from sqlspec.adapters.psycopg._typing import psycopg_sql as pg_sql +from sqlspec.adapters.psycopg.core import build_postgres_storage_params_sql from sqlspec.config import ADKConfig from sqlspec.extensions.adk import ( BaseAsyncADKStore, @@ -706,7 +707,7 @@ async def _sessions_table_ddl(self) -> str: self._session_table, self._session_table, self._session_table, - _postgres_table_options(_adk_config(self._config)), + build_postgres_storage_params_sql(_adk_config(self._config), default_fillfactor=80, section="adk"), ) async def _events_table_ddl(self) -> str: @@ -723,19 +724,23 @@ async def _events_table_ddl(self) -> str: self._events_table, covering_columns, generated_indexes, - _postgres_table_options(adk_config, include_autovacuum=True), + build_postgres_storage_params_sql( + adk_config, include_autovacuum=True, default_fillfactor=80, section="adk" + ), self._events_table, self._events_table, ) async def _app_states_table_ddl(self) -> str: return _ADK_APP_STATES_TABLE_DDL_TEMPLATE.format( - self._app_state_table, _postgres_table_options(_adk_config(self._config)) + self._app_state_table, + build_postgres_storage_params_sql(_adk_config(self._config), default_fillfactor=80, section="adk"), ) async def _user_states_table_ddl(self) -> str: return _ADK_USER_STATES_TABLE_DDL_TEMPLATE.format( - self._user_state_table, _postgres_table_options(_adk_config(self._config)) + self._user_state_table, + build_postgres_storage_params_sql(_adk_config(self._config), default_fillfactor=80, section="adk"), ) async def _metadata_table_ddl(self) -> str: @@ -1256,7 +1261,7 @@ def _sessions_table_ddl(self) -> str: self._session_table, self._session_table, self._session_table, - _postgres_table_options(_adk_config(self._config)), + build_postgres_storage_params_sql(_adk_config(self._config), default_fillfactor=80, section="adk"), ) def _events_table_ddl(self) -> str: @@ -1273,19 +1278,23 @@ def _events_table_ddl(self) -> str: self._events_table, covering_columns, generated_indexes, - _postgres_table_options(adk_config, include_autovacuum=True), + build_postgres_storage_params_sql( + adk_config, include_autovacuum=True, default_fillfactor=80, section="adk" + ), self._events_table, self._events_table, ) def _app_states_table_ddl(self) -> str: return _ADK_APP_STATES_TABLE_DDL_TEMPLATE.format( - self._app_state_table, _postgres_table_options(_adk_config(self._config)) + self._app_state_table, + build_postgres_storage_params_sql(_adk_config(self._config), default_fillfactor=80, section="adk"), ) def _user_states_table_ddl(self) -> str: return _ADK_USER_STATES_TABLE_DDL_TEMPLATE.format( - self._user_state_table, _postgres_table_options(_adk_config(self._config)) + self._user_state_table, + build_postgres_storage_params_sql(_adk_config(self._config), default_fillfactor=80, section="adk"), ) def _metadata_table_ddl(self) -> str: @@ -2022,33 +2031,6 @@ def _adk_config(config: Any) -> PsycopgADKConfig: return cast("PsycopgADKConfig", adk_config) -def _postgres_table_options(adk_config: PsycopgADKConfig, *, include_autovacuum: bool = False) -> str: - options = [_postgres_fillfactor_option(adk_config)] - if include_autovacuum: - options.extend(_postgres_autovacuum_options(adk_config)) - return f" WITH ({', '.join(options)})" - - -def _postgres_fillfactor_option(adk_config: PsycopgADKConfig) -> str: - value = adk_config.get("fillfactor", 80) - if not isinstance(value, int) or isinstance(value, bool) or value not in range(10, 101): - msg = "extension_config['adk']['fillfactor'] must be an integer from 10 to 100" - raise ValueError(msg) - return f"fillfactor = {value}" - - -def _postgres_autovacuum_options(adk_config: PsycopgADKConfig) -> "list[str]": - options: list[str] = [] - for key in ("autovacuum_vacuum_scale_factor", "autovacuum_analyze_scale_factor"): - value = adk_config.get(key) - if value is not None: - if not isinstance(value, (int, float)) or isinstance(value, bool) or not 0 <= float(value) <= 1: - msg = f"extension_config['adk']['{key}'] must be a number from 0 to 1" - raise ValueError(msg) - options.append(f"{key} = {float(value):g}") - return options - - def _postgres_memory_indexes( memory_table: str, use_fts: bool, diff --git a/sqlspec/adapters/psycopg/config.py b/sqlspec/adapters/psycopg/config.py index 80813c2cd..44c324f1b 100644 --- a/sqlspec/adapters/psycopg/config.py +++ b/sqlspec/adapters/psycopg/config.py @@ -1,6 +1,8 @@ """Psycopg database configuration with direct field-based configuration.""" +import asyncio import inspect +import threading from collections.abc import Mapping from contextlib import suppress from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, NotRequired, Self, TypedDict, cast @@ -94,7 +96,6 @@ "min_size", "name", "num_workers", - "null_pool", "open", "reconnect_failed", "reconnect_timeout", @@ -188,8 +189,6 @@ class PsycopgPoolParams(PsycopgConnectionParams): """Callback invoked when pool reconnection attempts exceed ``reconnect_timeout``.""" kwargs: NotRequired["dict[str, Any]"] """Explicit connection keyword mapping passed to ``psycopg.connect``.""" - null_pool: NotRequired[bool] - """Use ``NullConnectionPool`` / ``AsyncNullConnectionPool`` so idle connections are never retained.""" class PsycopgDriverFeatures(TypedDict): @@ -221,16 +220,14 @@ class PsycopgDriverFeatures(TypedDict): """AlloyDB network path (``PUBLIC``, ``PRIVATE``, or ``PSC``; defaults to ``PRIVATE``).""" enable_alloydb_columnar: NotRequired[bool] """Enable ``google_columnar_engine.enable_columnar_scan=on`` via startup ``options`` or per-transaction ``SET LOCAL``.""" + enable_pgcopy_binary: NotRequired[bool] + """Enable binary streaming format for Arrow ingestion via PGCOPY binary protocol.""" pooler_mode: NotRequired[Literal["session", "transaction", "statement", "auto"]] """External pooler compatibility mode; sets ``prepare_threshold=None`` in ``transaction`` or ``statement`` mode.""" - pgbouncer: NotRequired[bool] - """Alias for ``pooler_mode='transaction'`` disabling server-side prepared statements.""" - pgbouncer_mode: NotRequired[bool] - """Alias for ``pooler_mode='transaction'`` disabling server-side prepared statements.""" direct_dsn: NotRequired["str | Mapping[str, Any] | None"] """Direct PostgreSQL DSN or connection mapping bypassing the pooler for migrations and ``LISTEN/NOTIFY``.""" - null_pool: NotRequired[bool] - """Enable ``NullConnectionPool`` / ``AsyncNullConnectionPool`` for serverless or external pooler deployments.""" + enable_pooling: NotRequired[bool] + """Enable connection pooling (defaults to ``True``; ``False`` uses ``NullConnectionPool`` / ``AsyncNullConnectionPool``).""" def build_connection_config(connection_config: "PsycopgPoolParams | Mapping[str, Any] | None") -> dict[str, Any]: @@ -347,18 +344,14 @@ def __init__( statement_config, driver_features = apply_driver_features(statement_config, driver_features) features_dict = dict(driver_features) if driver_features else {} - self._user_connection_hook: Callable[[PsycopgSyncConnection], None] | None = features_dict.get( - "on_connection_create" + self._user_connection_hook: Callable[[PsycopgSyncConnection], None] | None = features_dict.pop( + "on_connection_create", None ) features_dict.setdefault("enable_alloydb", False) features_dict.setdefault("enable_alloydb_iam_auth", False) features_dict.setdefault("alloydb_ip_type", "PRIVATE") pooler_mode = features_dict.get("pooler_mode", "session") - is_transaction_pooler = bool( - pooler_mode in {"transaction", "statement"} - or features_dict.get("pgbouncer") - or features_dict.get("pgbouncer_mode") - ) + is_transaction_pooler = pooler_mode in {"transaction", "statement"} if is_transaction_pooler: connection_config.setdefault("prepare_threshold", None) if features_dict.get("enable_alloydb_columnar") and not is_transaction_pooler: @@ -368,6 +361,7 @@ def __init__( connection_config["options"] = columnar_opt elif columnar_opt not in str(existing_options): connection_config["options"] = f"{existing_options} {columnar_opt}" + self._probe_lock: threading.Lock = threading.Lock() self._pgvector_available: bool | None = None self._paradedb_available: bool | None = None self._pg_textsearch_available: bool | None = None @@ -402,18 +396,14 @@ def derive_direct_config(self) -> Self: conn_cfg.update(build_connection_config(direct_dsn)) else: conn_cfg["conninfo"] = str(direct_dsn) - is_pooled = bool( - self.driver_features.get("pooler_mode") in {"transaction", "statement"} - or self.driver_features.get("pgbouncer") - or self.driver_features.get("pgbouncer_mode") - ) + is_pooled = self.driver_features.get("pooler_mode") in {"transaction", "statement"} if is_pooled and conn_cfg.get("prepare_threshold") is None: conn_cfg.pop("prepare_threshold", None) features = dict(self.driver_features) features["pooler_mode"] = "session" - features["pgbouncer"] = False - features["pgbouncer_mode"] = False features["direct_dsn"] = None + if self._user_connection_hook is not None: + features["on_connection_create"] = self._user_connection_hook return type(self)( connection_config=conn_cfg, migration_config=dict(self.migration_config) if self.migration_config else None, @@ -559,7 +549,6 @@ def _setup_alloydb_connector( def _create_pool(self) -> "ConnectionPool": """Create the actual connection pool.""" all_config = dict(self.connection_config) - all_config.pop("null_pool", None) auth_token_provider = all_config.pop("auth_token_provider", None) default_max_lifetime = _AUTH_TOKEN_DEFAULT_MAX_LIFETIME if auth_token_provider is not None else 3600.0 @@ -612,9 +601,9 @@ def _chained_configure(conn: "PsycopgSyncConnection") -> None: self._setup_alloydb_connector(all_config, pool_parameters) conninfo = None - is_null_pool = bool(self.connection_config.get("null_pool") or self.driver_features.get("null_pool")) - pool_cls = NullConnectionPool if is_null_pool else ConnectionPool - if is_null_pool: + pooling_disabled = not bool(self.driver_features.get("enable_pooling", True)) + pool_cls = NullConnectionPool if pooling_disabled else ConnectionPool + if pooling_disabled: pool_parameters.pop("min_size", None) if conninfo: @@ -636,38 +625,51 @@ def _configure_connection(self, conn: "PsycopgSyncConnection") -> None: set_json_loads(deserializer, conn) if self._pgvector_available is None: - detected_extensions: set[str] = set() - extensions = build_postgres_extension_probe_names(self.driver_features) - probe_succeeded = True - if extensions: - try: - cursor = conn.execute( - "SELECT extname FROM pg_extension WHERE extname = ANY(%s::text[])", (extensions,) - ) - results = cursor.fetchall() - detected_extensions = { - cast("str", r["extname"] if isinstance(r, Mapping) else r[0]) for r in results - } - except Exception as exc: - probe_succeeded = False - self._pg_textsearch_probe_error = exc + with self._probe_lock: + if self._pgvector_available is None: + detected_extensions: set[str] = set() + extensions = build_postgres_extension_probe_names(self.driver_features) + probe_succeeded = True + if extensions: + try: + cursor = conn.execute( + "SELECT extname FROM pg_extension WHERE extname = ANY(%s::text[])", (extensions,) + ) + results = cursor.fetchall() + detected_extensions = { + cast("str", r["extname"] if isinstance(r, Mapping) else r[0]) for r in results + } + except Exception as exc: + probe_succeeded = False + self._pg_textsearch_probe_error = exc + if not conn.autocommit: + with suppress(Exception): + conn.rollback() + if probe_succeeded: + self.statement_config, self._pgvector_available, self._paradedb_available = ( + resolve_postgres_extension_state( + self.statement_config, self.driver_features, detected_extensions + ) + ) + self._pg_textsearch_available = is_postgres_extension_active( + self.driver_features, "pg_textsearch" + ) if not conn.autocommit: - with suppress(Exception): - conn.rollback() - if probe_succeeded: - self.statement_config, self._pgvector_available, self._paradedb_available = ( - resolve_postgres_extension_state(self.statement_config, self.driver_features, detected_extensions) - ) - self._pg_textsearch_available = is_postgres_extension_active(self.driver_features, "pg_textsearch") + conn.rollback() if self._pgvector_available: register_pgvector_sync(conn) if self._user_connection_hook is not None: - self._user_connection_hook(conn) - - if not conn.autocommit: - conn.rollback() + try: + self._user_connection_hook(conn) + if not conn.autocommit: + conn.commit() + except BaseException: + if not conn.autocommit: + with suppress(Exception): + conn.rollback() + raise def _ensure_pg_textsearch_available(self) -> None: if self._pg_textsearch_available: @@ -811,15 +813,11 @@ def __init__( if features_dict.get("enable_alloydb", False): msg = "Google AlloyDB connector integration (enable_alloydb) is only supported on PsycopgSyncConfig" raise ImproperConfigurationError(msg) - self._user_connection_hook: Callable[[PsycopgAsyncConnection], Awaitable[None]] | None = features_dict.get( - "on_connection_create" + self._user_connection_hook: Callable[[PsycopgAsyncConnection], Awaitable[None]] | None = features_dict.pop( + "on_connection_create", None ) pooler_mode = features_dict.get("pooler_mode", "session") - is_transaction_pooler = bool( - pooler_mode in {"transaction", "statement"} - or features_dict.get("pgbouncer") - or features_dict.get("pgbouncer_mode") - ) + is_transaction_pooler = pooler_mode in {"transaction", "statement"} if is_transaction_pooler: connection_config.setdefault("prepare_threshold", None) if features_dict.get("enable_alloydb_columnar") and not is_transaction_pooler: @@ -829,6 +827,7 @@ def __init__( connection_config["options"] = columnar_opt elif columnar_opt not in str(existing_options): connection_config["options"] = f"{existing_options} {columnar_opt}" + self._probe_lock: asyncio.Lock | None = None self._pgvector_available: bool | None = None self._paradedb_available: bool | None = None self._pg_textsearch_available: bool | None = None @@ -860,18 +859,14 @@ def derive_direct_config(self) -> Self: conn_cfg.update(build_connection_config(direct_dsn)) else: conn_cfg["conninfo"] = str(direct_dsn) - is_pooled = bool( - self.driver_features.get("pooler_mode") in {"transaction", "statement"} - or self.driver_features.get("pgbouncer") - or self.driver_features.get("pgbouncer_mode") - ) + is_pooled = self.driver_features.get("pooler_mode") in {"transaction", "statement"} if is_pooled and conn_cfg.get("prepare_threshold") is None: conn_cfg.pop("prepare_threshold", None) features = dict(self.driver_features) features["pooler_mode"] = "session" - features["pgbouncer"] = False - features["pgbouncer_mode"] = False features["direct_dsn"] = None + if self._user_connection_hook is not None: + features["on_connection_create"] = self._user_connection_hook return type(self)( connection_config=conn_cfg, migration_config=dict(self.migration_config) if self.migration_config else None, @@ -881,7 +876,7 @@ def derive_direct_config(self) -> Self: extension_config=dict(self.extension_config) if self.extension_config else None, ) - async def create_connection(self) -> "PsycopgAsyncConnection": # pyright: ignore + async def create_connection(self) -> "PsycopgAsyncConnection": """Open a standalone connection owned by the caller. The connection is not bound to the pool, so closing it actually closes @@ -971,7 +966,6 @@ async def _create_pool(self) -> "AsyncConnectionPool": """Create the actual async connection pool.""" all_config = dict(self.connection_config) - all_config.pop("null_pool", None) auth_token_provider = all_config.pop("auth_token_provider", None) default_max_lifetime = _AUTH_TOKEN_DEFAULT_MAX_LIFETIME if auth_token_provider is not None else 3600.0 @@ -1018,9 +1012,9 @@ async def _chained_async_configure(conn: "PsycopgAsyncConnection") -> None: kwargs = all_config.pop("kwargs", {}) all_config.update(kwargs) - is_null_pool = bool(self.connection_config.get("null_pool") or self.driver_features.get("null_pool")) - pool_cls = AsyncNullConnectionPool if is_null_pool else AsyncConnectionPool - if is_null_pool: + pooling_disabled = not bool(self.driver_features.get("enable_pooling", True)) + pool_cls = AsyncNullConnectionPool if pooling_disabled else AsyncConnectionPool + if pooling_disabled: pool_parameters.pop("min_size", None) if conninfo: @@ -1045,38 +1039,55 @@ async def _configure_async_connection(self, conn: "PsycopgAsyncConnection") -> N set_json_loads(deserializer, conn) if self._pgvector_available is None: - detected_extensions: set[str] = set() - extensions = build_postgres_extension_probe_names(self.driver_features) - probe_succeeded = True - if extensions: - try: - cursor = await conn.execute( - "SELECT extname FROM pg_extension WHERE extname = ANY(%s::text[])", (extensions,) - ) - results = await cursor.fetchall() - detected_extensions = { - cast("str", r["extname"] if isinstance(r, Mapping) else r[0]) for r in results - } - except Exception as exc: - probe_succeeded = False - self._pg_textsearch_probe_error = exc + if self._probe_lock is None: + self._probe_lock = asyncio.Lock() + async with self._probe_lock: + if self._pgvector_available is None: + detected_extensions: set[str] = set() + extensions = build_postgres_extension_probe_names(self.driver_features) + probe_succeeded = True + if extensions: + try: + cursor = await conn.execute( + "SELECT extname FROM pg_extension WHERE extname = ANY(%s::text[])", (extensions,) + ) + results = await cursor.fetchall() + detected_extensions = { + cast("str", r["extname"] if isinstance(r, Mapping) else r[0]) for r in results + } + except Exception as exc: + probe_succeeded = False + self._pg_textsearch_probe_error = exc + if not conn.autocommit: + with suppress(Exception): + await conn.rollback() + if probe_succeeded: + self.statement_config, self._pgvector_available, self._paradedb_available = ( + resolve_postgres_extension_state( + self.statement_config, self.driver_features, detected_extensions + ) + ) + self._pg_textsearch_available = is_postgres_extension_active( + self.driver_features, "pg_textsearch" + ) if not conn.autocommit: - with suppress(Exception): - await conn.rollback() - if probe_succeeded: - self.statement_config, self._pgvector_available, self._paradedb_available = ( - resolve_postgres_extension_state(self.statement_config, self.driver_features, detected_extensions) - ) - self._pg_textsearch_available = is_postgres_extension_active(self.driver_features, "pg_textsearch") + await conn.rollback() if self._pgvector_available: await register_pgvector_async(conn) if self._user_connection_hook is not None: - await self._user_connection_hook(conn) - - if not conn.autocommit: - await conn.rollback() + try: + hook_res = self._user_connection_hook(conn) + if inspect.isawaitable(hook_res): + await hook_res + if not conn.autocommit: + await conn.commit() + except BaseException: + if not conn.autocommit: + with suppress(Exception): + await conn.rollback() + raise def _ensure_pg_textsearch_available(self) -> None: if self._pg_textsearch_available: diff --git a/sqlspec/adapters/psycopg/core.py b/sqlspec/adapters/psycopg/core.py index 2de7d9a12..6539f8bdc 100644 --- a/sqlspec/adapters/psycopg/core.py +++ b/sqlspec/adapters/psycopg/core.py @@ -16,16 +16,12 @@ StatementConfig, build_statement_config_from_profile, ) -from sqlspec.core.config_runtime import ( - build_postgres_extension_probe_names, - is_postgres_extension_active, - resolve_postgres_extension_state, - resolve_runtime_statement_config, -) from sqlspec.driver import ExecutionResult, rows_to_dicts from sqlspec.exceptions import ( CheckViolationError, ConnectionTimeoutError, + DatabaseConnectionError, + DataError, DeadlockError, ForeignKeyViolationError, ImproperConfigurationError, @@ -34,9 +30,11 @@ OperationalError, OperationCancelledError, PermissionDeniedError, + QueryTimeoutError, SerializationConflictError, SQLParsingError, SQLSpecError, + TransactionError, UniqueViolationError, _classify_timeout_or_cancellation, map_sqlstate_to_exception, @@ -54,6 +52,7 @@ from sqlspec.core import ParameterStyleConfig, StackOperation __all__ = ( + "OPEN_TRANSACTION_STATUSES", "TRANSACTION_STATUS_ACTIVE", "TRANSACTION_STATUS_IDLE", "TRANSACTION_STATUS_INERROR", @@ -67,23 +66,19 @@ "build_async_pipeline_execution_result", "build_copy_from_command", "build_pipeline_execution_result", - "build_postgres_extension_probe_names", + "build_postgres_storage_params_sql", "build_profile", "build_statement_config", "build_truncate_command", - "collect_rows", "create_mapped_exception", "default_statement_config", "driver_profile", "escape_literal_percent", "execute_with_optional_parameters", "execute_with_optional_parameters_async", - "is_postgres_extension_active", "pipeline_supported", "resolve_many_rowcount", - "resolve_postgres_extension_state", "resolve_rowcount", - "resolve_runtime_statement_config", ) TRANSACTION_STATUS_IDLE = 0 @@ -91,6 +86,11 @@ TRANSACTION_STATUS_INTRANS = 2 TRANSACTION_STATUS_INERROR = 3 TRANSACTION_STATUS_UNKNOWN = 4 +OPEN_TRANSACTION_STATUSES: Final[tuple[int, ...]] = ( + TRANSACTION_STATUS_ACTIVE, + TRANSACTION_STATUS_INTRANS, + TRANSACTION_STATUS_INERROR, +) class PreparedStackOperation(NamedTuple): @@ -134,6 +134,39 @@ def build_truncate_command(table: str) -> "PsycopgComposed": return PsycopgSQL("TRUNCATE TABLE {}").format(_compose_table_identifier(table)) +def build_postgres_storage_params_sql( + settings: "dict[str, object] | Mapping[str, object] | None", + *, + include_autovacuum: bool = False, + default_fillfactor: int | None = None, + default_vacuum_scale: float | None = None, + default_analyze_scale: float | None = None, + section: str = "storage", +) -> str: + """Format PostgreSQL table storage parameters (WITH fillfactor = ...) clause.""" + options: list[str] = [] + config_dict = dict(settings) if settings else {} + fillfactor = config_dict.get("fillfactor", default_fillfactor) + if fillfactor is not None: + if not isinstance(fillfactor, int) or isinstance(fillfactor, bool) or fillfactor not in range(10, 101): + msg = f"extension_config['{section}']['fillfactor'] must be an integer from 10 to 100" + raise ValueError(msg) + options.append(f"fillfactor = {fillfactor}") + if include_autovacuum: + scales = ( + ("autovacuum_vacuum_scale_factor", default_vacuum_scale), + ("autovacuum_analyze_scale_factor", default_analyze_scale), + ) + for key, default_val in scales: + value = config_dict.get(key, default_val) + if value is not None: + if not isinstance(value, (int, float)) or isinstance(value, bool) or not 0 <= float(value) <= 1: + msg = f"extension_config['{section}']['{key}'] must be a number from 0 to 1" + raise ValueError(msg) + options.append(f"{key} = {float(value):g}") + return f" WITH ({', '.join(options)})" if options else "" + + def build_profile() -> "DriverParameterProfile": """Create the psycopg driver parameter profile.""" @@ -181,12 +214,10 @@ def apply_driver_features( features.setdefault("enable_paradedb", True) features.setdefault("enable_pg_textsearch", True) features.setdefault("enable_alloydb_columnar", False) - pgbouncer_enabled = bool(features.get("pgbouncer") or features.get("pgbouncer_mode")) - features["pgbouncer"] = pgbouncer_enabled - features["pgbouncer_mode"] = pgbouncer_enabled + features.setdefault("enable_pgcopy_binary", False) raw_pooler_mode = features.get("pooler_mode") if raw_pooler_mode is None: - features["pooler_mode"] = "transaction" if pgbouncer_enabled else "session" + features["pooler_mode"] = "session" else: normalized_mode = str(raw_pooler_mode).lower() if normalized_mode not in {"session", "transaction", "statement", "auto"}: @@ -194,10 +225,7 @@ def apply_driver_features( f"Invalid pooler_mode: {raw_pooler_mode!r}. Expected 'session', 'transaction', 'statement', or 'auto'." ) raise ImproperConfigurationError(msg) - features["pooler_mode"] = "transaction" if normalized_mode == "auto" and pgbouncer_enabled else normalized_mode - if features["pooler_mode"] in {"transaction", "statement"}: - features["pgbouncer"] = True - features["pgbouncer_mode"] = True + features["pooler_mode"] = "session" if normalized_mode == "auto" else normalized_mode features.setdefault("direct_dsn", None) parameter_config = _parameter_config(driver_profile, serializer, deserializer) @@ -270,27 +298,12 @@ def __init__(self, driver: Any, sql: str, parameters: Any, chunk_size: int) -> N def start(self) -> None: handler = self._driver.handle_database_exceptions() - with handler: - transaction = self._driver.connection.transaction() - transaction.__enter__() - self._transaction = transaction - try: - cursor = self._driver.connection.cursor(name=f"sqlspec_stream_{uuid4().hex}") - cursor.itersize = self._chunk_size - execute_with_optional_parameters(cursor, self._sql, self._parameters) - self._cursor = cursor - except BaseException as exc: - self._transaction = None - with contextlib.suppress(Exception): - transaction.__exit__(type(exc), exc, exc.__traceback__) - raise + self._driver._run_with_exception_handler(handler, self._start) self._driver._check_pending_exception(handler) def fetch_chunk(self) -> "list[dict[str, Any]]": handler = self._driver.handle_database_exceptions() - rows: list[Any] = [] - with handler: - rows = self._cursor.fetchmany(self._chunk_size) + rows = self._driver._run_with_exception_handler(handler, self._cursor.fetchmany, self._chunk_size) self._driver._check_pending_exception(handler) if not rows: return [] @@ -313,6 +326,21 @@ def close(self, error: bool = False) -> None: else: transaction.__exit__(None, None, None) + def _start(self) -> None: + transaction = self._driver.connection.transaction() + transaction.__enter__() + self._transaction = transaction + try: + cursor = self._driver.connection.cursor(name=f"sqlspec_stream_{uuid4().hex}") + cursor.itersize = self._chunk_size + execute_with_optional_parameters(cursor, self._sql, self._parameters) + self._cursor = cursor + except BaseException as exc: + self._transaction = None + with contextlib.suppress(Exception): + transaction.__exit__(type(exc), exc, exc.__traceback__) + raise + class PsycopgAsyncStreamSource: """Compiled async chunk source streaming dict rows from a psycopg server-side named cursor. @@ -570,17 +598,22 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpec Returns: A SQLSpec exception that wraps the original error """ - del logger + _ = logger error_type = type(error) mapped_error = _EXCEPTION_MAPPING.get(error_type) if mapped_error is None: mapped_error = _resolve_exception_mapping(error_type) if mapped_error is not None: mapped_error_code, error_class, description = mapped_error - if mapped_error_code == "57014": + sqlstate_code = getattr(error, "sqlstate", None) or mapped_error_code + if sqlstate_code == "57014": termination_class = _classify_timeout_or_cancellation(str(error)) or OperationalError - return _create_postgres_error(error, mapped_error_code, termination_class, description) - return _create_postgres_error(error, mapped_error_code, error_class, description) + return _create_postgres_error(error, sqlstate_code, termination_class, description) + if sqlstate_code == "55P03": + return _create_postgres_error(error, sqlstate_code, QueryTimeoutError, description) + if sqlstate_code.startswith("25"): + return _create_postgres_error(error, sqlstate_code, TransactionError, description) + return _create_postgres_error(error, sqlstate_code, error_class, description) sqlstate_attr = error.sqlstate if has_sqlstate(error) else None error_code = sqlstate_attr if sqlstate_attr is not None else None @@ -588,6 +621,10 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpec if error_code == "57014": termination_class = _classify_timeout_or_cancellation(str(error)) or OperationalError return _create_postgres_error(error, error_code, termination_class, "query terminated") + if error_code == "55P03": + return _create_postgres_error(error, error_code, QueryTimeoutError, "lock timeout") + if error_code.startswith("25"): + return _create_postgres_error(error, error_code, TransactionError, "invalid transaction state") exc_class = map_sqlstate_to_exception(error_code) if exc_class: return _create_postgres_error(error, error_code, exc_class, "database error") @@ -614,7 +651,14 @@ def escape_literal_percent(sql: str, parameters: Any, validator: "ParameterValid def _register_exception_mappings() -> None: + from sqlspec.typing import import_optional_attr + + pool_timeout = import_optional_attr("psycopg_pool", "PoolTimeout") + if pool_timeout is not None: + _EXCEPTION_MAPPING[pool_timeout] = ("", ConnectionTimeoutError, "connection pool timeout") + _EXCEPTION_MAPPING.update({ + pg_errors.LockNotAvailable: ("55P03", QueryTimeoutError, "lock not available"), pg_errors.UniqueViolation: ("23505", UniqueViolationError, "unique constraint violation"), pg_errors.ForeignKeyViolation: ("23503", ForeignKeyViolationError, "foreign key constraint violation"), pg_errors.NotNullViolation: ("23502", NotNullViolationError, "not-null constraint violation"), @@ -627,6 +671,10 @@ def _register_exception_mappings() -> None: pg_errors.AdminShutdown: ("57P01", ConnectionTimeoutError, "admin shutdown"), pg_errors.CannotConnectNow: ("57P03", ConnectionTimeoutError, "cannot connect now"), pg_errors.SyntaxError: ("42601", SQLParsingError, "SQL syntax error"), + psycopg.OperationalError: ("", DatabaseConnectionError, "operational error"), + psycopg.DataError: ("", DataError, "data error"), + psycopg.NotSupportedError: ("", SQLSpecError, "not supported"), + psycopg.InterfaceError: ("", DatabaseConnectionError, "interface error"), }) diff --git a/sqlspec/adapters/psycopg/data_dictionary.py b/sqlspec/adapters/psycopg/data_dictionary.py index e0c245ea0..a4413254f 100644 --- a/sqlspec/adapters/psycopg/data_dictionary.py +++ b/sqlspec/adapters/psycopg/data_dictionary.py @@ -28,7 +28,12 @@ system_metadata_gated_result, unsupported_system_metadata_capability, ) -from sqlspec.data_dictionary.dialects.postgres import resolve_postgres_json_type +from sqlspec.data_dictionary.dialects.postgres import ( + POSTGRES_METADATA_DOMAINS, + POSTGRES_SUPPORTED_DOMAINS, + POSTGRES_SYSTEM_METADATA_QUERIES, + resolve_postgres_json_type, +) from sqlspec.driver import AsyncDataDictionaryBase, SyncDataDictionaryBase if TYPE_CHECKING: @@ -38,37 +43,6 @@ __all__ = ("PsycopgAsyncDataDictionary", "PsycopgSyncDataDictionary") -_POSTGRES_METADATA_DOMAINS: tuple[str, ...] = ( - "schemas", - "objects", - "tables", - "columns", - "constraints", - "indexes", - "views", - "materialized_views", - "sequences", - "types", - "routines", - "triggers", - "comments", - "privileges", - "dependencies", - "extensions", - "partitions", - "ddl", - "system", -) -_POSTGRES_SUPPORTED_DOMAINS = frozenset(d for d in _POSTGRES_METADATA_DOMAINS if d != "system") -_POSTGRES_SYSTEM_METADATA_QUERIES = { - "settings": "settings", - "statement_history": "pg_stat_statements", - "pg_stat_statements": "pg_stat_statements", - "table_statistics": "table_stats", - "table_stats": "table_stats", -} - - @mypyc_attr(allow_interpreted_subclasses=True, native_class=False) class PsycopgSyncDataDictionary(SyncDataDictionaryBase): """PostgreSQL-specific sync data dictionary.""" @@ -89,7 +63,7 @@ def get_system_metadata_capabilities( ) -> tuple[SystemMetadataCapability, ...]: """Get PostgreSQL opt-in system metadata capability disclosures.""" _ = driver - requested_domains = tuple(_POSTGRES_SYSTEM_METADATA_QUERIES) if domains is None else tuple(domains) + requested_domains = tuple(POSTGRES_SYSTEM_METADATA_QUERIES) if domains is None else tuple(domains) return tuple(_postgres_system_metadata_capability(domain) for domain in requested_domains) def get_schemas(self, driver: "PsycopgSyncDriver") -> MetadataResult: @@ -163,7 +137,7 @@ def get_system_metadata( if gate_result.capability.support != MetadataSupport.SUPPORTED: return gate_result - query_name = _POSTGRES_SYSTEM_METADATA_QUERIES.get(metadata_request.domain) + query_name = POSTGRES_SYSTEM_METADATA_QUERIES.get(metadata_request.domain) if query_name is None: return gate_result parameters: dict[str, object] = {} @@ -367,7 +341,7 @@ async def get_system_metadata_capabilities( ) -> tuple[SystemMetadataCapability, ...]: """Get PostgreSQL opt-in system metadata capability disclosures.""" _ = driver - requested_domains = tuple(_POSTGRES_SYSTEM_METADATA_QUERIES) if domains is None else tuple(domains) + requested_domains = tuple(POSTGRES_SYSTEM_METADATA_QUERIES) if domains is None else tuple(domains) return tuple(_postgres_system_metadata_capability(domain) for domain in requested_domains) async def get_schemas(self, driver: "PsycopgAsyncDriver") -> MetadataResult: @@ -443,7 +417,7 @@ async def get_system_metadata( if gate_result.capability.support != MetadataSupport.SUPPORTED: return gate_result - query_name = _POSTGRES_SYSTEM_METADATA_QUERIES.get(metadata_request.domain) + query_name = POSTGRES_SYSTEM_METADATA_QUERIES.get(metadata_request.domain) if query_name is None: return gate_result parameters: dict[str, object] = {} @@ -637,7 +611,7 @@ def _postgres_metadata_capability(domain: str) -> MetadataCapability: risks=(MetadataRisk.EXPENSIVE, MetadataRisk.PRIVILEGED), warnings=("System metadata is opt-in and disabled by default.",), ) - if domain in _POSTGRES_SUPPORTED_DOMAINS: + if domain in POSTGRES_SUPPORTED_DOMAINS: return MetadataCapability( domain=domain, support=MetadataSupport.SUPPORTED, @@ -648,7 +622,7 @@ def _postgres_metadata_capability(domain: str) -> MetadataCapability: def _postgres_metadata_profile(adapter: str, domains: Sequence[str] | None) -> MetadataCapabilityProfile: - requested_domains = _POSTGRES_METADATA_DOMAINS if domains is None else tuple(domains) + requested_domains = POSTGRES_METADATA_DOMAINS if domains is None else tuple(domains) return MetadataCapabilityProfile( "postgres", adapter=adapter, @@ -707,7 +681,7 @@ def _ddl_result_from_rows( def _postgres_system_metadata_capability(domain: str) -> SystemMetadataCapability: - if domain not in _POSTGRES_SYSTEM_METADATA_QUERIES: + if domain not in POSTGRES_SYSTEM_METADATA_QUERIES: return unsupported_system_metadata_capability(domain) return SystemMetadataCapability( domain, diff --git a/sqlspec/adapters/psycopg/driver.py b/sqlspec/adapters/psycopg/driver.py index 3ddebc4eb..d94369abd 100644 --- a/sqlspec/adapters/psycopg/driver.py +++ b/sqlspec/adapters/psycopg/driver.py @@ -1,8 +1,9 @@ """PostgreSQL psycopg driver implementation.""" from collections.abc import Sized -from contextlib import AsyncExitStack, ExitStack +from contextlib import AsyncExitStack, ExitStack, suppress from itertools import chain +from time import perf_counter from typing import TYPE_CHECKING, Any, Literal, LiteralString, cast from sqlspec.adapters.psycopg._typing import ( @@ -16,10 +17,10 @@ PsycopgSyncSessionContext, ) from sqlspec.adapters.psycopg._typing import psycopg_module as psycopg +from sqlspec.adapters.psycopg._typing import psycopg_sql as pg_sql from sqlspec.adapters.psycopg.core import ( - TRANSACTION_STATUS_ACTIVE, + OPEN_TRANSACTION_STATUSES, TRANSACTION_STATUS_INERROR, - TRANSACTION_STATUS_INTRANS, PipelineCursorEntry, PreparedStackOperation, PsycopgAsyncStreamSource, @@ -63,7 +64,7 @@ ) from sqlspec.exceptions import SQLSpecError, StackExecutionError from sqlspec.utils.logging import get_logger -from sqlspec.utils.text import normalize_identifier, quote_identifier +from sqlspec.utils.text import normalize_identifier from sqlspec.utils.type_guards import is_readable, resolve_row_format if TYPE_CHECKING: @@ -73,6 +74,7 @@ from sqlspec.core import ArrowResult from sqlspec.driver import CachedQuery, ExecutionResult from sqlspec.storage import StorageBridgeJob, StorageDestination, StorageFormat, StorageTelemetry + from sqlspec.typing import ArrowRecordBatch __all__ = ( @@ -302,16 +304,37 @@ def dispatch_execute_script(self, cursor: Any, statement: "SQL") -> "ExecutionRe sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) statements = self.split_script_statements(sql, statement.statement_config, strip_trailing_semicolon=True) + statement_count = len(statements) successful_count = 0 last_cursor = cursor + runtime = self._observability + has_observers = runtime is not None and runtime.has_statement_observers - for stmt in statements: + for idx, stmt in enumerate(statements): stmt_sql = escape_literal_percent(stmt, prepared_parameters, self.statement_config.parameter_validator) + stmt_started = perf_counter() if has_observers else 0.0 + self._pending_statement_profile = None execute_with_optional_parameters(cursor, stmt_sql, prepared_parameters) + if has_observers: + stmt_duration = perf_counter() - stmt_started + stmt_profile = self._pending_statement_profile + self._pending_statement_profile = None + self._emit_script_statement_event( + sql=stmt, + parameters=prepared_parameters, + operation="SCRIPT", + rows_affected=resolve_rowcount(cursor), + duration_s=stmt_duration, + started_at=stmt_started, + profile=stmt_profile, + script_ordinal=idx, + script_statement_count=statement_count, + ) successful_count += 1 + self._pending_statement_profile = None return self.create_execution_result( - last_cursor, statement_count=len(statements), successful_statements=successful_count, is_script_result=True + last_cursor, statement_count=statement_count, successful_statements=successful_count, is_script_result=True ) def dispatch_special_handling(self, cursor: Any, statement: "SQL") -> "SQLResult | None": @@ -380,17 +403,21 @@ def dispatch_special_handling(self, cursor: Any, statement: "SQL") -> "SQLResult def begin(self) -> None: """Begin a database transaction on the current connection.""" + if getattr(getattr(self.connection, "info", None), "transaction_status", None) == TRANSACTION_STATUS_INERROR: + self.connection.rollback() if self._connection_in_transaction(): return + restore_autocommit = False try: - if self.connection.info.transaction_status == TRANSACTION_STATUS_INERROR: - self.connection.rollback() restore_autocommit = self.connection.autocommit if restore_autocommit: self.connection.autocommit = False if self._should_set_local_columnar(): self.connection.execute("SET LOCAL google_columnar_engine.enable_columnar_scan = 'on'") except psycopg.Error as e: + if restore_autocommit: + with suppress(Exception): + self.connection.autocommit = True mapped = create_mapped_exception(e) if type(mapped) is not SQLSpecError: raise mapped from e @@ -428,22 +455,20 @@ def rollback(self) -> None: def set_migration_session_schema(self, schema: str) -> None: """Set the PostgreSQL search path for migration SQL.""" normalized_schema = normalize_identifier(schema, "postgres") - quoted_schema = quote_identifier(normalized_schema) - sql = cast("LiteralString", f'SET LOCAL search_path TO {quoted_schema}, "$user", public') # type: ignore[redundant-cast] + sql = PsycopgSQL('SET LOCAL search_path TO {}, "$user", public').format(pg_sql.Identifier(normalized_schema)) with self.with_cursor(self.connection) as cursor: cursor.execute(sql) def set_migration_non_transactional_schema(self, schema: str) -> None: """Set the PostgreSQL search path for non-transactional migration SQL.""" normalized_schema = normalize_identifier(schema, "postgres") - quoted_schema = quote_identifier(normalized_schema) - sql = cast("LiteralString", f'SET search_path TO {quoted_schema}, "$user", public') # type: ignore[redundant-cast] + sql = PsycopgSQL('SET search_path TO {}, "$user", public').format(pg_sql.Identifier(normalized_schema)) with self.with_cursor(self.connection) as cursor: cursor.execute(sql) def reset_migration_session_schema(self) -> None: """Reset the PostgreSQL search path after non-transactional migration SQL.""" - sql = cast("LiteralString", "RESET search_path") # type: ignore[redundant-cast] + sql = "RESET search_path" with self.with_cursor(self.connection) as cursor: cursor.execute(sql) @@ -551,10 +576,9 @@ def load_from_arrow( if overwrite: truncate_sql = build_truncate_command(table) exc_handler = self.handle_database_exceptions() - with self.with_cursor(self.connection) as cursor, exc_handler: - cursor.execute(truncate_sql) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + with self.with_cursor(self.connection) as cursor: + self._run_with_exception_handler(exc_handler, cursor.execute, truncate_sql) + self._check_pending_exception(exc_handler) columns, batch_iter = self._coerce_arrow_batches(source) first_batch = next((b for b in batch_iter if b.num_rows > 0), None) metrics: dict[str, int] = {"rows_processed": 0, "bytes_processed": 0} @@ -563,47 +587,38 @@ def load_from_arrow( copy_sql = build_copy_from_command(table, columns, binary=use_binary) all_batches = chain((first_batch,), (b for b in batch_iter if b.num_rows > 0)) exc_handler = self.handle_database_exceptions() - with ExitStack() as stack: - stack.enter_context(exc_handler) - cursor = stack.enter_context(self.with_cursor(self.connection)) - copy_ctx = stack.enter_context(cursor.copy(copy_sql)) - if use_binary: - for chunk in self._iter_pgcopy_binary_chunks(all_batches, columns=columns, metrics=metrics): - copy_ctx.write(chunk) - else: - for batch in all_batches: - metrics["rows_processed"] += int(batch.num_rows) - metrics["bytes_processed"] += int(batch.nbytes) - records = self._arrow_batch_to_rows(batch, columns) - if self._arrow_rows_need_preparation(batch): - records = cast( - "list[Any]", - self.prepare_driver_parameters(records, self.statement_config, is_many=True), - ) - for record in records: - copy_ctx.write_row(record) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + self._run_with_exception_handler( + exc_handler, self._copy_arrow_batches, copy_sql, all_batches, columns, metrics, use_binary + ) + self._check_pending_exception(exc_handler) telemetry_payload = self._ingest_telemetry_from_batches(metrics["rows_processed"], metrics["bytes_processed"]) telemetry_payload["destination"] = table self._attach_partition_telemetry(telemetry_payload, partitioner) return self._storage_job(telemetry_payload, telemetry) - def load_from_storage( + def _copy_arrow_batches( self, - table: str, - source: "StorageDestination", - *, - file_format: "StorageFormat", - partitioner: "dict[str, object] | None" = None, - overwrite: bool = False, - **kwargs: Any, - ) -> "StorageBridgeJob": - """Load staged artifacts into PostgreSQL via COPY.""" - - return super().load_from_storage( - table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs - ) + copy_sql: "PsycopgComposed", + batches: "abc.Iterable[ArrowRecordBatch]", + columns: list[str], + metrics: dict[str, int], + use_binary: bool, + ) -> None: + with self.with_cursor(self.connection) as cursor, cursor.copy(copy_sql) as copy_ctx: + if use_binary: + for chunk in self._iter_pgcopy_binary_chunks(batches, columns=columns, metrics=metrics): + copy_ctx.write(chunk) + else: + for batch in batches: + metrics["rows_processed"] += int(batch.num_rows) + metrics["bytes_processed"] += int(batch.nbytes) + records = self._arrow_batch_to_rows(batch, columns) + if self._arrow_rows_need_preparation(batch): + records = cast( + "list[Any]", self.prepare_driver_parameters(records, self.statement_config, is_many=True) + ) + for record in records: + copy_ctx.write_row(record) def collect_rows(self, cursor: Any, fetched: "list[Any]") -> "tuple[list[Any], list[str], int]": """Collect psycopg sync rows for the direct execution path.""" @@ -646,7 +661,7 @@ def _execute_stack_pipeline( cursor = resource_stack.enter_context(self.with_cursor(self.connection)) try: - sql = cast("LiteralString | PsycopgSQL | PsycopgComposed", prepared.sql) # type: ignore[redundant-cast] + sql = prepared.sql if prepared.parameters: cursor.execute(sql, prepared.parameters) else: @@ -711,10 +726,9 @@ def _resolve_column_names(self, description: Any) -> list[str]: def _connection_in_transaction(self) -> bool: """Check if connection is in transaction.""" - return self._transaction_active or self.connection.info.transaction_status in { - TRANSACTION_STATUS_ACTIVE, - TRANSACTION_STATUS_INTRANS, - } + info = getattr(self.connection, "info", None) + status = getattr(info, "transaction_status", None) + return self._transaction_active or status in OPEN_TRANSACTION_STATUSES def _restore_original_autocommit(self) -> None: """Restore autocommit after a completed SQLSpec-owned transaction.""" @@ -724,6 +738,9 @@ def _restore_original_autocommit(self) -> None: try: self.connection.autocommit = True except psycopg.Error as e: + mapped = create_mapped_exception(e) + if type(mapped) is not SQLSpecError: + raise mapped from e msg = f"Failed to restore autocommit: {e}" raise SQLSpecError(msg) from e @@ -731,11 +748,7 @@ def _should_set_local_columnar(self) -> bool: """Return True when AlloyDB columnar GUC must be set transaction-locally.""" if not self.driver_features.get("enable_alloydb_columnar"): return False - return bool( - self.driver_features.get("pooler_mode") in {"transaction", "statement"} - or self.driver_features.get("pgbouncer") - or self.driver_features.get("pgbouncer_mode") - ) + return self.driver_features.get("pooler_mode") in {"transaction", "statement"} class PsycopgAsyncExceptionHandler(BaseAsyncExceptionHandler): @@ -867,16 +880,37 @@ async def dispatch_execute_script(self, cursor: Any, statement: "SQL") -> "Execu sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) statements = self.split_script_statements(sql, statement.statement_config, strip_trailing_semicolon=True) + statement_count = len(statements) successful_count = 0 last_cursor = cursor + runtime = self._observability + has_observers = runtime is not None and runtime.has_statement_observers - for stmt in statements: + for idx, stmt in enumerate(statements): stmt_sql = escape_literal_percent(stmt, prepared_parameters, self.statement_config.parameter_validator) + stmt_started = perf_counter() if has_observers else 0.0 + self._pending_statement_profile = None await execute_with_optional_parameters_async(cursor, stmt_sql, prepared_parameters) + if has_observers: + stmt_duration = perf_counter() - stmt_started + stmt_profile = self._pending_statement_profile + self._pending_statement_profile = None + self._emit_script_statement_event( + sql=stmt, + parameters=prepared_parameters, + operation="SCRIPT", + rows_affected=resolve_rowcount(cursor), + duration_s=stmt_duration, + started_at=stmt_started, + profile=stmt_profile, + script_ordinal=idx, + script_statement_count=statement_count, + ) successful_count += 1 + self._pending_statement_profile = None return self.create_execution_result( - last_cursor, statement_count=len(statements), successful_statements=successful_count, is_script_result=True + last_cursor, statement_count=statement_count, successful_statements=successful_count, is_script_result=True ) async def dispatch_special_handling(self, cursor: Any, statement: "SQL") -> "SQLResult | None": @@ -946,17 +980,21 @@ async def dispatch_special_handling(self, cursor: Any, statement: "SQL") -> "SQL async def begin(self) -> None: """Begin a database transaction on the current connection.""" + if getattr(getattr(self.connection, "info", None), "transaction_status", None) == TRANSACTION_STATUS_INERROR: + await self.connection.rollback() if self._connection_in_transaction(): return + restore_autocommit = False try: - if self.connection.info.transaction_status == TRANSACTION_STATUS_INERROR: - await self.connection.rollback() restore_autocommit = self.connection.autocommit if restore_autocommit: await self.connection.set_autocommit(False) if self._should_set_local_columnar(): await self.connection.execute("SET LOCAL google_columnar_engine.enable_columnar_scan = 'on'") except psycopg.Error as e: + if restore_autocommit: + with suppress(Exception): + await self.connection.set_autocommit(True) mapped = create_mapped_exception(e) if type(mapped) is not SQLSpecError: raise mapped from e @@ -994,22 +1032,20 @@ async def rollback(self) -> None: async def set_migration_session_schema(self, schema: str) -> None: """Set the PostgreSQL search path for migration SQL.""" normalized_schema = normalize_identifier(schema, "postgres") - quoted_schema = quote_identifier(normalized_schema) - sql = cast("LiteralString", f'SET LOCAL search_path TO {quoted_schema}, "$user", public') # type: ignore[redundant-cast] + sql = PsycopgSQL('SET LOCAL search_path TO {}, "$user", public').format(pg_sql.Identifier(normalized_schema)) async with self.with_cursor(self.connection) as cursor: await cursor.execute(sql) async def set_migration_non_transactional_schema(self, schema: str) -> None: """Set the PostgreSQL search path for non-transactional migration SQL.""" normalized_schema = normalize_identifier(schema, "postgres") - quoted_schema = quote_identifier(normalized_schema) - sql = cast("LiteralString", f'SET search_path TO {quoted_schema}, "$user", public') # type: ignore[redundant-cast] + sql = PsycopgSQL('SET search_path TO {}, "$user", public').format(pg_sql.Identifier(normalized_schema)) async with self.with_cursor(self.connection) as cursor: await cursor.execute(sql) async def reset_migration_session_schema(self) -> None: """Reset the PostgreSQL search path after non-transactional migration SQL.""" - sql = cast("LiteralString", "RESET search_path") # type: ignore[redundant-cast] + sql = "RESET search_path" async with self.with_cursor(self.connection) as cursor: await cursor.execute(sql) @@ -1122,10 +1158,9 @@ async def load_from_arrow( if overwrite: truncate_sql = build_truncate_command(table) exc_handler = self.handle_database_exceptions() - async with self.with_cursor(self.connection) as cursor, exc_handler: - await cursor.execute(truncate_sql) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + async with self.with_cursor(self.connection) as cursor: + await self._run_with_exception_handler(exc_handler, cursor.execute, truncate_sql) + self._check_pending_exception(exc_handler) columns, batch_iter = self._coerce_arrow_batches(source) first_batch = next((b for b in batch_iter if b.num_rows > 0), None) metrics: dict[str, int] = {"rows_processed": 0, "bytes_processed": 0} @@ -1134,47 +1169,38 @@ async def load_from_arrow( copy_sql = build_copy_from_command(table, columns, binary=use_binary) all_batches = chain((first_batch,), (b for b in batch_iter if b.num_rows > 0)) exc_handler = self.handle_database_exceptions() - async with AsyncExitStack() as stack: - await stack.enter_async_context(exc_handler) - cursor = await stack.enter_async_context(self.with_cursor(self.connection)) - copy_ctx = await stack.enter_async_context(cursor.copy(copy_sql)) - if use_binary: - for chunk in self._iter_pgcopy_binary_chunks(all_batches, columns=columns, metrics=metrics): - await copy_ctx.write(chunk) - else: - for batch in all_batches: - metrics["rows_processed"] += int(batch.num_rows) - metrics["bytes_processed"] += int(batch.nbytes) - records = self._arrow_batch_to_rows(batch, columns) - if self._arrow_rows_need_preparation(batch): - records = cast( - "list[Any]", - self.prepare_driver_parameters(records, self.statement_config, is_many=True), - ) - for record in records: - await copy_ctx.write_row(record) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + await self._run_with_exception_handler( + exc_handler, self._copy_arrow_batches, copy_sql, all_batches, columns, metrics, use_binary + ) + self._check_pending_exception(exc_handler) telemetry_payload = self._ingest_telemetry_from_batches(metrics["rows_processed"], metrics["bytes_processed"]) telemetry_payload["destination"] = table self._attach_partition_telemetry(telemetry_payload, partitioner) return self._storage_job(telemetry_payload, telemetry) - async def load_from_storage( + async def _copy_arrow_batches( self, - table: str, - source: "StorageDestination", - *, - file_format: "StorageFormat", - partitioner: "dict[str, object] | None" = None, - overwrite: bool = False, - **kwargs: Any, - ) -> "StorageBridgeJob": - """Load staged artifacts asynchronously.""" - - return await super().load_from_storage( - table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs - ) + copy_sql: "PsycopgComposed", + batches: "abc.Iterable[ArrowRecordBatch]", + columns: list[str], + metrics: dict[str, int], + use_binary: bool, + ) -> None: + async with self.with_cursor(self.connection) as cursor, cursor.copy(copy_sql) as copy_ctx: + if use_binary: + for chunk in self._iter_pgcopy_binary_chunks(batches, columns=columns, metrics=metrics): + await copy_ctx.write(chunk) + else: + for batch in batches: + metrics["rows_processed"] += int(batch.num_rows) + metrics["bytes_processed"] += int(batch.nbytes) + records = self._arrow_batch_to_rows(batch, columns) + if self._arrow_rows_need_preparation(batch): + records = cast( + "list[Any]", self.prepare_driver_parameters(records, self.statement_config, is_many=True) + ) + for record in records: + await copy_ctx.write_row(record) def collect_rows(self, cursor: Any, fetched: "list[Any]") -> "tuple[list[Any], list[str], int]": """Collect psycopg async rows for the direct execution path.""" @@ -1217,7 +1243,7 @@ async def _execute_stack_pipeline( cursor = await resource_stack.enter_async_context(self.with_cursor(self.connection)) try: - sql = cast("LiteralString | PsycopgSQL | PsycopgComposed", prepared.sql) # type: ignore[redundant-cast] + sql = prepared.sql if prepared.parameters: await cursor.execute(sql, prepared.parameters) else: @@ -1282,10 +1308,9 @@ def _resolve_column_names(self, description: Any) -> list[str]: def _connection_in_transaction(self) -> bool: """Check if connection is in transaction.""" - return self._transaction_active or self.connection.info.transaction_status in { - TRANSACTION_STATUS_ACTIVE, - TRANSACTION_STATUS_INTRANS, - } + info = getattr(self.connection, "info", None) + status = getattr(info, "transaction_status", None) + return self._transaction_active or status in OPEN_TRANSACTION_STATUSES async def _restore_original_autocommit(self) -> None: """Restore autocommit after a completed SQLSpec-owned transaction.""" @@ -1295,6 +1320,9 @@ async def _restore_original_autocommit(self) -> None: try: await self.connection.set_autocommit(True) except psycopg.Error as e: + mapped = create_mapped_exception(e) + if type(mapped) is not SQLSpecError: + raise mapped from e msg = f"Failed to restore autocommit: {e}" raise SQLSpecError(msg) from e @@ -1302,11 +1330,7 @@ def _should_set_local_columnar(self) -> bool: """Return True when AlloyDB columnar GUC must be set transaction-locally.""" if not self.driver_features.get("enable_alloydb_columnar"): return False - return bool( - self.driver_features.get("pooler_mode") in {"transaction", "statement"} - or self.driver_features.get("pgbouncer") - or self.driver_features.get("pgbouncer_mode") - ) + return self.driver_features.get("pooler_mode") in {"transaction", "statement"} def _attribute_pipeline_failure( diff --git a/sqlspec/adapters/psycopg/events/_hub.py b/sqlspec/adapters/psycopg/events/_hub.py index d0634504e..b7deea7a7 100644 --- a/sqlspec/adapters/psycopg/events/_hub.py +++ b/sqlspec/adapters/psycopg/events/_hub.py @@ -14,6 +14,7 @@ import logging import queue as stdlib_queue import threading +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, cast from weakref import WeakKeyDictionary @@ -173,16 +174,9 @@ def _resolve_listener_config(self) -> "PsycopgAsyncConfig": return self._listener_config raw_features = getattr(self._config, "driver_features", None) driver_features = raw_features if isinstance(raw_features, dict) else {} - pooler_mode = str( - driver_features.get("pooler_mode") - or ( - "transaction" - if driver_features.get("pgbouncer") or driver_features.get("pgbouncer_mode") - else "session" - ) - ).lower() + pooler_mode = str(driver_features.get("pooler_mode") or "session").lower() raw_direct_dsn = driver_features.get("direct_dsn") - direct_dsn = raw_direct_dsn if isinstance(raw_direct_dsn, str) and raw_direct_dsn else None + direct_dsn = raw_direct_dsn if isinstance(raw_direct_dsn, (str, Mapping)) and raw_direct_dsn else None if pooler_mode in {"transaction", "statement"} and not direct_dsn: logger.warning( "PostgreSQL LISTEN/NOTIFY cannot run under pooler_mode=%r without configuring driver_features['direct_dsn']", @@ -440,16 +434,9 @@ def _resolve_listener_config(self) -> "PsycopgSyncConfig": return self._listener_config raw_features = getattr(self._config, "driver_features", None) driver_features = raw_features if isinstance(raw_features, dict) else {} - pooler_mode = str( - driver_features.get("pooler_mode") - or ( - "transaction" - if driver_features.get("pgbouncer") or driver_features.get("pgbouncer_mode") - else "session" - ) - ).lower() + pooler_mode = str(driver_features.get("pooler_mode") or "session").lower() raw_direct_dsn = driver_features.get("direct_dsn") - direct_dsn = raw_direct_dsn if isinstance(raw_direct_dsn, str) and raw_direct_dsn else None + direct_dsn = raw_direct_dsn if isinstance(raw_direct_dsn, (str, Mapping)) and raw_direct_dsn else None if pooler_mode in {"transaction", "statement"} and not direct_dsn: logger.warning( "PostgreSQL LISTEN/NOTIFY cannot run under pooler_mode=%r without configuring driver_features['direct_dsn']", diff --git a/sqlspec/adapters/psycopg/events/backend.py b/sqlspec/adapters/psycopg/events/backend.py index 7a683ebb6..a7dbafb9a 100644 --- a/sqlspec/adapters/psycopg/events/backend.py +++ b/sqlspec/adapters/psycopg/events/backend.py @@ -34,7 +34,7 @@ "create_event_backend", ) -logger = get_logger("sqlspec.events.psycopg") +logger = get_logger("sqlspec.adapters.psycopg.events.backend") _MIN_LISTENER_POOL_SIZE = 2 _MAX_SEEN_MARKERS = 1_024 @@ -363,14 +363,15 @@ def dequeue(self, channel: str, poll_interval: float) -> EventMessage | None: return self._dequeue_reconciled(channel) self._queue._reset_empty_poll_delay(channel) self._runtime.increment_metric("events.listener.wakeup") - batch_size = _extract_batch_size(payload) - is_new, should_drain = self._marker_state.register(channel, _extract_marker_id(payload), batch_size) + parsed_payload = _parse_notify_payload(payload) + batch_size = _extract_batch_size(parsed_payload) + is_new, should_drain = self._marker_state.register(channel, _extract_marker_id(parsed_payload), batch_size) if is_new and batch_size > 1: self._runtime.increment_metric("events.marker.coalesced", batch_size - 1) if not should_drain: self._runtime.increment_metric("events.marker.duplicate") continue - event_id = _extract_event_id(payload) + event_id = _extract_event_id(parsed_payload) if event_id is not None: event = self._queue.dequeue_by_event_id(event_id) if event is not None: @@ -534,14 +535,15 @@ async def dequeue(self, channel: str, poll_interval: float) -> EventMessage | No return await self._dequeue_reconciled(channel) self._queue._reset_empty_poll_delay(channel) self._runtime.increment_metric("events.listener.wakeup") - batch_size = _extract_batch_size(payload) - is_new, should_drain = self._marker_state.register(channel, _extract_marker_id(payload), batch_size) + parsed_payload = _parse_notify_payload(payload) + batch_size = _extract_batch_size(parsed_payload) + is_new, should_drain = self._marker_state.register(channel, _extract_marker_id(parsed_payload), batch_size) if is_new and batch_size > 1: self._runtime.increment_metric("events.marker.coalesced", batch_size - 1) if not should_drain: self._runtime.increment_metric("events.marker.duplicate") continue - event_id = _extract_event_id(payload) + event_id = _extract_event_id(parsed_payload) if event_id is not None: event = await self._queue.dequeue_by_event_id(event_id) if event is not None: @@ -685,34 +687,40 @@ def _validate_listener_pool_capacity(config: "PsycopgAsyncConfig | PsycopgSyncCo raise ImproperConfigurationError(msg) -def _extract_event_id(payload: "str | None") -> "str | None": - if not payload: +def _parse_notify_payload(payload: "dict[str, Any] | str | None") -> "dict[str, Any] | None": + if isinstance(payload, dict): + return payload + if not payload or not isinstance(payload, str): return None - raw = from_json(payload) - if isinstance(raw, dict): - event_id = raw.get("event_id") - return event_id if isinstance(event_id, str) else None - return None + try: + raw = from_json(payload) + return raw if isinstance(raw, dict) else None + except (TypeError, ValueError): + return None + + +def _extract_event_id(payload: "dict[str, Any] | str | None") -> "str | None": + parsed = payload if isinstance(payload, dict) else _parse_notify_payload(payload) + if parsed is None: + return None + event_id = parsed.get("event_id") + return event_id if isinstance(event_id, str) else None -def _extract_batch_size(payload: "str | None") -> int: - if not payload: +def _extract_batch_size(payload: "dict[str, Any] | str | None") -> int: + parsed = payload if isinstance(payload, dict) else _parse_notify_payload(payload) + if parsed is None: return 0 - raw = from_json(payload) - if isinstance(raw, dict): - batch_size = raw.get("batch_size") - return batch_size if isinstance(batch_size, int) else 0 - return 0 + batch_size = parsed.get("batch_size") + return batch_size if isinstance(batch_size, int) else 0 -def _extract_marker_id(payload: "str | None") -> "str | None": - if not payload: +def _extract_marker_id(payload: "dict[str, Any] | str | None") -> "str | None": + parsed = payload if isinstance(payload, dict) else _parse_notify_payload(payload) + if parsed is None: return None - raw = from_json(payload) - if isinstance(raw, dict): - marker_id = raw.get("marker_id", raw.get("event_id")) - return marker_id if isinstance(marker_id, str) else None - return None + marker_id = parsed.get("marker_id", parsed.get("event_id")) + return marker_id if isinstance(marker_id, str) else None def _record_dequeue_result(runtime: Any, event: "EventMessage | None") -> None: diff --git a/sqlspec/adapters/psycopg/events/store.py b/sqlspec/adapters/psycopg/events/store.py index 6a3b327f2..36f1d720a 100644 --- a/sqlspec/adapters/psycopg/events/store.py +++ b/sqlspec/adapters/psycopg/events/store.py @@ -3,6 +3,7 @@ from typing import NotRequired from sqlspec.adapters.psycopg.config import PsycopgAsyncConfig, PsycopgSyncConfig +from sqlspec.adapters.psycopg.core import build_postgres_storage_params_sql from sqlspec.config import EventsConfig from sqlspec.extensions.events import BaseEventQueueStore @@ -50,7 +51,7 @@ def _column_types(self) -> "tuple[str, str, str]": def _table_clause(self) -> str: """Return explicitly configured PostgreSQL table tuning options.""" - return _postgres_table_options(self.settings) + return build_postgres_storage_params_sql(self.settings, include_autovacuum=True, section="events") class PsycopgAsyncEventQueueStore(BaseEventQueueStore[PsycopgAsyncConfig]): @@ -78,27 +79,4 @@ def _column_types(self) -> "tuple[str, str, str]": def _table_clause(self) -> str: """Return explicitly configured PostgreSQL table tuning options.""" - return _postgres_table_options(self.settings) - - -def _postgres_table_options(settings: "dict[str, object]") -> str: - options = _postgres_option_values(settings) - return f" WITH ({', '.join(options)})" if options else "" - - -def _postgres_option_values(settings: "dict[str, object]") -> "list[str]": - options: list[str] = [] - fillfactor = settings.get("fillfactor") - if fillfactor is not None: - if not isinstance(fillfactor, int) or isinstance(fillfactor, bool) or fillfactor not in range(10, 101): - msg = "extension_config['events']['fillfactor'] must be an integer from 10 to 100" - raise ValueError(msg) - options.append(f"fillfactor = {fillfactor}") - for key in ("autovacuum_vacuum_scale_factor", "autovacuum_analyze_scale_factor"): - value = settings.get(key) - if value is not None: - if not isinstance(value, (int, float)) or isinstance(value, bool) or not 0 <= float(value) <= 1: - msg = f"extension_config['events']['{key}'] must be a number from 0 to 1" - raise ValueError(msg) - options.append(f"{key} = {float(value):g}") - return options + return build_postgres_storage_params_sql(self.settings, include_autovacuum=True, section="events") diff --git a/sqlspec/adapters/psycopg/litestar/store.py b/sqlspec/adapters/psycopg/litestar/store.py index 3d7a2e7ab..865ab6fd6 100644 --- a/sqlspec/adapters/psycopg/litestar/store.py +++ b/sqlspec/adapters/psycopg/litestar/store.py @@ -7,6 +7,8 @@ from typing import TYPE_CHECKING, Any, NotRequired, cast from sqlspec.adapters.psycopg._typing import psycopg_dict_row as dict_row +from sqlspec.adapters.psycopg._typing import psycopg_errors as errors +from sqlspec.adapters.psycopg.core import build_postgres_storage_params_sql from sqlspec.config import LitestarConfig from sqlspec.extensions.litestar.store import BaseSQLSpecStore from sqlspec.utils.sync_tools import async_ @@ -92,27 +94,30 @@ async def get(self, key: str, renew_for: "int | timedelta | None" = None) -> "by AND (expires_at IS NULL OR expires_at > CURRENT_TIMESTAMP) """ - conn_context = self._config.provide_connection() - async with conn_context as conn: - async with conn.cursor(row_factory=dict_row) as cur: - await cur.execute(cast("Any", sql), (key,)) - row = await cur.fetchone() - - if row is None: - return None - - if renew_for is not None and row["expires_at"] is not None: - new_expires_at = self._calculate_expires_at(renew_for) - if new_expires_at is not None: - update_sql = f""" - UPDATE {self._table_name} - SET expires_at = %s, updated_at = CURRENT_TIMESTAMP - WHERE session_id = %s - """ - await conn.execute(cast("Any", update_sql), (new_expires_at, key)) - await conn.commit() - - return bytes(row["data"]) + try: + conn_context = self._config.provide_connection() + async with conn_context as conn: + async with conn.cursor(row_factory=dict_row) as cur: + await cur.execute(cast("Any", sql), (key,)) + row = await cur.fetchone() + + if row is None: + return None + + if renew_for is not None and row["expires_at"] is not None: + new_expires_at = self._calculate_expires_at(renew_for) + if new_expires_at is not None: + update_sql = f""" + UPDATE {self._table_name} + SET expires_at = %s, updated_at = CURRENT_TIMESTAMP + WHERE session_id = %s + """ + await conn.execute(cast("Any", update_sql), (new_expires_at, key)) + await conn.commit() + + return bytes(row["data"]) + except errors.UndefinedTable: + return None async def set(self, key: str, value: "str | bytes", expires_in: "int | timedelta | None" = None) -> None: """Store a session value. @@ -148,20 +153,26 @@ async def delete(self, key: str) -> None: """ sql = f"DELETE FROM {self._table_name} WHERE session_id = %s" - conn_context = self._config.provide_connection() - async with conn_context as conn: - await conn.execute(cast("Any", sql), (key,)) - await conn.commit() + try: + conn_context = self._config.provide_connection() + async with conn_context as conn: + await conn.execute(cast("Any", sql), (key,)) + await conn.commit() + except errors.UndefinedTable: + return async def delete_all(self) -> None: """Delete all sessions from the store.""" sql = f"DELETE FROM {self._table_name}" - conn_context = self._config.provide_connection() - async with conn_context as conn: - await conn.execute(cast("Any", sql)) - await conn.commit() - self._log_delete_all() + try: + conn_context = self._config.provide_connection() + async with conn_context as conn: + await conn.execute(cast("Any", sql)) + await conn.commit() + self._log_delete_all() + except errors.UndefinedTable: + return async def exists(self, key: str) -> bool: """Check if a session key exists and is not expired. @@ -178,11 +189,14 @@ async def exists(self, key: str) -> bool: AND (expires_at IS NULL OR expires_at > CURRENT_TIMESTAMP) """ - conn_context = self._config.provide_connection() - async with conn_context as conn, conn.cursor() as cur: - await cur.execute(cast("Any", sql), (key,)) - result = await cur.fetchone() - return result is not None + try: + conn_context = self._config.provide_connection() + async with conn_context as conn, conn.cursor() as cur: + await cur.execute(cast("Any", sql), (key,)) + result = await cur.fetchone() + return result is not None + except errors.UndefinedTable: + return False async def expires_in(self, key: str) -> "int | None": """Get the time in seconds until the session expires. @@ -198,22 +212,25 @@ async def expires_in(self, key: str) -> "int | None": WHERE session_id = %s """ - conn_context = self._config.provide_connection() - async with conn_context as conn: - async with conn.cursor(row_factory=dict_row) as cur: - await cur.execute(cast("Any", sql), (key,)) - row = await cur.fetchone() + try: + conn_context = self._config.provide_connection() + async with conn_context as conn: + async with conn.cursor(row_factory=dict_row) as cur: + await cur.execute(cast("Any", sql), (key,)) + row = await cur.fetchone() - if row is None or row["expires_at"] is None: - return None + if row is None or row["expires_at"] is None: + return None - expires_at = row["expires_at"] - now = datetime.now(UTC) - if expires_at <= now: - return 0 + expires_at = row["expires_at"] + now = datetime.now(UTC) + if expires_at <= now: + return 0 - delta = expires_at - now - return int(delta.total_seconds()) + delta = expires_at - now + return int(delta.total_seconds()) + except errors.UndefinedTable: + return None async def delete_expired(self) -> int: """Delete all expired sessions. @@ -223,14 +240,17 @@ async def delete_expired(self) -> int: """ sql = f"DELETE FROM {self._table_name} WHERE expires_at <= CURRENT_TIMESTAMP" - conn_context = self._config.provide_connection() - async with conn_context as conn, conn.cursor() as cur: - await cur.execute(cast("Any", sql)) - await conn.commit() - count = cur.rowcount if cur.rowcount and cur.rowcount > 0 else 0 - if count > 0: - self._log_delete_expired(count) - return count + try: + conn_context = self._config.provide_connection() + async with conn_context as conn, conn.cursor() as cur: + await cur.execute(cast("Any", sql)) + await conn.commit() + count = cur.rowcount if cur.rowcount and cur.rowcount > 0 else 0 + if count > 0: + self._log_delete_expired(count) + return count + except errors.UndefinedTable: + return 0 def _table_ddl(self) -> str: """Get PostgreSQL CREATE TABLE SQL with optimized schema. @@ -238,7 +258,14 @@ def _table_ddl(self) -> str: Returns: SQL statement to create the sessions table with proper indexes. """ - fillfactor, vacuum_scale, analyze_scale = _postgres_litestar_tuning(self._config) + tuning_clause = build_postgres_storage_params_sql( + self._config.extension_config.get("litestar", {}), + default_fillfactor=80, + default_vacuum_scale=0.05, + default_analyze_scale=0.02, + include_autovacuum=True, + section="litestar", + ) return f""" CREATE TABLE IF NOT EXISTS {self._table_name} ( session_id TEXT PRIMARY KEY, @@ -246,15 +273,10 @@ def _table_ddl(self) -> str: expires_at TIMESTAMPTZ, created_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP, updated_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP - ) WITH (fillfactor = {fillfactor}); + ){tuning_clause}; CREATE INDEX IF NOT EXISTS idx_{self._table_name}_expires_at ON {self._table_name}(expires_at) WHERE expires_at IS NOT NULL; - - ALTER TABLE {self._table_name} SET ( - autovacuum_vacuum_scale_factor = {vacuum_scale:g}, - autovacuum_analyze_scale_factor = {analyze_scale:g} - ); """ def _drop_table_sql(self) -> "list[str]": @@ -376,7 +398,14 @@ def _table_ddl(self) -> str: Returns: SQL statement to create the sessions table with proper indexes. """ - fillfactor, vacuum_scale, analyze_scale = _postgres_litestar_tuning(self._config) + tuning_clause = build_postgres_storage_params_sql( + self._config.extension_config.get("litestar", {}), + default_fillfactor=80, + default_vacuum_scale=0.05, + default_analyze_scale=0.02, + include_autovacuum=True, + section="litestar", + ) return f""" CREATE TABLE IF NOT EXISTS {self._table_name} ( session_id TEXT PRIMARY KEY, @@ -384,15 +413,10 @@ def _table_ddl(self) -> str: expires_at TIMESTAMPTZ, created_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP, updated_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP - ) WITH (fillfactor = {fillfactor}); + ){tuning_clause}; CREATE INDEX IF NOT EXISTS idx_{self._table_name}_expires_at ON {self._table_name}(expires_at) WHERE expires_at IS NOT NULL; - - ALTER TABLE {self._table_name} SET ( - autovacuum_vacuum_scale_factor = {vacuum_scale:g}, - autovacuum_analyze_scale_factor = {analyze_scale:g} - ); """ def _drop_table_sql(self) -> "list[str]": @@ -419,26 +443,29 @@ def _get(self, key: str, renew_for: "int | timedelta | None" = None) -> "bytes | AND (expires_at IS NULL OR expires_at > CURRENT_TIMESTAMP) """ - with self._config.provide_connection() as conn: - with conn.cursor(row_factory=dict_row) as cur: - cur.execute(cast("Any", sql), (key,)) - row = cur.fetchone() - - if row is None: - return None - - if renew_for is not None and row["expires_at"] is not None: - new_expires_at = self._calculate_expires_at(renew_for) - if new_expires_at is not None: - update_sql = f""" - UPDATE {self._table_name} - SET expires_at = %s, updated_at = CURRENT_TIMESTAMP - WHERE session_id = %s - """ - conn.execute(cast("Any", update_sql), (new_expires_at, key)) - conn.commit() - - return bytes(row["data"]) + try: + with self._config.provide_connection() as conn: + with conn.cursor(row_factory=dict_row) as cur: + cur.execute(cast("Any", sql), (key,)) + row = cur.fetchone() + + if row is None: + return None + + if renew_for is not None and row["expires_at"] is not None: + new_expires_at = self._calculate_expires_at(renew_for) + if new_expires_at is not None: + update_sql = f""" + UPDATE {self._table_name} + SET expires_at = %s, updated_at = CURRENT_TIMESTAMP + WHERE session_id = %s + """ + conn.execute(cast("Any", update_sql), (new_expires_at, key)) + conn.commit() + + return bytes(row["data"]) + except errors.UndefinedTable: + return None def _set(self, key: str, value: "str | bytes", expires_in: "int | timedelta | None" = None) -> None: """Synchronous implementation of set.""" @@ -463,18 +490,24 @@ def _delete(self, key: str) -> None: """Synchronous implementation of delete.""" sql = f"DELETE FROM {self._table_name} WHERE session_id = %s" - with self._config.provide_connection() as conn: - conn.execute(cast("Any", sql), (key,)) - conn.commit() + try: + with self._config.provide_connection() as conn: + conn.execute(cast("Any", sql), (key,)) + conn.commit() + except errors.UndefinedTable: + return def _delete_all(self) -> None: """Synchronous implementation of delete_all.""" sql = f"DELETE FROM {self._table_name}" - with self._config.provide_connection() as conn: - conn.execute(cast("Any", sql)) - conn.commit() - self._log_delete_all() + try: + with self._config.provide_connection() as conn: + conn.execute(cast("Any", sql)) + conn.commit() + self._log_delete_all() + except errors.UndefinedTable: + return def _exists(self, key: str) -> bool: """Synchronous implementation of exists.""" @@ -484,10 +517,13 @@ def _exists(self, key: str) -> bool: AND (expires_at IS NULL OR expires_at > CURRENT_TIMESTAMP) """ - with self._config.provide_connection() as conn, conn.cursor() as cur: - cur.execute(cast("Any", sql), (key,)) - result = cur.fetchone() - return result is not None + try: + with self._config.provide_connection() as conn, conn.cursor() as cur: + cur.execute(cast("Any", sql), (key,)) + result = cur.fetchone() + return result is not None + except errors.UndefinedTable: + return False def _expires_in(self, key: str) -> "int | None": """Synchronous implementation of expires_in.""" @@ -496,47 +532,37 @@ def _expires_in(self, key: str) -> "int | None": WHERE session_id = %s """ - with self._config.provide_connection() as conn: - with conn.cursor(row_factory=dict_row) as cur: - cur.execute(cast("Any", sql), (key,)) - row = cur.fetchone() + try: + with self._config.provide_connection() as conn: + with conn.cursor(row_factory=dict_row) as cur: + cur.execute(cast("Any", sql), (key,)) + row = cur.fetchone() - if row is None or row["expires_at"] is None: - return None + if row is None or row["expires_at"] is None: + return None - expires_at = row["expires_at"] - now = datetime.now(UTC) + expires_at = row["expires_at"] + now = datetime.now(UTC) - if expires_at <= now: - return 0 + if expires_at <= now: + return 0 - delta = expires_at - now - return int(delta.total_seconds()) + delta = expires_at - now + return int(delta.total_seconds()) + except errors.UndefinedTable: + return None def _delete_expired(self) -> int: """Synchronous implementation of delete_expired.""" sql = f"DELETE FROM {self._table_name} WHERE expires_at <= CURRENT_TIMESTAMP" - with self._config.provide_connection() as conn, conn.cursor() as cur: - cur.execute(cast("Any", sql)) - conn.commit() - count = cur.rowcount if cur.rowcount and cur.rowcount > 0 else 0 - if count > 0: - self._log_delete_expired(count) - return count - - -def _postgres_litestar_tuning(config: Any) -> "tuple[int, float, float]": - settings = cast("dict[str, Any]", config.extension_config.get("litestar", {})) - fillfactor = settings.get("fillfactor", 80) - if not isinstance(fillfactor, int) or isinstance(fillfactor, bool) or fillfactor not in range(10, 101): - msg = "extension_config['litestar']['fillfactor'] must be an integer from 10 to 100" - raise ValueError(msg) - scales: list[float] = [] - for key, default in (("autovacuum_vacuum_scale_factor", 0.05), ("autovacuum_analyze_scale_factor", 0.02)): - value = settings.get(key, default) - if not isinstance(value, (int, float)) or isinstance(value, bool) or not 0 <= float(value) <= 1: - msg = f"extension_config['litestar']['{key}'] must be a number from 0 to 1" - raise ValueError(msg) - scales.append(float(value)) - return fillfactor, scales[0], scales[1] + try: + with self._config.provide_connection() as conn, conn.cursor() as cur: + cur.execute(cast("Any", sql)) + conn.commit() + count = cur.rowcount if cur.rowcount and cur.rowcount > 0 else 0 + if count > 0: + self._log_delete_expired(count) + return count + except errors.UndefinedTable: + return 0 diff --git a/sqlspec/adapters/psycopg/type_converter.py b/sqlspec/adapters/psycopg/type_converter.py index 5870ef344..8d0a8018b 100644 --- a/sqlspec/adapters/psycopg/type_converter.py +++ b/sqlspec/adapters/psycopg/type_converter.py @@ -23,7 +23,10 @@ logger = get_logger(__name__) -_pgvector_psycopg: Any | None = import_optional("pgvector.psycopg") + +def _pgvector_psycopg_module() -> Any: + """Return the ``pgvector.psycopg`` module when installed, otherwise ``None``.""" + return import_optional("pgvector.psycopg") def register_pgvector_sync(connection: "Connection[Any]") -> None: @@ -35,7 +38,7 @@ def register_pgvector_sync(connection: "Connection[Any]") -> None: Args: connection: Psycopg sync connection. """ - pgvector_psycopg = _pgvector_psycopg + pgvector_psycopg = _pgvector_psycopg_module() if pgvector_psycopg is None: return @@ -58,7 +61,7 @@ async def register_pgvector_async(connection: "AsyncConnection[Any]") -> None: Args: connection: Psycopg async connection. """ - pgvector_psycopg = _pgvector_psycopg + pgvector_psycopg = _pgvector_psycopg_module() if pgvector_psycopg is None: return diff --git a/sqlspec/adapters/pymssql/_typing.py b/sqlspec/adapters/pymssql/_typing.py index 7a1ec10d7..ebc8084f7 100644 --- a/sqlspec/adapters/pymssql/_typing.py +++ b/sqlspec/adapters/pymssql/_typing.py @@ -5,18 +5,13 @@ """ import contextlib +import importlib from typing import TYPE_CHECKING, Any -import pymssql as _pymssql # pyright: ignore[reportMissingTypeStubs] -from pymssql import Connection as _PymssqlConnection # pyright: ignore[reportMissingTypeStubs] -from pymssql import Cursor as _PymssqlRawCursor # pyright: ignore[reportMissingTypeStubs] +import pymssql as PYMSSQL_MODULE # noqa: N812 +from pymssql import Connection, Cursor from pymssql import Error as PymssqlError -from sqlspec.utils.module_loader import import_optional_attr - -PYMSSQL_MODULE = _pymssql -_DRIVER_CLS: Any = None - if TYPE_CHECKING: from collections.abc import Callable from types import TracebackType @@ -27,13 +22,12 @@ from sqlspec.adapters.pymssql.driver import PymssqlDriver from sqlspec.core import StatementConfig - PymssqlConnection: TypeAlias = _PymssqlConnection - PymssqlRawCursor: TypeAlias = _PymssqlRawCursor - -if not TYPE_CHECKING: + PymssqlConnection: TypeAlias = Connection + PymssqlRawCursor: TypeAlias = Cursor +else: PymssqlQueryParams = Any - PymssqlConnection = _PymssqlConnection - PymssqlRawCursor = _PymssqlRawCursor + PymssqlConnection = Connection + PymssqlRawCursor = Cursor __all__ = ( "PYMSSQL_MODULE", @@ -76,6 +70,9 @@ def __exit__(self, *_: Any) -> None: self.cursor.close() +_DRIVER_CLS: "type[PymssqlDriver] | None" = None + + class PymssqlSessionContext: """Sync context manager for pymssql sessions.""" @@ -83,6 +80,7 @@ class PymssqlSessionContext: "_acquire_connection", "_connection", "_driver", + "_driver_cls", "_driver_features", "_prepare_driver", "_release_connection", @@ -96,21 +94,27 @@ def __init__( statement_config: "StatementConfig", driver_features: "dict[str, Any]", prepare_driver: "Callable[[PymssqlDriver], PymssqlDriver]", + driver_cls: "type[PymssqlDriver] | None" = None, ) -> None: self._acquire_connection = acquire_connection self._release_connection = release_connection self._statement_config = statement_config self._driver_features = driver_features self._prepare_driver = prepare_driver + self._driver_cls = driver_cls self._connection: Any = None self._driver: PymssqlDriver | None = None def __enter__(self) -> "PymssqlDriver": global _DRIVER_CLS - if _DRIVER_CLS is None: - _DRIVER_CLS = import_optional_attr("sqlspec.adapters.pymssql.driver", "PymssqlDriver") + driver_cls = self._driver_cls or _DRIVER_CLS + if driver_cls is None: + module = importlib.import_module("sqlspec.adapters.pymssql.driver") + _DRIVER_CLS = module.PymssqlDriver + driver_cls = _DRIVER_CLS + assert driver_cls is not None self._connection = self._acquire_connection() - self._driver = _DRIVER_CLS( + self._driver = driver_cls( connection=self._connection, statement_config=self._statement_config, driver_features=self._driver_features ) assert self._driver is not None diff --git a/sqlspec/adapters/pymssql/adk/store.py b/sqlspec/adapters/pymssql/adk/store.py index 2a2974b59..71040934a 100644 --- a/sqlspec/adapters/pymssql/adk/store.py +++ b/sqlspec/adapters/pymssql/adk/store.py @@ -49,6 +49,8 @@ class PymssqlADKStore(BaseSyncADKStore["PymssqlConfig"]): connector_name: ClassVar[str] = "pymssql" __slots__ = ("_json_column_type", "_native_json") + _max_table_name_parts = 2 + def __init__(self, config: "PymssqlConfig") -> None: super().__init__(config) adk_config = _adk_config(config) @@ -433,6 +435,8 @@ class PymssqlADKMemoryStore(BaseSyncADKMemoryStore["PymssqlConfig"]): __slots__ = ("_json_column_type", "_native_json") + _max_table_name_parts = 2 + def __init__(self, config: "PymssqlConfig") -> None: super().__init__(config) adk_config = _adk_config(config) diff --git a/sqlspec/adapters/pymssql/config.py b/sqlspec/adapters/pymssql/config.py index 7491a74ef..529f93f2f 100644 --- a/sqlspec/adapters/pymssql/config.py +++ b/sqlspec/adapters/pymssql/config.py @@ -89,6 +89,10 @@ class PymssqlDriverFeatures(TypedDict): """Enable durable table-backed SQLSpec event queue support.""" events_backend: NotRequired[Literal["poll_queue"]] """Event channel backend (``poll_queue`` is the sole SQL Server backend).""" + enable_pooling: NotRequired[bool] + """Enable driver-level connection pooling (defaults to ``True``).""" + enable_uuid_conversion: NotRequired[bool] + """Enable automatic UUID string-to-object conversions.""" class PymssqlConnectionContext(SyncPoolConnectionContext): diff --git a/sqlspec/adapters/pymssql/core.py b/sqlspec/adapters/pymssql/core.py index f6aafdfde..2feefddb1 100644 --- a/sqlspec/adapters/pymssql/core.py +++ b/sqlspec/adapters/pymssql/core.py @@ -19,6 +19,7 @@ SerializationConflictError, SQLParsingError, SQLSpecError, + TransactionError, UniqueViolationError, ) from sqlspec.utils.serializers import from_json, to_json @@ -58,9 +59,19 @@ _ERROR_NUMBER_PATTERN: Final[re.Pattern[str]] = re.compile(r"(?:\(([-]?\d+)(?:,|\))|\bMsg\s+([-]?\d+)\b)") _MSSQL_CONSTRAINT_547: Final[int] = 547 _COLUMN_CACHE_MAX_SIZE: Final[int] = 256 +_ODBC_SQLSTATE_MAP: Final[dict[str, tuple[type[SQLSpecError], str]]] = { + "08S01": (DatabaseConnectionError, "communication link failure"), + "08001": (DatabaseConnectionError, "unable to connect to data source"), + "HYT00": (QueryTimeoutError, "query timeout expired"), + "HY008": (QueryTimeoutError, "operation canceled"), +} _ERROR_CODE_MAPPING: Final[dict[int, tuple[type[SQLSpecError], str]]] = { + 220: (DataError, "arithmetic overflow error"), + 258: (QueryTimeoutError, "timeout waiting for latch"), 2601: (UniqueViolationError, "unique constraint violation"), 2627: (UniqueViolationError, "unique constraint violation"), + 3617: (QueryTimeoutError, "query canceled by user or attention event"), + 3930: (TransactionError, "current transaction cannot be committed"), 515: (NotNullViolationError, "not-null constraint violation"), 18456: (PermissionDeniedError, "permission denied"), 229: (PermissionDeniedError, "permission denied"), @@ -85,14 +96,20 @@ 8115: (DataError, "arithmetic overflow error"), 8152: (DataError, "string or binary data would be truncated"), 1105: (OperationalError, "operational error"), + 20003: (QueryTimeoutError, "query timeout expired (FreeTDS)"), + 20002: (DatabaseConnectionError, "adaptive Server connection failed (FreeTDS)"), + 20009: (DatabaseConnectionError, "unable to connect: host unreachable or server down (FreeTDS)"), + 20047: (DatabaseConnectionError, "connection closed or broken (FreeTDS)"), } def quote_bracket_identifier(identifier: str) -> str: """Quote a T-SQL identifier with square brackets.""" cleaned = identifier.strip() - if (cleaned.startswith("[") and cleaned.endswith("]")) or (cleaned.startswith('"') and cleaned.endswith('"')): - cleaned = cleaned[1:-1].replace("]]", "]").replace('""', '"') + if cleaned.startswith("[") and cleaned.endswith("]"): + cleaned = cleaned[1:-1].replace("]]", "]") + elif cleaned.startswith('"') and cleaned.endswith('"'): + cleaned = cleaned[1:-1].replace('""', '"') return f"[{cleaned.replace(']', ']]')}]" @@ -122,13 +139,13 @@ def format_identifier(identifier: str) -> str: msg = "Table name must not be empty" raise SQLSpecError(msg) parts = split_qualified_identifier(cleaned, quote_chars='"', allow_bracket_quotes=True) - return ".".join(quote_bracket_identifier(part) for part in parts) + return ".".join([quote_bracket_identifier(part) for part in parts]) def build_insert_statement(table: str, columns: "list[str]") -> str: """Build a pymssql-compatible INSERT statement.""" - column_clause = ", ".join(quote_bracket_identifier(column) for column in columns) - placeholders = ", ".join("%s" for _ in columns) + column_clause = ", ".join([quote_bracket_identifier(column) for column in columns]) + placeholders = ", ".join(["%s"] * len(columns)) return f"INSERT INTO {format_identifier(table)} ({column_clause}) VALUES ({placeholders})" @@ -137,7 +154,13 @@ def escape_literal_percent(sql: str, parameters: Any, validator: "ParameterValid if not parameters or "%" not in sql: return sql keep = {info.position for info in validator.extract_parameters(sql)} - return "".join("%%" if char == "%" and index not in keep else char for index, char in enumerate(sql)) + result: list[str] = [] + for index, char in enumerate(sql): + if char == "%" and index not in keep: + result.append("%%") + else: + result.append(char) + return "".join(result) def normalize_execute_parameters(parameters: Any) -> Any: @@ -192,6 +215,9 @@ def apply_driver_features( features: dict[str, Any] = dict(driver_features) if driver_features else {} json_serializer = features.setdefault("json_serializer", to_json) json_deserializer = features.setdefault("json_deserializer", from_json) + features.setdefault("enable_pooling", True) + features.setdefault("events_backend", "poll_queue") + features.setdefault("enable_uuid_conversion", False) if json_serializer is not None and ( json_serializer is not to_json @@ -206,8 +232,32 @@ def apply_driver_features( return statement_config, features +_SQLSTATE_LENGTH: Final[int] = 5 + + +def extract_sqlstate(error: BaseException) -> str | None: + """Extract an ODBC SQLSTATE code from a driver error.""" + val = getattr(error, "sqlstate", None) + if isinstance(val, str) and val: + return val.strip().upper() + if hasattr(error, "args") and len(error.args) > 1: + second = error.args[1] + if isinstance(second, str) and len(second.strip()) == _SQLSTATE_LENGTH: + return second.strip().upper() + message = str(error).upper() + for code in _ODBC_SQLSTATE_MAP: + if code in message: + return code + return None + + def create_mapped_exception(error: BaseException, *, logger: "Logger | None" = None) -> SQLSpecError: """Map a SQL Server exception to SQLSpec's exception hierarchy.""" + sqlstate = extract_sqlstate(error) + if sqlstate is not None and sqlstate in _ODBC_SQLSTATE_MAP: + error_class, description = _ODBC_SQLSTATE_MAP[sqlstate] + return error_class(f"SQL Server SQLSTATE {sqlstate}: {description}. Original error: {error}") + error_number = extract_error_number(error) if error_number == _MSSQL_CONSTRAINT_547: @@ -230,6 +280,8 @@ def create_mapped_exception(error: BaseException, *, logger: "Logger | None" = N return constraint_exception exc_name = type(error).__name__ + if exc_name == "InterfaceError": + return DatabaseConnectionError(f"SQL Server interface error. Original error: {error}") if exc_name == "IntegrityError": return IntegrityError(f"SQL Server integrity error. Original error: {error}") if exc_name == "OperationalError": diff --git a/sqlspec/adapters/pymssql/data_dictionary.py b/sqlspec/adapters/pymssql/data_dictionary.py index a3bcfaabd..9e8ada99e 100644 --- a/sqlspec/adapters/pymssql/data_dictionary.py +++ b/sqlspec/adapters/pymssql/data_dictionary.py @@ -44,7 +44,7 @@ from sqlspec.adapters.pymssql.driver import PymssqlDriver from sqlspec.core import SQL - from sqlspec.data_dictionary._types import DialectConfig, MetadataCapabilityProfile + from sqlspec.data_dictionary import DialectConfig, MetadataCapabilityProfile __all__ = ("MssqlVersionInfo", "PymssqlSyncDataDictionary") @@ -277,6 +277,51 @@ def get_constraints( ) return MetadataResult("constraints", capability=capability, items=tuple(rows)) + def get_schemas(self, driver: "PymssqlDriver") -> MetadataResult: + """Get schema metadata.""" + capability = self.get_metadata_capabilities(driver, ("schemas",)).get("schemas") + self._log_schema_introspect(driver, schema_name=None, table_name=None, operation="schemas") + rows = driver.select(self.get_domain_query("schemas", "list")) + return MetadataResult("schemas", capability=capability, items=tuple(rows)) + + def get_objects(self, driver: "PymssqlDriver", schema: str | None = None) -> MetadataResult: + """Get database object metadata.""" + schema_name = self.resolve_connection_schema(driver, schema) + capability = self.get_metadata_capabilities(driver, ("objects",)).get("objects") + self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="objects") + rows = driver.select(self.get_domain_query("objects", "by_schema"), schema_name=schema_name) + return MetadataResult("objects", capability=capability, items=tuple(rows)) + + def get_table_details(self, driver: "PymssqlDriver", table: str, schema: str | None = None) -> MetadataResult: + """Get rich table metadata.""" + schema_name = self.resolve_connection_schema(driver, schema) + capability = self.get_metadata_capabilities(driver, ("tables",)).get("tables") + self._log_schema_introspect(driver, schema_name=schema_name, table_name=table, operation="table_details") + rows = driver.select( + self.get_domain_query("tables", "details_by_table"), schema_name=schema_name, table_name=table + ) + return MetadataResult("tables", capability=capability, items=tuple(rows)) + + def get_routines(self, driver: "PymssqlDriver", schema: str | None = None) -> MetadataResult: + """Get routine metadata.""" + schema_name = self.resolve_connection_schema(driver, schema) + capability = self.get_metadata_capabilities(driver, ("routines",)).get("routines") + self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="routines") + rows = driver.select(self.get_domain_query("routines", "by_schema"), schema_name=schema_name) + return MetadataResult("routines", capability=capability, items=tuple(rows)) + + def get_privileges( + self, driver: "PymssqlDriver", object_name: str | None = None, schema: str | None = None + ) -> MetadataResult: + """Get privilege metadata.""" + schema_name = self.resolve_connection_schema(driver, schema) + capability = self.get_metadata_capabilities(driver, ("privileges",)).get("privileges") + self._log_schema_introspect(driver, schema_name=schema_name, table_name=object_name, operation="privileges") + rows = driver.select(self.get_domain_query("permissions", "by_schema"), schema_name=schema_name) + if object_name is not None: + rows = [r for r in rows if _row_value(r, "object_name") == object_name] + return MetadataResult("privileges", capability=capability, items=tuple(rows)) + def get_views(self, driver: "PymssqlDriver", schema: str | None = None) -> MetadataResult: """Get view metadata for a schema.""" schema_name = self.resolve_connection_schema(driver, schema) @@ -297,11 +342,12 @@ def get_dependencies( self, driver: "PymssqlDriver", object_name: str | None = None, schema: str | None = None ) -> MetadataResult: """Get SQL expression dependency metadata for a schema.""" - _ = object_name schema_name = self.resolve_connection_schema(driver, schema) capability = self.get_metadata_capabilities(driver, ("dependencies",)).get("dependencies") - self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="dependencies") + self._log_schema_introspect(driver, schema_name=schema_name, table_name=object_name, operation="dependencies") rows = driver.select(self.get_domain_query("dependencies", "by_schema"), schema_name=schema_name) + if object_name is not None: + rows = [r for r in rows if _row_value(r, "referencing_object_name") == object_name] return MetadataResult("dependencies", capability=capability, items=tuple(rows)) def get_system_metadata( diff --git a/sqlspec/adapters/pymssql/driver.py b/sqlspec/adapters/pymssql/driver.py index db5cfef23..a431ff9d1 100644 --- a/sqlspec/adapters/pymssql/driver.py +++ b/sqlspec/adapters/pymssql/driver.py @@ -215,6 +215,8 @@ def begin(self) -> None: ``BEGIN TRANSACTION``, and the matching ``commit()`` or ``rollback()`` ends it with T-SQL because pymssql ignores those calls under autocommit. """ + if self._transaction_active: + return try: explicit = bool(self.connection.autocommit_state) if explicit: @@ -451,13 +453,16 @@ def _execute_cache_hit( self, sql: str, params: "tuple[Any, ...] | list[Any] | dict[str, Any]", cached: "CachedQuery" ) -> "SQLResult": normalized_params = normalize_execute_parameters(params) if isinstance(params, list) else params - if not normalized_params or ( + param_count = len(normalized_params) if isinstance(normalized_params, (list, tuple)) else -1 + has_literal_percent = ( "%" in cached.compiled_sql + and (param_count < 0 or cached.compiled_sql.count("%") != param_count) and escape_literal_percent( cached.compiled_sql, normalized_params, self.statement_config.parameter_validator ) != cached.compiled_sql - ): + ) + if not normalized_params or has_literal_percent: statement = self._cached_statement( sql, normalized_params, diff --git a/sqlspec/adapters/pymssql/litestar/store.py b/sqlspec/adapters/pymssql/litestar/store.py index f4360a99c..ce5fbae28 100644 --- a/sqlspec/adapters/pymssql/litestar/store.py +++ b/sqlspec/adapters/pymssql/litestar/store.py @@ -133,7 +133,7 @@ def _set(self, key: str, value: "str | bytes", expires_in: "int | timedelta | No expires_at = self._calculate_expires_at(expires_in) table_ref = _table_ref(self._table_name) sql = f""" - MERGE INTO {table_ref} AS target + MERGE INTO {table_ref} WITH (HOLDLOCK) AS target USING (SELECT %s AS session_id, CONVERT(VARBINARY(MAX), %s) AS data, %s AS expires_at) AS src ON target.session_id = src.session_id WHEN MATCHED THEN diff --git a/sqlspec/adapters/pymssql/pool.py b/sqlspec/adapters/pymssql/pool.py index 5bc87a61b..3b61c5350 100644 --- a/sqlspec/adapters/pymssql/pool.py +++ b/sqlspec/adapters/pymssql/pool.py @@ -94,14 +94,18 @@ def get_connection(self) -> "Generator[PymssqlConnection, None, None]": Yields: A thread-local database connection. """ + depth = getattr(self._thread_local, "depth", 0) connection = self.acquire() + self._thread_local.depth = depth + 1 try: yield connection except Exception: - with contextlib.suppress(Exception): - self._close_thread_connection() + if getattr(self._thread_local, "depth", 1) <= 1: + with contextlib.suppress(Exception): + self._close_thread_connection() raise finally: + self._thread_local.depth = max(0, getattr(self._thread_local, "depth", 1) - 1) self.release(connection) def close(self) -> None: diff --git a/sqlspec/adapters/pymysql/_typing.py b/sqlspec/adapters/pymysql/_typing.py index c7571d2b7..5458aae6e 100644 --- a/sqlspec/adapters/pymysql/_typing.py +++ b/sqlspec/adapters/pymysql/_typing.py @@ -38,8 +38,7 @@ class PyMysqlServerStatusProtocol(Protocol): PyMysqlConnect: TypeAlias = type["PyMysqlConnection"] PyMysqlFieldType: TypeAlias = PyMysqlFieldTypeProtocol PyMysqlServerStatus: TypeAlias = PyMysqlServerStatusProtocol - -if not TYPE_CHECKING: +else: PyMysqlConnect = pymysql.connect PyMysqlFieldType = _PYMYSQL_FIELD_TYPE PyMysqlServerStatus = _PYMYSQL_SERVER_STATUS diff --git a/sqlspec/adapters/pymysql/adk/store.py b/sqlspec/adapters/pymysql/adk/store.py index 88fd93f6b..592c8930c 100644 --- a/sqlspec/adapters/pymysql/adk/store.py +++ b/sqlspec/adapters/pymysql/adk/store.py @@ -431,14 +431,19 @@ def _search_entries( """ params = (*scope_params, f"%{query}%", limit_value) - with self._config.provide_connection() as conn: - cursor = conn.cursor() - try: - cursor.execute(sql, params) - rows = cursor.fetchall() - columns = [col[0] for col in cursor.description or []] - finally: - cursor.close() + try: + with self._config.provide_connection() as conn: + cursor = conn.cursor() + try: + cursor.execute(sql, params) + rows = cursor.fetchall() + columns = [col[0] for col in cursor.description or []] + finally: + cursor.close() + except PyMysqlMySQLError as exc: + if _is_mysql_table_missing(exc): + return [] + raise records: list[StoredMemory] = [] for row in rows: @@ -458,14 +463,19 @@ def _delete_entries_by_session(self, session_id: str) -> int: raise RuntimeError(msg) sql = f"DELETE FROM {self._memory_table} WHERE session_id = %s" - with self._config.provide_connection() as conn: - cursor = conn.cursor() - try: - cursor.execute(sql, (session_id,)) - conn.commit() - return cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0 - finally: - cursor.close() + try: + with self._config.provide_connection() as conn: + cursor = conn.cursor() + try: + cursor.execute(sql, (session_id,)) + conn.commit() + return cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0 + finally: + cursor.close() + except PyMysqlMySQLError as exc: + if _is_mysql_table_missing(exc): + return 0 + raise def _delete_entries_older_than(self, days: int, app_name: "str | None" = None, scope: "str | None" = None) -> int: if not self._enabled: @@ -482,14 +492,19 @@ def _delete_entries_older_than(self, days: int, app_name: "str | None" = None, s params.append(scope) where_sql = " AND ".join(clauses) sql = f"DELETE FROM {self._memory_table} WHERE {where_sql}" - with self._config.provide_connection() as conn: - cursor = conn.cursor() - try: - cursor.execute(sql, tuple(params)) - conn.commit() - return cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0 - finally: - cursor.close() + try: + with self._config.provide_connection() as conn: + cursor = conn.cursor() + try: + cursor.execute(sql, tuple(params)) + conn.commit() + return cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0 + finally: + cursor.close() + except PyMysqlMySQLError as exc: + if _is_mysql_table_missing(exc): + return 0 + raise def _create_tables(store: PyMysqlADKStore) -> None: diff --git a/sqlspec/adapters/pymysql/core.py b/sqlspec/adapters/pymysql/core.py index b89db729d..2d0da3c4d 100644 --- a/sqlspec/adapters/pymysql/core.py +++ b/sqlspec/adapters/pymysql/core.py @@ -3,7 +3,7 @@ import contextlib from collections.abc import Callable, Mapping, Sequence, Sized from datetime import timedelta -from typing import TYPE_CHECKING, Any, Final, Literal, cast +from typing import TYPE_CHECKING, Any, Final, Literal from sqlspec.adapters.pymysql._typing import PyMysqlSSCursor from sqlspec.core import DriverParameterProfile, ParameterStyle, StatementConfig, build_statement_config_from_profile @@ -17,6 +17,7 @@ ForeignKeyViolationError, IntegrityError, NotNullViolationError, + OperationalError, PermissionDeniedError, QueryTimeoutError, SQLParsingError, @@ -28,13 +29,13 @@ from sqlspec.utils.serializers import from_json, to_json from sqlspec.utils.text import quote_backtick_identifier, split_qualified_identifier from sqlspec.utils.type_converters import build_uuid_coercions -from sqlspec.utils.type_guards import has_cursor_metadata, has_lastrowid, has_rowcount, has_sqlstate, has_type_code +from sqlspec.utils.type_guards import has_lastrowid, has_rowcount, has_sqlstate, has_type_code if TYPE_CHECKING: from sqlspec.core.parameters import ParameterValidator __all__ = ( - "PymysqlStreamSource", + "PyMysqlStreamSource", "apply_driver_features", "build_insert_statement", "build_load_data_statement", @@ -45,8 +46,6 @@ "collect_stream_rows", "create_mapped_exception", "default_statement_config", - "detect_json_columns", - "detect_json_columns_from_description", "driver_profile", "encode_records_for_local_infile", "escape_literal_percent", @@ -54,7 +53,6 @@ "normalize_execute_many_parameters", "normalize_execute_parameters", "normalize_lastrowid", - "resolve_column_names", "resolve_many_rowcount", "resolve_row_plan", "resolve_rowcount", @@ -80,7 +78,10 @@ MYSQL_SYNTAX_ERROR_MIN: Final[int] = 1064 MYSQL_SYNTAX_ERROR_MAX_EXCLUSIVE: Final[int] = 1100 -_MYSQL_MIGRATION_ERROR_CODES: Final[frozenset[int]] = frozenset((1061, 1091)) +_MYSQL_DDL_ERROR_DISPATCH: dict[int, tuple[type[SQLSpecError], str]] = { + 1061: (OperationalError, "duplicate key name"), + 1091: (OperationalError, "cannot drop column or key"), +} _MYSQL_SQLSTATE_EXACT_DISPATCH: Final[dict[str, tuple[type[SQLSpecError], str]]] = { "23505": (UniqueViolationError, "unique constraint violation"), "23503": (ForeignKeyViolationError, "foreign key constraint violation"), @@ -145,10 +146,19 @@ } -class PymysqlStreamSource: +class PyMysqlStreamSource: """Compiled chunk source streaming dict rows from a PyMySQL unbuffered ``SSCursor``.""" - __slots__ = ("_chunk_size", "_cursor", "_driver", "_json_type_codes", "_parameters", "_row_plan", "_sql") + __slots__ = ( + "_chunk_size", + "_cursor", + "_deserializer", + "_driver", + "_json_type_codes", + "_parameters", + "_row_plan", + "_sql", + ) def __init__(self, driver: Any, sql: str, parameters: Any, chunk_size: int, json_type_codes: "set[int]") -> None: self._driver = driver @@ -158,36 +168,35 @@ def __init__(self, driver: Any, sql: str, parameters: Any, chunk_size: int, json self._cursor: Any = None self._json_type_codes = json_type_codes self._row_plan: tuple[list[str], list[int] | None] | None = None + deserializer = getattr(driver, "_json_deserializer", None) + if deserializer is None and hasattr(driver, "driver_features") and isinstance(driver.driver_features, dict): + deserializer = driver.driver_features.get("json_deserializer") + self._deserializer: Callable[[Any], Any] = deserializer or from_json def start(self) -> None: handler = self._driver.handle_database_exceptions() - with handler: - try: - cursor = self._driver.connection.cursor(PyMysqlSSCursor) - self._cursor = cursor - cursor.execute(self._sql, normalize_execute_parameters(self._parameters)) - self._row_plan = resolve_row_plan(self._cursor.description, self._json_type_codes) - except BaseException: - self.close(error=True) - raise + self._driver._run_with_exception_handler(handler, self._start) self._driver._check_pending_exception(handler) + def _start(self) -> None: + try: + cursor = self._driver.connection.cursor(PyMysqlSSCursor) + self._cursor = cursor + cursor.execute(self._sql, normalize_execute_parameters(self._parameters)) + self._row_plan = resolve_row_plan(self._cursor.description, self._json_type_codes) + except BaseException: + self.close(error=True) + raise + def fetch_chunk(self) -> "list[dict[str, Any]]": handler = self._driver.handle_database_exceptions() - rows: list[Any] = [] - with handler: - rows = self._cursor.fetchmany(self._chunk_size) + rows = self._driver._run_with_exception_handler(handler, self._cursor.fetchmany, self._chunk_size) self._driver._check_pending_exception(handler) if not rows: return [] if self._row_plan is None: self._row_plan = resolve_row_plan(self._cursor.description, self._json_type_codes) - deserializer = cast( - "Callable[[Any], Any]", - getattr(self._driver, "_json_deserializer", None) - or self._driver.driver_features.get("json_deserializer", from_json), - ) - return collect_stream_rows(rows, self._row_plan, deserializer) + return collect_stream_rows(rows, self._row_plan, self._deserializer) def close(self, error: bool = False) -> None: cursor = self._cursor @@ -199,7 +208,7 @@ def close(self, error: bool = False) -> None: def build_profile() -> "DriverParameterProfile": """Create the PyMySQL driver parameter profile.""" - coercions: dict[type, Callable[[Any], Any]] = {bool: bool_to_int, **build_uuid_coercions()} + coercions: dict[type, Callable[[Any], Any]] = {bool: int, **build_uuid_coercions()} return DriverParameterProfile( name="PyMySQL", default_style=ParameterStyle.QMARK, @@ -254,14 +263,14 @@ def format_identifier(identifier: str) -> str: msg = "Table name must not be empty" raise SQLSpecError(msg) parts = split_qualified_identifier(cleaned, quote_chars="`", allow_bracket_quotes=False) - formatted = ".".join(quote_backtick_identifier(part) for part in parts) + formatted = ".".join([quote_backtick_identifier(part) for part in parts]) return formatted or quote_backtick_identifier(cleaned) def build_insert_statement(table: str, columns: "list[str]") -> str: """Construct an INSERT statement with positional placeholders.""" - column_clause = ", ".join(quote_backtick_identifier(column) for column in columns) - placeholders = ", ".join("%s" for _ in columns) + column_clause = ", ".join([quote_backtick_identifier(column) for column in columns]) + placeholders = ", ".join(["%s" for _ in columns]) return f"INSERT INTO {format_identifier(table)} ({column_clause}) VALUES ({placeholders})" @@ -294,7 +303,11 @@ def can_use_local_infile_records(arrow_table: Any, records: "Sequence[Sequence[A all_safe_types = False if all_safe_types: return True - return not any(isinstance(value, (bytes, bytearray, memoryview, timedelta)) for row in records for value in row) + for row in records: + for value in row: + if isinstance(value, (bytes, bytearray, memoryview, timedelta)): + return False + return True def encode_records_for_local_infile(records: "list[tuple[Any, ...]]") -> bytes: @@ -338,13 +351,6 @@ def normalize_execute_many_parameters(parameters: Any) -> Any: return parameters -def resolve_column_names(description: "Sequence[Any] | None") -> "list[str]": - """Extract column names from cursor description.""" - if not description: - return [] - return [col[0] for col in description] - - def resolve_row_plan( description: "Sequence[Any] | None", json_type_codes: "set[int]" ) -> "tuple[list[str], list[int] | None]": @@ -354,7 +360,7 @@ def resolve_row_plan( column_names: list[str] = [] if not json_type_codes: - column_names.extend(column[0] for column in description) + column_names.extend([column[0] for column in description]) return column_names, None json_indexes: list[int] = [] @@ -370,24 +376,6 @@ def resolve_row_plan( return column_names, json_indexes or None -def detect_json_columns_from_description( - description: "Sequence[Any] | None", json_type_codes: "set[int]" -) -> "list[int]": - """Identify JSON column indexes from pre-fetched cursor description metadata.""" - return resolve_row_plan(description, json_type_codes)[1] or [] - - -def detect_json_columns( - cursor: Any, json_type_codes: "set[int]", description: "Sequence[Any] | None" = None -) -> "list[int]": - """Identify JSON column indexes from cursor metadata.""" - if description is None: - if not has_cursor_metadata(cursor): - return [] - description = cursor.description - return detect_json_columns_from_description(description, json_type_codes) - - def collect_stream_rows( rows: "Sequence[Any]", row_plan: "tuple[list[str], list[int] | None]", @@ -469,13 +457,9 @@ def normalize_lastrowid(cursor: Any) -> int | None: return last_id if isinstance(last_id, int) else None -def bool_to_int(value: bool) -> int: - """Coerce boolean values to integer for MySQL parameter compatibility.""" - return int(value) - - -def create_mapped_exception(error: Any, *, logger: Any | None = None) -> "SQLSpecError | bool": +def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpecError: """Map MySQL driver exceptions to SQLSpec errors.""" + _ = logger error_code = getattr(error, "errno", None) if error_code is None and isinstance(error, Exception) and error.args: value = error.args[0] @@ -484,10 +468,9 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> "SQLSpe sqlstate = error.sqlstate if has_sqlstate(error) else None sqlstate_prefix = sqlstate[:2] if isinstance(sqlstate, str) and sqlstate else None - if error_code in _MYSQL_MIGRATION_ERROR_CODES: - if logger is not None: - logger.warning("MySQL expected migration error (ignoring): %s", error) - return True + dispatch = _MYSQL_DDL_ERROR_DISPATCH.get(error_code) if error_code is not None else None + if dispatch is not None: + return _create_mysql_error(error, sqlstate, error_code, dispatch[0], dispatch[1]) dispatch = _MYSQL_SQLSTATE_EXACT_DISPATCH.get(sqlstate) if sqlstate is not None else None if dispatch is not None: @@ -543,7 +526,9 @@ def escape_literal_percent(sql: str, parameters: Any, validator: "ParameterValid if not parameters or "%" not in sql: return sql keep = {info.position for info in validator.extract_parameters(sql)} - return "".join("%%" if char == "%" and index not in keep else char for index, char in enumerate(sql)) + if len(keep) == sql.count("%"): + return sql + return "".join(["%%" if char == "%" and index not in keep else char for index, char in enumerate(sql)]) def build_load_data_statement( @@ -553,10 +538,10 @@ def build_load_data_statement( table_sql = format_identifier(table) if escape_percent: table_sql = table_sql.replace("%", "%%") - column_list = ", ".join( + column_list = ", ".join([ quote_backtick_identifier(column).replace("%", "%%") if escape_percent else quote_backtick_identifier(column) for column in columns - ) + ]) return ( f"LOAD DATA LOCAL INFILE {placeholder} INTO TABLE {table_sql} " "CHARACTER SET utf8mb4 FIELDS TERMINATED BY '\\t' ESCAPED BY '\\\\' " diff --git a/sqlspec/adapters/pymysql/driver.py b/sqlspec/adapters/pymysql/driver.py index 47f15d481..a6f09e05d 100644 --- a/sqlspec/adapters/pymysql/driver.py +++ b/sqlspec/adapters/pymysql/driver.py @@ -15,7 +15,7 @@ PyMysqlSessionContext, ) from sqlspec.adapters.pymysql.core import ( - PymysqlStreamSource, + PyMysqlStreamSource, build_insert_statement, build_load_data_statement, can_use_local_infile_records, @@ -83,10 +83,7 @@ def _handle_exception(self, exc_type: "type[BaseException] | None", exc_val: "Ba if exc_type is None: return False if issubclass(exc_type, PyMysqlMySQLError): - result = create_mapped_exception(exc_val, logger=logger) - if result is True: - return True - self.pending_exception = cast("Exception", result) + self.pending_exception = create_mapped_exception(exc_val) return True return False @@ -122,8 +119,8 @@ def data_dictionary(self) -> "PyMysqlDataDictionary": return self._data_dictionary def dispatch_execute(self, cursor: Any, statement: "SQL") -> "ExecutionResult": - sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) - sql = escape_literal_percent(sql, prepared_parameters, self.statement_config.parameter_validator) + sql, prepared_parameters = self._compiled_sql(statement, statement.statement_config) + sql = escape_literal_percent(sql, prepared_parameters, statement.statement_config.parameter_validator) cursor.execute(sql, normalize_execute_parameters(prepared_parameters)) if statement.returns_rows() or cursor.description is not None: @@ -150,15 +147,13 @@ def dispatch_execute(self, cursor: Any, statement: "SQL") -> "ExecutionResult": return self.create_execution_result(cursor, rowcount_override=affected_rows, last_inserted_id=last_id) def dispatch_execute_many(self, cursor: Any, statement: "SQL") -> "ExecutionResult": - sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) + sql, prepared_parameters = self._compiled_sql(statement, statement.statement_config) + validator = statement.statement_config.parameter_validator match = PYMYSQL_INSERT_VALUES_PATTERN.match(sql) if match: - sql = ( - escape_literal_percent(match.group(1), prepared_parameters, self.statement_config.parameter_validator) - + sql[match.end(1) :] - ) + sql = escape_literal_percent(match.group(1), prepared_parameters, validator) + sql[match.end(1) :] else: - sql = escape_literal_percent(sql, prepared_parameters, self.statement_config.parameter_validator) + sql = escape_literal_percent(sql, prepared_parameters, validator) prepared_parameters = normalize_execute_many_parameters(prepared_parameters) parameter_count = len(prepared_parameters) if isinstance(prepared_parameters, Sized) else None @@ -168,47 +163,34 @@ def dispatch_execute_many(self, cursor: Any, statement: "SQL") -> "ExecutionResu return self.create_execution_result(cursor, rowcount_override=affected_rows, is_many_result=True) def dispatch_execute_script(self, cursor: Any, statement: "SQL") -> "ExecutionResult": - sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) - statements = self.split_script_statements(sql, statement.statement_config, strip_trailing_semicolon=True) - - successful_count = 0 - last_cursor = cursor + """Execute SQL script with statement splitting and parameter handling. - for stmt in statements: - stmt_sql = escape_literal_percent(stmt, prepared_parameters, self.statement_config.parameter_validator) - cursor.execute(stmt_sql, normalize_execute_parameters(prepared_parameters)) - successful_count += 1 + Args: + cursor: PyMySQL cursor object + statement: SQL script to execute - return self.create_execution_result( - last_cursor, statement_count=len(statements), successful_statements=successful_count, is_script_result=True - ) + Returns: + ExecutionResult: Script execution results with statement count + """ + return super().dispatch_execute_script(cursor, statement) def begin(self) -> None: try: self.connection.begin() except PyMysqlMySQLError as exc: - mapped = create_mapped_exception(exc, logger=logger) - if isinstance(mapped, bool): - return - raise mapped from exc + raise create_mapped_exception(exc) from exc def commit(self) -> None: try: self.connection.commit() except PyMysqlMySQLError as exc: - mapped = create_mapped_exception(exc, logger=logger) - if isinstance(mapped, bool): - return - raise mapped from exc + raise create_mapped_exception(exc) from exc def rollback(self) -> None: try: self.connection.rollback() except PyMysqlMySQLError as exc: - mapped = create_mapped_exception(exc, logger=logger) - if isinstance(mapped, bool): - return - raise mapped from exc + raise create_mapped_exception(exc) from exc def with_cursor(self, connection: "PyMysqlConnection") -> "PyMysqlCursor": return PyMysqlCursor(connection) @@ -217,9 +199,9 @@ def dispatch_select_stream(self, statement: "SQL", chunk_size: int) -> "SyncRowS """Return a native PyMySQL row stream backed by an unbuffered ``SSCursor``.""" if not statement.returns_rows(): return None - sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) - sql = escape_literal_percent(sql, prepared_parameters, self.statement_config.parameter_validator) - return SyncRowStream(PymysqlStreamSource(self, sql, prepared_parameters, chunk_size, PYMYSQL_JSON_TYPE_CODES)) + sql, prepared_parameters = self._compiled_sql(statement, statement.statement_config) + sql = escape_literal_percent(sql, prepared_parameters, statement.statement_config.parameter_validator) + return SyncRowStream(PyMysqlStreamSource(self, sql, prepared_parameters, chunk_size, PYMYSQL_JSON_TYPE_CODES)) def handle_database_exceptions(self) -> "PyMysqlExceptionHandler": return PyMysqlExceptionHandler() @@ -284,17 +266,17 @@ def load_from_arrow( if overwrite: statement = f"TRUNCATE TABLE {format_identifier(table)}" exc_handler = self.handle_database_exceptions() - with exc_handler, self.with_cursor(self.connection) as cursor: - cursor.execute(statement) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + self._run_with_exception_handler(exc_handler, self._execute_cursor_statement, statement) + self._check_pending_exception(exc_handler) columns, records = self._arrow_table_to_rows(arrow_table) if records: needs_preparation = self._arrow_rows_need_preparation(arrow_table) - use_infile = bool( - self.driver_features.get("enable_local_infile_bulk_load") - ) and can_use_local_infile_records(arrow_table, records) + use_infile = ( + bool(self.driver_features.get("enable_local_infile_bulk_load")) + and not needs_preparation + and can_use_local_infile_records(arrow_table, records) + ) if use_infile: payload = encode_records_for_local_infile(records) fd, tmp_name = tempfile.mkstemp(suffix=".tsv") @@ -303,10 +285,8 @@ def load_from_arrow( tmp.write(payload) load_sql = build_load_data_statement(table, columns, escape_percent=True) exc_handler = self.handle_database_exceptions() - with exc_handler, self.with_cursor(self.connection) as cursor: - cursor.execute(load_sql, (tmp_name,)) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + self._run_with_exception_handler(exc_handler, self._execute_cursor_statement, load_sql, (tmp_name,)) + self._check_pending_exception(exc_handler) finally: Path(tmp_name).unlink(missing_ok=True) else: @@ -317,29 +297,26 @@ def load_from_arrow( else records ) exc_handler = self.handle_database_exceptions() - with exc_handler, self.with_cursor(self.connection) as cursor: - cursor.executemany(insert_sql, cast("Any", prepared_records)) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + self._run_with_exception_handler( + exc_handler, self._executemany_cursor_statement, insert_sql, prepared_records + ) + self._check_pending_exception(exc_handler) telemetry_payload = self._ingest_telemetry(arrow_table) telemetry_payload["destination"] = table self._attach_partition_telemetry(telemetry_payload, partitioner) return self._storage_job(telemetry_payload, telemetry) - def load_from_storage( - self, - table: str, - source: "StorageDestination", - *, - file_format: "StorageFormat", - partitioner: "dict[str, object] | None" = None, - overwrite: bool = False, - **kwargs: Any, - ) -> "StorageBridgeJob": - return super().load_from_storage( - table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs - ) + def _execute_cursor_statement(self, statement: str, parameters: Any = None) -> None: + with self.with_cursor(self.connection) as cursor: + if parameters is None: + cursor.execute(statement) + else: + cursor.execute(statement, parameters) + + def _executemany_cursor_statement(self, statement: str, parameters: Any) -> None: + with self.with_cursor(self.connection) as cursor: + cursor.executemany(statement, parameters) def collect_rows(self, cursor: Any, fetched: "list[Any]") -> "tuple[list[Any], list[str], int]": """Collect PyMySQL rows for the direct execution path.""" diff --git a/sqlspec/adapters/pymysql/litestar/store.py b/sqlspec/adapters/pymysql/litestar/store.py index a10d61ba1..01aa98397 100644 --- a/sqlspec/adapters/pymysql/litestar/store.py +++ b/sqlspec/adapters/pymysql/litestar/store.py @@ -87,10 +87,7 @@ def _table_ddl(self) -> str: """ def _drop_table_sql(self) -> "list[str]": - return [ - f"DROP INDEX idx_{self._table_name}_expires_at ON {self._table_name}", - f"DROP TABLE IF EXISTS {self._table_name}", - ] + return [f"DROP TABLE IF EXISTS {self._table_name}"] def _create_table(self) -> None: sql = self._table_ddl() @@ -142,10 +139,10 @@ def _set(self, key: str, value: "str | bytes", expires_in: "int | timedelta | No sql = f""" INSERT INTO {self._table_name} (session_id, data, expires_at) - VALUES (%s, %s, %s) AS new + VALUES (%s, %s, %s) ON DUPLICATE KEY UPDATE - data = new.data, - expires_at = new.expires_at, + data = VALUES(data), + expires_at = VALUES(expires_at), updated_at = UTC_TIMESTAMP(6) """ diff --git a/sqlspec/adapters/pymysql/pool.py b/sqlspec/adapters/pymysql/pool.py index b55b18448..d205d13a3 100644 --- a/sqlspec/adapters/pymysql/pool.py +++ b/sqlspec/adapters/pymysql/pool.py @@ -30,6 +30,7 @@ class PyMysqlConnectionPool: "_connection_factory", "_connection_parameters", "_connection_registry", + "_connection_threads", "_generation", "_health_check_interval", "_on_connection_create", @@ -60,6 +61,7 @@ def __init__( self._connection_factory = connection_factory self._thread_local = threading.local() self._connection_registry: set[PyMysqlConnection] = set() + self._connection_threads: dict[PyMysqlConnection, threading.Thread] = {} self._checked_out = 0 self._generation = 0 self._registry_lock = threading.Lock() @@ -107,13 +109,7 @@ def get_connection(self) -> "Generator[PyMysqlConnection, None, None]": connection = self.acquire() try: yield connection - except BaseException: - with self._registry_lock: - self._checked_out = max(0, self._checked_out - 1) - with contextlib.suppress(Exception): - self._close_thread_connection() - raise - else: + finally: self.release(connection) def close(self) -> None: @@ -122,6 +118,7 @@ def close(self) -> None: with self._registry_lock: orphaned = list(self._connection_registry) self._connection_registry.clear() + self._connection_threads.clear() self._generation += 1 for connection in orphaned: with contextlib.suppress(Exception): @@ -161,6 +158,7 @@ def release(self, connection: PyMysqlConnection) -> None: def size(self) -> int: """Report total active connections managed by this pool.""" + self._prune_dead_thread_connections() with self._registry_lock: return len(self._connection_registry) @@ -168,11 +166,25 @@ def checked_out(self) -> int: with self._registry_lock: return self._checked_out + def _prune_dead_thread_connections(self) -> None: + dead_connections: list[PyMysqlConnection] = [] + with self._registry_lock: + for conn, owner_thread in list(self._connection_threads.items()): + if not owner_thread.is_alive(): + self._connection_registry.discard(conn) + self._connection_threads.pop(conn, None) + dead_connections.append(conn) + for conn in dead_connections: + with contextlib.suppress(Exception): + conn.close() + def _create_connection(self) -> PyMysqlConnection: + self._prune_dead_thread_connections() connection = self.new_connection() with self._registry_lock: self._connection_registry.add(connection) + self._connection_threads[connection] = threading.current_thread() return connection @@ -238,6 +250,7 @@ def _retire_connection(self, connection: PyMysqlConnection) -> None: """Close a pool-owned connection and drop it from the shutdown registry.""" with self._registry_lock: self._connection_registry.discard(connection) + self._connection_threads.pop(connection, None) with contextlib.suppress(Exception): connection.close() diff --git a/sqlspec/adapters/spanner/_typing.py b/sqlspec/adapters/spanner/_typing.py index 3f09222df..f188773ec 100644 --- a/sqlspec/adapters/spanner/_typing.py +++ b/sqlspec/adapters/spanner/_typing.py @@ -5,104 +5,70 @@ """ from importlib import import_module -from typing import TYPE_CHECKING, Any, cast - -from sqlspec.typing import import_optional_attr - - -class _UnavailableSpannerGoogleAPICallError(Exception): - """Fallback Spanner API exception when google-api-core is unavailable.""" - - -class _UnavailableSpannerTransaction: - """Fallback Spanner transaction class when google-cloud-spanner is unavailable.""" - - __slots__ = () - - -class _UnavailableSpannerAsyncTransaction: - """Fallback Spanner async transaction class when google-cloud-spanner is unavailable.""" - - __slots__ = () - +from typing import TYPE_CHECKING, Any, TypeAlias, cast + +from google.api_core import exceptions as spanner_exceptions +from google.api_core.client_info import ClientInfo as SpannerClientInfo +from google.api_core.client_options import ClientOptions as SpannerClientOptions +from google.api_core.exceptions import GoogleAPICallError as SpannerGoogleAPICallError +from google.api_core.exceptions import GoogleAPIError as SpannerGoogleAPIError +from google.api_core.exceptions import NotFound as SpannerNotFound +from google.api_core.retry import Retry as SpannerRetry +from google.auth.credentials import Credentials as SpannerCredentials +from google.cloud.spanner_admin_database_v1.types import DatabaseDialect as SpannerDatabaseDialect +from google.cloud.spanner_admin_database_v1.types import EncryptionConfig as SpannerEncryptionConfig +from google.cloud.spanner_v1 import Client as SpannerClient +from google.cloud.spanner_v1 import DirectedReadOptions as SpannerDirectedReadOptions +from google.cloud.spanner_v1 import ExecuteSqlRequest as SpannerExecuteSqlRequest +from google.cloud.spanner_v1 import RequestOptions as SpannerRequestOptions +from google.cloud.spanner_v1 import param_types as spanner_param_types +from google.cloud.spanner_v1._async.batch import Batch as SpannerAsyncBatch +from google.cloud.spanner_v1._async.batch import MutationGroups as SpannerAsyncMutationGroups +from google.cloud.spanner_v1._async.client import Client as SpannerAsyncClient +from google.cloud.spanner_v1._async.database import Database as SpannerAsyncDatabase +from google.cloud.spanner_v1._async.database import SnapshotCheckout as SpannerAsyncSnapshotCheckout +from google.cloud.spanner_v1._async.database_sessions_manager import TransactionType as SpannerAsyncTransactionType +from google.cloud.spanner_v1._async.pool import AbstractSessionPool as SpannerAsyncAbstractSessionPool +from google.cloud.spanner_v1._async.pool import BurstyPool as SpannerAsyncBurstyPool +from google.cloud.spanner_v1._async.pool import FixedSizePool as SpannerAsyncFixedSizePool +from google.cloud.spanner_v1._async.pool import PingingPool as SpannerAsyncPingingPool +from google.cloud.spanner_v1._async.pool import TransactionPingingPool as SpannerAsyncTransactionPingingPool +from google.cloud.spanner_v1._async.session import Session as SpannerAsyncSession +from google.cloud.spanner_v1._async.snapshot import Snapshot as SpannerAsyncSnapshot +from google.cloud.spanner_v1._async.streamed import StreamedResultSet as SpannerAsyncStreamedResultSet +from google.cloud.spanner_v1._async.transaction import Transaction as SpannerAsyncTransaction +from google.cloud.spanner_v1.batch import Batch as SpannerBatch +from google.cloud.spanner_v1.batch import MutationGroups as SpannerMutationGroups +from google.cloud.spanner_v1.data_types import Interval as SpannerInterval +from google.cloud.spanner_v1.data_types import JsonObject as SpannerJsonObject +from google.cloud.spanner_v1.database import Database as SpannerDatabase +from google.cloud.spanner_v1.database import SnapshotCheckout +from google.cloud.spanner_v1.database_sessions_manager import TransactionType as SpannerTransactionType +from google.cloud.spanner_v1.pool import AbstractSessionPool as SpannerAbstractSessionPool +from google.cloud.spanner_v1.pool import BurstyPool as SpannerBurstyPool +from google.cloud.spanner_v1.pool import FixedSizePool as SpannerFixedSizePool +from google.cloud.spanner_v1.pool import PingingPool as SpannerPingingPool +from google.cloud.spanner_v1.pool import TransactionPingingPool as SpannerTransactionPingingPool +from google.cloud.spanner_v1.transaction import DefaultTransactionOptions as SpannerDefaultTransactionOptions +from google.cloud.spanner_v1.transaction import Transaction as SpannerTransaction +from google.cloud.spanner_v1.types.type import TypeCode as SpannerTypeCode if TYPE_CHECKING: from collections.abc import Awaitable, Callable from types import TracebackType - from typing import TypeAlias - - from google.api_core import exceptions as spanner_exceptions - from google.api_core.client_info import ClientInfo as SpannerClientInfo - from google.api_core.client_options import ClientOptions as SpannerClientOptions - from google.api_core.exceptions import GoogleAPICallError as _SpannerGoogleAPICallError - from google.api_core.exceptions import NotFound as SpannerNotFound - from google.api_core.retry import Retry as SpannerRetry - from google.auth.credentials import Credentials as SpannerCredentials - from google.cloud.spanner_admin_database_v1.types import DatabaseDialect as SpannerDatabaseDialect - from google.cloud.spanner_admin_database_v1.types import EncryptionConfig as SpannerEncryptionConfig - from google.cloud.spanner_v1 import Client as SpannerClient - from google.cloud.spanner_v1 import DirectedReadOptions as SpannerDirectedReadOptions - from google.cloud.spanner_v1 import ExecuteSqlRequest as SpannerExecuteSqlRequest - from google.cloud.spanner_v1 import RequestOptions as SpannerRequestOptions - from google.cloud.spanner_v1 import param_types as spanner_param_types - from google.cloud.spanner_v1._async.batch import Batch as SpannerAsyncBatch - from google.cloud.spanner_v1._async.batch import MutationGroups as SpannerAsyncMutationGroups - from google.cloud.spanner_v1._async.client import Client as SpannerAsyncClient - from google.cloud.spanner_v1._async.database import Database as SpannerAsyncDatabase - from google.cloud.spanner_v1._async.database import SnapshotCheckout as SpannerAsyncSnapshotCheckout - from google.cloud.spanner_v1._async.database_sessions_manager import TransactionType as SpannerAsyncTransactionType - from google.cloud.spanner_v1._async.pool import AbstractSessionPool as SpannerAsyncAbstractSessionPool - from google.cloud.spanner_v1._async.pool import BurstyPool as SpannerAsyncBurstyPool - from google.cloud.spanner_v1._async.pool import FixedSizePool as SpannerAsyncFixedSizePool - from google.cloud.spanner_v1._async.pool import PingingPool as SpannerAsyncPingingPool - from google.cloud.spanner_v1._async.pool import TransactionPingingPool as SpannerAsyncTransactionPingingPool - from google.cloud.spanner_v1._async.session import Session as SpannerAsyncSession - from google.cloud.spanner_v1._async.snapshot import Snapshot as SpannerAsyncSnapshot - from google.cloud.spanner_v1._async.streamed import StreamedResultSet as SpannerAsyncStreamedResultSet - from google.cloud.spanner_v1._async.transaction import Transaction as _SpannerAsyncTransaction - from google.cloud.spanner_v1.batch import Batch as SpannerBatch - from google.cloud.spanner_v1.batch import MutationGroups as SpannerMutationGroups - from google.cloud.spanner_v1.data_types import Interval as SpannerInterval - from google.cloud.spanner_v1.data_types import JsonObject as SpannerJsonObject - from google.cloud.spanner_v1.database import Database as SpannerDatabase - from google.cloud.spanner_v1.database import SnapshotCheckout - from google.cloud.spanner_v1.database_sessions_manager import TransactionType as SpannerTransactionType - from google.cloud.spanner_v1.pool import AbstractSessionPool as SpannerAbstractSessionPool - from google.cloud.spanner_v1.pool import BurstyPool as SpannerBurstyPool - from google.cloud.spanner_v1.pool import FixedSizePool as SpannerFixedSizePool - from google.cloud.spanner_v1.pool import PingingPool as SpannerPingingPool - from google.cloud.spanner_v1.pool import TransactionPingingPool as SpannerTransactionPingingPool + from google.cloud.spanner_v1.snapshot import Snapshot - from google.cloud.spanner_v1.transaction import DefaultTransactionOptions as SpannerDefaultTransactionOptions - from google.cloud.spanner_v1.transaction import Transaction as _SpannerTransaction - from google.cloud.spanner_v1.types.type import TypeCode as SpannerTypeCode from sqlspec.adapters.spanner.driver import SpannerAsyncDriver, SpannerSyncDriver from sqlspec.core import StatementConfig - SpannerSyncConnection: TypeAlias = Snapshot | SnapshotCheckout | _SpannerTransaction | SpannerBatch + SpannerSyncConnection: TypeAlias = Snapshot | SnapshotCheckout | SpannerTransaction | SpannerBatch SpannerAsyncConnection: TypeAlias = ( - SpannerAsyncSnapshot | SpannerAsyncSnapshotCheckout | _SpannerAsyncTransaction | SpannerAsyncBatch + SpannerAsyncSnapshot | SpannerAsyncSnapshotCheckout | SpannerAsyncTransaction | SpannerAsyncBatch ) - SpannerGoogleAPICallError: TypeAlias = _SpannerGoogleAPICallError - SpannerTransaction: TypeAlias = _SpannerTransaction - SpannerAsyncTransaction: TypeAlias = _SpannerAsyncTransaction - - -if not TYPE_CHECKING: +else: SpannerSyncConnection = Any SpannerAsyncConnection = Any - SpannerGoogleAPICallError = ( - import_optional_attr("google.api_core.exceptions", "GoogleAPICallError") - or _UnavailableSpannerGoogleAPICallError - ) - SpannerTransaction = ( - import_optional_attr("google.cloud.spanner_v1.transaction", "Transaction") or _UnavailableSpannerTransaction - ) - SpannerAsyncTransaction = ( - import_optional_attr("google.cloud.spanner_v1._async.transaction", "Transaction") - or _UnavailableSpannerAsyncTransaction - ) __all__ = ( @@ -138,6 +104,7 @@ class _UnavailableSpannerAsyncTransaction: "SpannerExecuteSqlRequest", "SpannerFixedSizePool", "SpannerGoogleAPICallError", + "SpannerGoogleAPIError", "SpannerInterval", "SpannerJsonObject", "SpannerMutationGroups", @@ -297,63 +264,6 @@ async def __aexit__( return None -_LAZY_DRIVER_EXPORTS: dict[str, tuple[str, str]] = { - "SpannerAsyncTransactionType": ("google.cloud.spanner_v1._async.database_sessions_manager", "TransactionType"), - "SpannerInterval": ("google.cloud.spanner_v1.data_types", "Interval"), - "SpannerAsyncAbstractSessionPool": ("google.cloud.spanner_v1._async.pool", "AbstractSessionPool"), - "SpannerAsyncBatch": ("google.cloud.spanner_v1._async.batch", "Batch"), - "SpannerAsyncBurstyPool": ("google.cloud.spanner_v1._async.pool", "BurstyPool"), - "SpannerAsyncClient": ("google.cloud.spanner_v1._async.client", "Client"), - "SpannerAsyncDatabase": ("google.cloud.spanner_v1._async.database", "Database"), - "SpannerAsyncFixedSizePool": ("google.cloud.spanner_v1._async.pool", "FixedSizePool"), - "SpannerAsyncMutationGroups": ("google.cloud.spanner_v1._async.batch", "MutationGroups"), - "SpannerAsyncPingingPool": ("google.cloud.spanner_v1._async.pool", "PingingPool"), - "SpannerAsyncSession": ("google.cloud.spanner_v1._async.session", "Session"), - "SpannerAsyncSnapshot": ("google.cloud.spanner_v1._async.snapshot", "Snapshot"), - "SpannerAsyncStreamedResultSet": ("google.cloud.spanner_v1._async.streamed", "StreamedResultSet"), - "SpannerAsyncTransactionPingingPool": ("google.cloud.spanner_v1._async.pool", "TransactionPingingPool"), - "SpannerBatch": ("google.cloud.spanner_v1.batch", "Batch"), - "SpannerBurstyPool": ("google.cloud.spanner_v1.pool", "BurstyPool"), - "SpannerFixedSizePool": ("google.cloud.spanner_v1.pool", "FixedSizePool"), - "SpannerMutationGroups": ("google.cloud.spanner_v1.batch", "MutationGroups"), - "SpannerPingingPool": ("google.cloud.spanner_v1.pool", "PingingPool"), - "SpannerTransactionPingingPool": ("google.cloud.spanner_v1.pool", "TransactionPingingPool"), - "spanner_exceptions": ("google.api_core", "exceptions"), - "SpannerNotFound": ("google.api_core.exceptions", "NotFound"), - "SpannerClient": ("google.cloud.spanner_v1", "Client"), - "spanner_param_types": ("google.cloud.spanner_v1", "param_types"), - "SpannerJsonObject": ("google.cloud.spanner_v1.data_types", "JsonObject"), - "SpannerTransactionType": ("google.cloud.spanner_v1.database_sessions_manager", "TransactionType"), - "SpannerTypeCode": ("google.cloud.spanner_v1.types.type", "TypeCode"), - "SpannerClientInfo": ("google.api_core.client_info", "ClientInfo"), - "SpannerClientOptions": ("google.api_core.client_options", "ClientOptions"), - "SpannerRetry": ("google.api_core.retry", "Retry"), - "SpannerCredentials": ("google.auth.credentials", "Credentials"), - "SpannerDatabaseDialect": ("google.cloud.spanner_admin_database_v1.types", "DatabaseDialect"), - "SpannerEncryptionConfig": ("google.cloud.spanner_admin_database_v1.types", "EncryptionConfig"), - "SpannerDirectedReadOptions": ("google.cloud.spanner_v1", "DirectedReadOptions"), - "SpannerExecuteSqlRequest": ("google.cloud.spanner_v1", "ExecuteSqlRequest"), - "SpannerRequestOptions": ("google.cloud.spanner_v1", "RequestOptions"), - "SpannerDatabase": ("google.cloud.spanner_v1.database", "Database"), - "SpannerAbstractSessionPool": ("google.cloud.spanner_v1.pool", "AbstractSessionPool"), - "SpannerDefaultTransactionOptions": ("google.cloud.spanner_v1.transaction", "DefaultTransactionOptions"), -} - - -def __getattr__(name: str) -> Any: - """Resolve optional driver symbols only when a consumer requests them.""" - target = _LAZY_DRIVER_EXPORTS.get(name) - if target is None: - msg = f"module {__name__!r} has no attribute {name!r}" - raise AttributeError(msg) - module_name, attribute = target - value = getattr(import_module(module_name), attribute, None) - if value is None: - msg = f"Cannot import {attribute!r} from {module_name!r}" - raise ImportError(msg) - return value - - def _get_sync_driver_type() -> "type[SpannerSyncDriver]": driver_module = import_module("sqlspec.adapters.spanner.driver") return cast("type[SpannerSyncDriver]", driver_module.SpannerSyncDriver) diff --git a/sqlspec/adapters/spanner/adk/artifact_store.py b/sqlspec/adapters/spanner/adk/artifact_store.py index c73d82f20..732532a8f 100644 --- a/sqlspec/adapters/spanner/adk/artifact_store.py +++ b/sqlspec/adapters/spanner/adk/artifact_store.py @@ -2,25 +2,23 @@ from typing import TYPE_CHECKING, Any, ClassVar, cast -from sqlspec.adapters.spanner._typing import SpannerGoogleAPICallError +from sqlspec.adapters.spanner._typing import SpannerGoogleAPICallError, spanner_param_types from sqlspec.adapters.spanner._typing import SpannerNotFound as NotFound -from sqlspec.adapters.spanner.adk.store import ( +from sqlspec.adapters.spanner.config import SpannerAsyncConfig, SpannerSyncConfig +from sqlspec.adapters.spanner.core import ( _DDL_TIMEOUT_SECONDS, - SPANNER_PARAM_TYPES, _adk_config, + _convert_json_row_value, _extract_spanner_json_codecs, _is_spanner_table_missing, _json_param_type, + _run_read, _run_read_async, _spanner_row_deletion_policy, _SpannerAsyncUpdateJob, _SpannerReadProtocol, _SpannerSyncUpdateJob, _SpannerWriteTransactionProtocol, -) -from sqlspec.adapters.spanner.config import SpannerAsyncConfig, SpannerSyncConfig -from sqlspec.adapters.spanner.core import ( - _convert_json_row_value, create_mapped_exception, execute_ddl_async, execute_ddl_sync, @@ -44,6 +42,8 @@ __all__ = ("USER_SCOPED_SESSION_ID", "SpannerAsyncADKArtifactStore", "SpannerSyncADKArtifactStore") +SPANNER_PARAM_TYPES = spanner_param_types + USER_SCOPED_SESSION_ID = "" @@ -465,12 +465,7 @@ def _existing_tables(self) -> "set[str]": def _run_read( self, sql: str, params: "dict[str, Any] | None" = None, types: "dict[str, Any] | None" = None ) -> "list[Any]": - try: - with self._config.provide_connection(transaction=False) as snapshot: - reader = cast("_SpannerReadProtocol", snapshot) - return list(reader.execute_sql(sql, params=params, param_types=types)) - except SpannerGoogleAPICallError as exc: - raise create_mapped_exception(exc) from exc + return _run_read(self._config, sql, params, types) class SpannerAsyncADKArtifactStore(BaseAsyncADKArtifactStore[SpannerAsyncConfig], _SpannerADKArtifactStoreMixin): diff --git a/sqlspec/adapters/spanner/config.py b/sqlspec/adapters/spanner/config.py index 7a03f6966..5828fde9d 100644 --- a/sqlspec/adapters/spanner/config.py +++ b/sqlspec/adapters/spanner/config.py @@ -46,6 +46,7 @@ from sqlspec.exceptions import ImproperConfigurationError from sqlspec.extensions.events import EventRuntimeHints from sqlspec.utils.config_tools import normalize_connection_config +from sqlspec.utils.module_loader import ensure_pyarrow if TYPE_CHECKING: from collections.abc import Awaitable, Callable, Mapping @@ -67,6 +68,7 @@ from sqlspec.config import ExtensionConfigs from sqlspec.core import StatementConfig from sqlspec.observability import ObservabilityConfig + from sqlspec.storage import StorageCapabilities __all__ = ( "SpannerAsyncConfig", @@ -249,6 +251,8 @@ class SpannerDriverFeatures(TypedDict): """Event queue backend name (``'poll_queue'``).""" enable_batch_write_api: "NotRequired[bool]" """Route ``load_from_arrow`` through Spanner ``batch_write`` mutation groups instead of single-transaction mutations (defaults to ``False``).""" + on_connection_create: "NotRequired[Callable[[Any], Any | None] | None]" + """Callback invoked after a new Spanner connection (Snapshot or Transaction) is created.""" def build_connection_config( @@ -287,15 +291,29 @@ def __enter__(self) -> "SpannerSyncConnection": try: txn = self._session.transaction() txn.__enter__() - except Exception: - manager.put_session(self._session) + except BaseException: + with contextlib.suppress(BaseException): + manager.put_session(self._session) self._session = None raise self._connection = cast("SpannerSyncConnection", txn) - return self._connection - - self._session = cast("Any", database).snapshot(multi_use=True) - self._connection = cast("SpannerSyncConnection", self._session.__enter__()) + else: + self._session = cast("Any", database).snapshot(multi_use=True) + self._connection = cast("SpannerSyncConnection", self._session.__enter__()) + try: + hook = getattr(self._config, "_user_connection_hook", None) or getattr( + self._config, "on_connection_create", None + ) + if hook: + modified = hook(self._connection) + if modified is not None: + self._connection = cast("SpannerSyncConnection", modified) + except BaseException as exc: + with contextlib.suppress(BaseException): + self.__exit__(type(exc), exc, exc.__traceback__) + self._connection = None + self._session = None + raise return self._connection def __exit__( @@ -360,8 +378,8 @@ class SpannerSyncConfig(SyncDatabaseConfig["SpannerSyncConnection", "AbstractSes "type[SpannerSyncConnection]", SpannerSyncConnection ) supports_transactional_ddl: ClassVar[bool] = False - supports_native_arrow_export: ClassVar[bool] = True - supports_native_arrow_import: ClassVar[bool] = True + supports_native_arrow_export: ClassVar[bool] = False + supports_native_arrow_import: ClassVar[bool] = False supports_native_parquet_export: ClassVar[bool] = False supports_native_parquet_import: ClassVar[bool] = False type_coercion_capabilities: ClassVar[TypeCoercionCapabilities] = TypeCoercionCapabilities( @@ -386,11 +404,16 @@ def __init__( bind_key: "str | None" = None, extension_config: "ExtensionConfigs | None" = None, observability_config: "ObservabilityConfig | None" = None, + on_connection_create: "Callable[[Any], Any | None] | None" = None, **kwargs: Any, ) -> None: self.connection_config, raw_driver_features = _prepare_settings( connection_config, driver_features, default_pool_type=FixedSizePool ) + user_connection_hook = on_connection_create or cast( + "Callable[[Any], Any | None] | None", raw_driver_features.pop("on_connection_create", None) + ) + self._user_connection_hook = user_connection_hook statement_config, processed_features = apply_driver_features( statement_config or default_statement_config, raw_driver_features ) @@ -410,6 +433,14 @@ def __init__( self._client: Client | None = None self._database: Database | None = None + def _build_storage_capabilities(self) -> "StorageCapabilities": + """Enable Arrow conversion and mutation ingestion when PyArrow is installed.""" + capabilities = super()._build_storage_capabilities() + arrow_ready = self._dependency_available(ensure_pyarrow) + capabilities["arrow_import_enabled"] = arrow_ready + capabilities["arrow_export_enabled"] = arrow_ready + return capabilities + def get_database(self) -> "Database": """Return the configured database, creating the pool and client on first use. @@ -678,15 +709,29 @@ async def __aenter__(self) -> "SpannerAsyncConnection": try: txn = self._session.transaction() await txn.__aenter__() - except Exception: - await manager.put_session(self._session) + except BaseException: + with contextlib.suppress(BaseException): + await manager.put_session(self._session) self._session = None raise self._connection = cast("SpannerAsyncConnection", txn) - return self._connection - - self._session = cast("Any", database).snapshot(multi_use=True) - self._connection = cast("SpannerAsyncConnection", await self._session.__aenter__()) + else: + self._session = cast("Any", database).snapshot(multi_use=True) + self._connection = cast("SpannerAsyncConnection", await self._session.__aenter__()) + try: + hook = getattr(self._config, "_user_connection_hook", None) or getattr( + self._config, "on_connection_create", None + ) + if hook: + modified = hook(self._connection) + if modified is not None: + self._connection = cast("SpannerAsyncConnection", modified) + except BaseException as exc: + with contextlib.suppress(BaseException): + await self.__aexit__(type(exc), exc, exc.__traceback__) + self._connection = None + self._session = None + raise return self._connection async def __aexit__( @@ -752,8 +797,8 @@ class SpannerAsyncConfig(AsyncDatabaseConfig["SpannerAsyncConnection", "AsyncAbs "type[SpannerAsyncConnection]", SpannerAsyncConnection ) supports_transactional_ddl: ClassVar[bool] = False - supports_native_arrow_export: ClassVar[bool] = True - supports_native_arrow_import: ClassVar[bool] = True + supports_native_arrow_export: ClassVar[bool] = False + supports_native_arrow_import: ClassVar[bool] = False supports_native_parquet_export: ClassVar[bool] = False supports_native_parquet_import: ClassVar[bool] = False type_coercion_capabilities: ClassVar[TypeCoercionCapabilities] = TypeCoercionCapabilities( @@ -778,11 +823,16 @@ def __init__( bind_key: "str | None" = None, extension_config: "ExtensionConfigs | None" = None, observability_config: "ObservabilityConfig | None" = None, + on_connection_create: "Callable[[Any], Any | None] | None" = None, **kwargs: Any, ) -> None: self.connection_config, raw_driver_features = _prepare_settings( connection_config, driver_features, default_pool_type=AsyncFixedSizePool ) + user_connection_hook = on_connection_create or cast( + "Callable[[Any], Any | None] | None", raw_driver_features.pop("on_connection_create", None) + ) + self._user_connection_hook = user_connection_hook statement_config, processed_features = apply_driver_features( statement_config or default_statement_config, raw_driver_features ) @@ -802,6 +852,14 @@ def __init__( self._client: AsyncClient | None = None self._database: AsyncDatabase | None = None + def _build_storage_capabilities(self) -> "StorageCapabilities": + """Enable Arrow conversion and mutation ingestion when PyArrow is installed.""" + capabilities = super()._build_storage_capabilities() + arrow_ready = self._dependency_available(ensure_pyarrow) + capabilities["arrow_import_enabled"] = arrow_ready + capabilities["arrow_export_enabled"] = arrow_ready + return capabilities + async def get_database(self) -> "AsyncDatabase": """Return the configured database, creating the pool and client on first use. diff --git a/sqlspec/adapters/spanner/core.py b/sqlspec/adapters/spanner/core.py index fe3455e1f..4dec2403d 100644 --- a/sqlspec/adapters/spanner/core.py +++ b/sqlspec/adapters/spanner/core.py @@ -2,9 +2,10 @@ import contextlib import re +from collections.abc import Callable, Iterable, Mapping from functools import partial from itertools import islice -from typing import TYPE_CHECKING, Any, Final, Literal, cast +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast import sqlglot from sqlglot import exp @@ -13,6 +14,7 @@ from sqlspec.adapters.spanner._typing import SpannerJsonObject as JsonObject from sqlspec.adapters.spanner._typing import SpannerTypeCode as TypeCode from sqlspec.adapters.spanner._typing import spanner_exceptions as api_exceptions +from sqlspec.adapters.spanner._typing import spanner_param_types from sqlspec.adapters.spanner.type_converter import coerce_params_for_spanner, infer_spanner_param_types from sqlspec.core import ( DriverParameterProfile, @@ -22,8 +24,11 @@ build_statement_config_from_profile, ) from sqlspec.driver import rows_to_dicts +from sqlspec.driver._exception_handler import _run_with_async_exception_handler, _run_with_sync_exception_handler from sqlspec.exceptions import ( + DataError, DeadlockError, + IntegrityError, NotFoundError, OperationalError, OperationCancelledError, @@ -38,7 +43,7 @@ from sqlspec.utils.serializers import from_json, to_json if TYPE_CHECKING: - from collections.abc import AsyncGenerator, Awaitable, Callable, Generator, Mapping, Sequence + from collections.abc import AsyncGenerator, Awaitable, Generator, Sequence from sqlspec.adapters.spanner._typing import SpannerDirectedReadOptions as DirectedReadOptions from sqlspec.adapters.spanner._typing import SpannerExecuteSqlRequest as ExecuteSqlRequest @@ -150,12 +155,24 @@ def build_execute_kwargs( return kwargs +_EXECUTE_OPTION_KEYS: Final[tuple[str, ...]] = ( + "request_options", + "query_options", + "directed_read_options", + "retry", + "timeout", + "last_statement", +) + + def pop_execute_options(kwargs: dict[str, Any]) -> "SpannerExecuteOptions | None": """Remove native execution options before processing SQL parameters.""" - if not any( - key in kwargs - for key in ("request_options", "query_options", "directed_read_options", "retry", "timeout", "last_statement") - ): + has_options = False + for key in _EXECUTE_OPTION_KEYS: + if key in kwargs: + has_options = True + break + if not has_options: return None return SpannerExecuteOptions( request_options=kwargs.pop("request_options", None), @@ -367,9 +384,14 @@ def collect_rows( Tuple of (rows, column_names). """ resolved_column_names = column_names if column_names is not None else [field.name for field in fields] - if column_plan is None and all(type(row) is tuple for row in rows): - return rows, resolved_column_names if column_plan is None: + all_tuples = True + for row in rows: + if type(row) is not tuple: + all_tuples = False + break + if all_tuples: + return rows, resolved_column_names return [tuple(row) for row in rows], resolved_column_names data: list[Any] = [] @@ -397,6 +419,21 @@ def create_arrow_data( return convert_dict_to_arrow(data, return_format=return_format) +_GRPC_STATUS_CODE_MAP: Final[dict[int, type[SQLSpecError]]] = { + 3: SQLParsingError, + 4: QueryTimeoutError, + 5: NotFoundError, + 6: UniqueViolationError, + 7: PermissionDeniedError, + 8: OperationalError, + 9: IntegrityError, + 10: DeadlockError, + 11: DataError, + 14: OperationalError, + 16: PermissionDeniedError, +} + + def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpecError: """Map Spanner exceptions to SQLSpec exceptions. @@ -416,6 +453,12 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpec A SQLSpec exception that wraps the original error """ del logger + if hasattr(error, "code") and not isinstance(error, BaseException): + status_code = error.code + if isinstance(status_code, int) and status_code in _GRPC_STATUS_CODE_MAP: + msg = str(getattr(error, "message", "batch DML failure") or "batch DML failure") + return _GRPC_STATUS_CODE_MAP[status_code](msg) + if isinstance(error, api_exceptions.AlreadyExists): return _create_spanner_error(error, UniqueViolationError, "resource already exists") @@ -433,12 +476,20 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpec if isinstance(error, api_exceptions.Aborted): return _create_spanner_error(error, DeadlockError, "transaction aborted") + if isinstance(error, api_exceptions.FailedPrecondition): + return _create_spanner_error(error, IntegrityError, "constraint violation or precondition failed") + + if isinstance(error, api_exceptions.OutOfRange): + return _create_spanner_error(error, DataError, "out of range") + if isinstance(error, api_exceptions.Cancelled): return _create_spanner_error(error, OperationCancelledError, "operation cancelled") if isinstance(error, api_exceptions.DeadlineExceeded): return _create_spanner_error(error, QueryTimeoutError, "deadline exceeded") - if isinstance(error, (api_exceptions.ServiceUnavailable, api_exceptions.TooManyRequests)): + if isinstance( + error, (api_exceptions.ServiceUnavailable, api_exceptions.TooManyRequests, api_exceptions.ResourceExhausted) + ): return _create_spanner_error(error, OperationalError, "service unavailable or rate limited") return _create_spanner_error(error, SQLSpecError, "error") @@ -864,22 +915,22 @@ def __init__( def start(self) -> None: handler = self._driver.handle_database_exceptions() - with handler: - result_set = self._driver.connection.execute_sql( - self._sql, params=self._params, param_types=self._param_types, **self._execute_kwargs - ) - self._result_set = result_set - self._row_iterator = cast("Generator[Any, None, None]", iter(result_set)) + _run_with_sync_exception_handler(handler, self._start_query) self._driver._check_pending_exception(handler) + def _start_query(self) -> None: + result_set = self._driver.connection.execute_sql( + self._sql, params=self._params, param_types=self._param_types, **self._execute_kwargs + ) + self._result_set = result_set + self._row_iterator = cast("Generator[Any, None, None]", iter(result_set)) + def fetch_chunk(self) -> "list[dict[str, Any]]": row_iterator = self._row_iterator if self._result_set is None or row_iterator is None: return [] handler = self._driver.handle_database_exceptions() - rows: list[Any] = [] - with handler: - rows = list(islice(row_iterator, self._chunk_size)) + rows: list[Any] | None = _run_with_sync_exception_handler(handler, list, islice(row_iterator, self._chunk_size)) self._driver._check_pending_exception(handler) if not rows: return [] @@ -912,29 +963,34 @@ def __init__( async def start(self) -> None: handler = self._driver.handle_database_exceptions() - async with handler: - result_set = await self._driver.connection.execute_sql( - self._sql, params=self._params, param_types=self._param_types, **self._execute_kwargs - ) - self._result_set = result_set - self._row_iterator = cast("AsyncGenerator[Any, None]", result_set.__aiter__()) + await _run_with_async_exception_handler(handler, self._start_query) self._driver._check_pending_exception(handler) + async def _start_query(self) -> None: + result_set = await self._driver.connection.execute_sql( + self._sql, params=self._params, param_types=self._param_types, **self._execute_kwargs + ) + self._result_set = result_set + self._row_iterator = cast("AsyncGenerator[Any, None]", result_set.__aiter__()) + async def fetch_chunk(self) -> "list[dict[str, Any]]": row_iterator = self._row_iterator if self._result_set is None or row_iterator is None: return [] handler = self._driver.handle_database_exceptions() - rows: list[Any] = [] - async with handler: - with contextlib.suppress(StopAsyncIteration): - while len(rows) < self._chunk_size: - rows.append(await row_iterator.__anext__()) + rows = await _run_with_async_exception_handler(handler, self._fetch_rows, row_iterator) self._driver._check_pending_exception(handler) if not rows: return [] return self._rows_to_chunk(rows) + async def _fetch_rows(self, row_iterator: "AsyncGenerator[Any, None]") -> "list[Any]": + rows: list[Any] = [] + with contextlib.suppress(StopAsyncIteration): + while len(rows) < self._chunk_size: + rows.append(await row_iterator.__anext__()) + return rows + async def close(self, error: bool = False) -> None: row_iterator = self._row_iterator self._row_iterator = None @@ -980,3 +1036,108 @@ def _create_spanner_error(error: Any, error_class: type[SQLSpecError], descripti driver_profile = build_profile() default_statement_config = build_statement_config() + +_DDL_TIMEOUT_SECONDS: Final[float] = 300.0 + + +def _json_param_type() -> Any: + """Return the SPANNER_PARAM_TYPES JSON type or fallback to STRING.""" + return getattr(spanner_param_types, "JSON", getattr(spanner_param_types, "STRING", None)) + + +def _extract_spanner_json_codecs(config: Any) -> "tuple[Callable[[Any], str], Callable[[str], Any]]": + driver_features = getattr(config, "driver_features", None) + features = driver_features if isinstance(driver_features, Mapping) else {} + serializer = features.get("json_serializer") + deserializer = features.get("json_deserializer") + return ( + cast("Callable[[Any], str]", serializer if callable(serializer) else to_json), + cast("Callable[[str], Any]", deserializer if callable(deserializer) else from_json), + ) + + +def _spanner_row_deletion_policy(adk_config: "Mapping[str, Any]", ttl_key: str, column: str) -> str: + retention = adk_config.get("retention") + if not isinstance(retention, dict): + return "" + ttl_seconds = retention.get(ttl_key) + if not isinstance(ttl_seconds, int) or ttl_seconds <= 0: + return "" + ttl_days = max(1, (ttl_seconds + 86_399) // 86_400) + return f",\nROW DELETION POLICY (OLDER_THAN({column}, INTERVAL {ttl_days} DAY))" + + +def _is_spanner_table_missing(exc: Exception, *table_names: str) -> bool: + message = str(exc).lower() + return "not found" in message and any(table_name.lower() in message for table_name in table_names) + + +class _SpannerReadProtocol(Protocol): + def execute_sql( + self, sql: str, params: "dict[str, Any] | None" = None, param_types: "dict[str, Any] | None" = None + ) -> "Iterable[Any]": ... + + +class _SpannerWriteTransactionProtocol(Protocol): + def execute_update( + self, sql: str, params: "dict[str, Any] | None" = None, param_types: "dict[str, Any] | None" = None + ) -> int: ... + + +async def _run_read_async( + database: Any, sql: str, params: "dict[str, Any] | None" = None, types: "dict[str, Any] | None" = None +) -> "list[Any]": + try: + async with cast("Any", database).snapshot() as snapshot: + result_set = await snapshot.execute_sql(sql, params=params, param_types=types) + return [row async for row in result_set] + except api_exceptions.GoogleAPICallError as exc: + raise create_mapped_exception(exc) from exc + + +def _adk_config(config: Any) -> "dict[str, Any]": + extension_config = getattr(config, "extension_config", None) + adk_config = extension_config.get("adk") if isinstance(extension_config, dict) else {} + if not isinstance(adk_config, dict): + return {} + return adk_config + + +def _run_read( + config: Any, sql: str, params: "dict[str, Any] | None" = None, types: "dict[str, Any] | None" = None +) -> "list[Any]": + try: + with config.provide_connection(transaction=False) as snapshot: + reader = cast("_SpannerReadProtocol", snapshot) + return list(reader.execute_sql(sql, params=params, param_types=types)) + except api_exceptions.GoogleAPICallError as exc: + raise create_mapped_exception(exc) from exc + + +class _SpannerSyncUpdateJob: + """Callable transaction work item for Spanner sync single DML updates.""" + + __slots__ = ("_params", "_sql", "_types") + + def __init__(self, sql: str, params: "dict[str, Any]", types: "dict[str, Any]") -> None: + self._sql = sql + self._params = params + self._types = types + + def __call__(self, transaction: Any) -> int: + writer = cast("_SpannerWriteTransactionProtocol", transaction) + return int(writer.execute_update(self._sql, params=self._params, param_types=self._types)) + + +class _SpannerAsyncUpdateJob: + """Callable async transaction work item for Spanner async single DML updates.""" + + __slots__ = ("_params", "_sql", "_types") + + def __init__(self, sql: str, params: "dict[str, Any]", types: "dict[str, Any]") -> None: + self._sql = sql + self._params = params + self._types = types + + async def __call__(self, transaction: Any) -> int: + return int(await transaction.execute_update(self._sql, params=self._params, param_types=self._types)) diff --git a/sqlspec/adapters/spanner/data_dictionary.py b/sqlspec/adapters/spanner/data_dictionary.py index c5e8e4ce3..37385387c 100644 --- a/sqlspec/adapters/spanner/data_dictionary.py +++ b/sqlspec/adapters/spanner/data_dictionary.py @@ -1,11 +1,11 @@ """Spanner metadata queries using INFORMATION_SCHEMA.""" import re -from typing import TYPE_CHECKING, Any, ClassVar, cast +from typing import TYPE_CHECKING, Any, ClassVar, Final, cast from mypy_extensions import mypyc_attr -from sqlspec.core import SQL +from sqlspec.core import SQL, TypedParameter from sqlspec.data_dictionary import ( ColumnMetadata, DDLResult, @@ -72,6 +72,15 @@ ) +_SPANNER_UNSUPPORTED_DOMAINS: Final[frozenset[str]] = frozenset({"schemas", "dependencies"}) + + +def _typed_str(val: "str | None") -> "str | TypedParameter": + if val is None: + return TypedParameter(None, str) + return val + + @mypyc_attr(allow_interpreted_subclasses=True, native_class=False) class SpannerSyncDataDictionary(SyncDataDictionaryBase): """Fetch table, column, and index metadata from Spanner.""" @@ -130,7 +139,9 @@ def get_tables(self, driver: "SpannerSyncDriver", schema: "str | None" = None) - schema_name = self.resolve_schema(schema) self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="tables") return driver.select( - self.get_query("tables", "by_schema", mode=self.mode), schema_name=schema_name, schema_type=TableMetadata + self.get_query("tables", "by_schema", mode=self.mode), + schema_name=_typed_str(schema_name), + schema_type=TableMetadata, ) def get_columns( @@ -142,15 +153,15 @@ def get_columns( self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="columns") return driver.select( self.get_query("columns", "by_schema", mode=self.mode), - schema_name=schema_name, + schema_name=_typed_str(schema_name), schema_type=ColumnMetadata, ) self._log_table_describe(driver, schema_name=schema_name, table_name=table, operation="columns") return driver.select( self.get_query("columns", "by_table", mode=self.mode), - table_name=table, - schema_name=schema_name, + table_name=_typed_str(table), + schema_name=_typed_str(schema_name), schema_type=ColumnMetadata, ) @@ -163,15 +174,15 @@ def get_indexes( self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="indexes") return driver.select( self.get_query("indexes", "by_schema", mode=self.mode), - schema_name=schema_name, + schema_name=_typed_str(schema_name), schema_type=IndexMetadata, ) self._log_table_describe(driver, schema_name=schema_name, table_name=table, operation="indexes") return driver.select( self.get_query("indexes", "by_table", mode=self.mode), - table_name=table, - schema_name=schema_name, + table_name=_typed_str(table), + schema_name=_typed_str(schema_name), schema_type=IndexMetadata, ) @@ -184,14 +195,14 @@ def get_foreign_keys( self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="foreign_keys") return driver.select( self.get_query("foreign_keys", "by_schema", mode=self.mode), - schema_name=schema_name, + schema_name=_typed_str(schema_name), schema_type=ForeignKeyMetadata, ) self._log_table_describe(driver, schema_name=schema_name, table_name=table, operation="foreign_keys") return driver.select( self.get_query("foreign_keys", "by_table", mode=self.mode), - table_name=table, - schema_name=schema_name, + table_name=_typed_str(table), + schema_name=_typed_str(schema_name), schema_type=ForeignKeyMetadata, ) @@ -203,12 +214,16 @@ def get_constraints( if table is None: self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="constraints") rows = driver.select( - self.get_query("constraints", "by_schema", mode=self.mode), schema_name=schema_name, table_name=None + self.get_query("constraints", "by_schema", mode=self.mode), + schema_name=_typed_str(schema_name), + table_name=_typed_str(None), ) else: self._log_table_describe(driver, schema_name=schema_name, table_name=table, operation="constraints") rows = driver.select( - self.get_query("constraints", "by_table", mode=self.mode), table_name=table, schema_name=schema_name + self.get_query("constraints", "by_table", mode=self.mode), + table_name=_typed_str(table), + schema_name=_typed_str(schema_name), ) return self._build_domain_result("constraints", rows) @@ -382,7 +397,9 @@ async def get_tables(self, driver: "SpannerAsyncDriver", schema: "str | None" = schema_name = self.resolve_schema(schema) self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="tables") return await driver.select( - self.get_query("tables", "by_schema", mode=self.mode), schema_name=schema_name, schema_type=TableMetadata + self.get_query("tables", "by_schema", mode=self.mode), + schema_name=_typed_str(schema_name), + schema_type=TableMetadata, ) async def get_columns( @@ -394,15 +411,15 @@ async def get_columns( self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="columns") return await driver.select( self.get_query("columns", "by_schema", mode=self.mode), - schema_name=schema_name, + schema_name=_typed_str(schema_name), schema_type=ColumnMetadata, ) self._log_table_describe(driver, schema_name=schema_name, table_name=table, operation="columns") return await driver.select( self.get_query("columns", "by_table", mode=self.mode), - table_name=table, - schema_name=schema_name, + table_name=_typed_str(table), + schema_name=_typed_str(schema_name), schema_type=ColumnMetadata, ) @@ -415,15 +432,15 @@ async def get_indexes( self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="indexes") return await driver.select( self.get_query("indexes", "by_schema", mode=self.mode), - schema_name=schema_name, + schema_name=_typed_str(schema_name), schema_type=IndexMetadata, ) self._log_table_describe(driver, schema_name=schema_name, table_name=table, operation="indexes") return await driver.select( self.get_query("indexes", "by_table", mode=self.mode), - table_name=table, - schema_name=schema_name, + table_name=_typed_str(table), + schema_name=_typed_str(schema_name), schema_type=IndexMetadata, ) @@ -436,14 +453,14 @@ async def get_foreign_keys( self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="foreign_keys") return await driver.select( self.get_query("foreign_keys", "by_schema", mode=self.mode), - schema_name=schema_name, + schema_name=_typed_str(schema_name), schema_type=ForeignKeyMetadata, ) self._log_table_describe(driver, schema_name=schema_name, table_name=table, operation="foreign_keys") return await driver.select( self.get_query("foreign_keys", "by_table", mode=self.mode), - table_name=table, - schema_name=schema_name, + table_name=_typed_str(table), + schema_name=_typed_str(schema_name), schema_type=ForeignKeyMetadata, ) @@ -455,12 +472,16 @@ async def get_constraints( if table is None: self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="constraints") rows = await driver.select( - self.get_query("constraints", "by_schema", mode=self.mode), schema_name=schema_name, table_name=None + self.get_query("constraints", "by_schema", mode=self.mode), + schema_name=_typed_str(schema_name), + table_name=_typed_str(None), ) else: self._log_table_describe(driver, schema_name=schema_name, table_name=table, operation="constraints") rows = await driver.select( - self.get_query("constraints", "by_table", mode=self.mode), table_name=table, schema_name=schema_name + self.get_query("constraints", "by_table", mode=self.mode), + table_name=_typed_str(table), + schema_name=_typed_str(schema_name), ) return self._build_domain_result("constraints", rows) @@ -586,6 +607,8 @@ def _spanner_capability_for_domain(domain: str, *, mode: str) -> MetadataCapabil risks=(MetadataRisk.VERSION_GATED,), warnings=(_SPANNER_POSTGRESQL_WARNING,), ) + if domain in _SPANNER_UNSUPPORTED_DOMAINS: + return MetadataCapability.unsupported(domain) if domain == "ddl": return MetadataCapability( domain=domain, diff --git a/sqlspec/adapters/spanner/driver.py b/sqlspec/adapters/spanner/driver.py index fe298de68..032e699e2 100644 --- a/sqlspec/adapters/spanner/driver.py +++ b/sqlspec/adapters/spanner/driver.py @@ -1,5 +1,6 @@ """Spanner driver implementation.""" +import contextvars from collections.abc import AsyncIterator, Iterator from typing import TYPE_CHECKING, Any, Literal, Protocol, cast, overload @@ -11,6 +12,7 @@ SpannerAsyncSessionContext, SpannerAsyncTransaction, SpannerGoogleAPICallError, + SpannerGoogleAPIError, SpannerSyncCursor, SpannerSyncSessionContext, SpannerTransaction, @@ -80,6 +82,10 @@ "SpannerSyncSessionContext", ) +_pending_spanner_options: contextvars.ContextVar[SpannerExecuteOptions | None] = contextvars.ContextVar( + "_pending_spanner_options", default=None +) + _READ_ONLY_SNAPSHOT_ERROR_MESSAGE = ( "Cannot execute DML in a read-only Snapshot context. " @@ -113,7 +119,7 @@ def _handle_exception(self, exc_type: "type[BaseException] | None", exc_val: "Ba if exc_type is None: return False - if isinstance(exc_val, SpannerGoogleAPICallError): + if isinstance(exc_val, (SpannerGoogleAPICallError, SpannerGoogleAPIError)): self.pending_exception = create_mapped_exception(exc_val) return True return False @@ -128,7 +134,7 @@ def _handle_exception(self, exc_type: "type[BaseException] | None", exc_val: "Ba if exc_type is None: return False - if isinstance(exc_val, SpannerGoogleAPICallError): + if isinstance(exc_val, (SpannerGoogleAPICallError, SpannerGoogleAPIError)): self.pending_exception = create_mapped_exception(exc_val) return True return False @@ -262,7 +268,9 @@ def dispatch_execute_many(self, cursor: "SpannerSyncConnection", statement: "SQL append_batch_arg((sql, coerced_params, param_types)) writer = cast("_SpannerWriteProtocol", cursor) - _status, row_counts = writer.batch_update(batch_args, **execute_kwargs) + status, row_counts = writer.batch_update(batch_args, **execute_kwargs) + if status is not None and getattr(status, "code", 0) != 0: + raise create_mapped_exception(status) total_rows = sum(row_counts) if row_counts else 0 return self.create_execution_result(cursor, rowcount_override=total_rows, is_many_result=True) @@ -324,6 +332,8 @@ def begin(self) -> None: def commit(self) -> None: """Commit the active transaction when it has work to commit.""" + if not self._owns_transaction: + return if isinstance(self.connection, SpannerTransaction): exc_handler = self.handle_database_exceptions() with exc_handler: @@ -334,6 +344,8 @@ def commit(self) -> None: def rollback(self) -> None: """Roll back the active transaction when it has begun.""" + if not self._owns_transaction: + return if isinstance(self.connection, SpannerTransaction): exc_handler = self.handle_database_exceptions() with exc_handler: @@ -387,12 +399,11 @@ def execute( execute_options = pop_execute_options(kwargs) if execute_options is None: return super().execute(statement, *parameters, statement_config=statement_config, **kwargs) - previous_options = self._pending_execute_options - self._pending_execute_options = execute_options + token = _pending_spanner_options.set(execute_options) try: return super().execute(statement, *parameters, statement_config=statement_config, **kwargs) finally: - self._pending_execute_options = previous_options + _pending_spanner_options.reset(token) def execute_many( self, @@ -407,12 +418,11 @@ def execute_many( execute_options = pop_execute_options(kwargs) if execute_options is None: return super().execute_many(statement, parameters, *filters, statement_config=statement_config, **kwargs) - previous_options = self._pending_execute_options - self._pending_execute_options = execute_options + token = _pending_spanner_options.set(execute_options) try: return super().execute_many(statement, parameters, *filters, statement_config=statement_config, **kwargs) finally: - self._pending_execute_options = previous_options + _pending_spanner_options.reset(token) def execute_script( self, @@ -426,12 +436,11 @@ def execute_script( execute_options = pop_execute_options(kwargs) if execute_options is None: return super().execute_script(statement, *parameters, statement_config=statement_config, **kwargs) - previous_options = self._pending_execute_options - self._pending_execute_options = execute_options + token = _pending_spanner_options.set(execute_options) try: return super().execute_script(statement, *parameters, statement_config=statement_config, **kwargs) finally: - self._pending_execute_options = previous_options + _pending_spanner_options.reset(token) @overload def select_stream( @@ -598,11 +607,15 @@ def load_from_arrow( msg = "Arrow import requires a Transaction context." raise SQLConversionError(msg) writer = cast("_SpannerWriteProtocol", conn) - for chunk in chunks: + for idx, chunk in enumerate(chunks): if conflict_keys and on_conflict == "nothing" and hasattr(writer, "insert_or_ignore"): writer.insert_or_ignore(table, columns, chunk) else: writer.insert_or_update(table, columns, chunk) + if self._owns_transaction and idx < len(chunks) - 1 and hasattr(writer, "commit"): + writer.commit() + self._renew_finished_transaction() + writer = cast("_SpannerWriteProtocol", self.connection) self._check_pending_exception(exc_handler) telemetry_payload = self._ingest_telemetry(arrow_table) @@ -717,14 +730,17 @@ def _apply_mutation( mutation_fn = getattr(self.connection, method_name) exc_handler = self.handle_database_exceptions() with exc_handler: - for chunk in chunks: + for idx, chunk in enumerate(chunks): + mutation_fn = getattr(self.connection, method_name) mutation_fn(table, column_list, chunk) + if self._owns_transaction and idx < len(chunks) - 1 and isinstance(self.connection, SpannerTransaction): + cast("_SpannerWriteProtocol", self.connection).commit() + self._renew_finished_transaction() self._check_pending_exception(exc_handler) def _execute_kwargs(self, *, for_read: bool = False, for_batch: bool = False) -> dict[str, Any]: - return build_execute_kwargs( - self.driver_features, self._pending_execute_options, for_read=for_read, for_batch=for_batch - ) + options = _pending_spanner_options.get() or self._pending_execute_options + return build_execute_kwargs(self.driver_features, options, for_read=for_read, for_batch=for_batch) def _chunk_mutation_rows(self, columns: "list[str]", records: "list[tuple[Any, ...]]") -> "list[list[list[Any]]]": return chunk_mutation_rows(columns, records, self._coerce_params) @@ -767,7 +783,15 @@ def _transaction_driver(self, transaction: "SpannerSyncConnection") -> "SpannerS def _connection_in_transaction(self) -> bool: """Check if connection is in transaction.""" - return False + if not self._owns_transaction: + return True + conn = self.connection + return bool( + isinstance(conn, (SpannerTransaction, SpannerAsyncTransaction)) + and (conn._transaction_id is not None or bool(conn._mutations)) + and getattr(conn, "committed", None) is None + and not getattr(conn, "rolled_back", False) + ) def _coerce_params(self, params: "dict[str, Any] | list[Any] | tuple[Any, ...] | None") -> "dict[str, Any] | None": return coerce_params( @@ -919,7 +943,9 @@ async def dispatch_execute_many(self, cursor: "SpannerAsyncConnection", statemen append_batch_arg((sql, coerced_params, param_types)) writer = cast("_SpannerAsyncWriteProtocol", cursor) - _status, row_counts = await writer.batch_update(batch_args, **execute_kwargs) + status, row_counts = await writer.batch_update(batch_args, **execute_kwargs) + if status is not None and getattr(status, "code", 0) != 0: + raise create_mapped_exception(status) total_rows = sum(row_counts) if row_counts else 0 return self.create_execution_result(cursor, rowcount_override=total_rows, is_many_result=True) @@ -984,6 +1010,8 @@ async def begin(self) -> None: async def commit(self) -> None: """Commit the active transaction when it has work to commit.""" + if not self._owns_transaction: + return if isinstance(self.connection, SpannerAsyncTransaction): exc_handler = self.handle_database_exceptions() async with exc_handler: @@ -994,6 +1022,8 @@ async def commit(self) -> None: async def rollback(self) -> None: """Roll back the active transaction when it has begun.""" + if not self._owns_transaction: + return if isinstance(self.connection, SpannerAsyncTransaction): exc_handler = self.handle_database_exceptions() async with exc_handler: @@ -1047,12 +1077,11 @@ async def execute( execute_options = pop_execute_options(kwargs) if execute_options is None: return await super().execute(statement, *parameters, statement_config=statement_config, **kwargs) - previous_options = self._pending_execute_options - self._pending_execute_options = execute_options + token = _pending_spanner_options.set(execute_options) try: return await super().execute(statement, *parameters, statement_config=statement_config, **kwargs) finally: - self._pending_execute_options = previous_options + _pending_spanner_options.reset(token) async def execute_many( self, @@ -1069,14 +1098,13 @@ async def execute_many( return await super().execute_many( statement, parameters, *filters, statement_config=statement_config, **kwargs ) - previous_options = self._pending_execute_options - self._pending_execute_options = execute_options + token = _pending_spanner_options.set(execute_options) try: return await super().execute_many( statement, parameters, *filters, statement_config=statement_config, **kwargs ) finally: - self._pending_execute_options = previous_options + _pending_spanner_options.reset(token) async def execute_script( self, @@ -1090,12 +1118,11 @@ async def execute_script( execute_options = pop_execute_options(kwargs) if execute_options is None: return await super().execute_script(statement, *parameters, statement_config=statement_config, **kwargs) - previous_options = self._pending_execute_options - self._pending_execute_options = execute_options + token = _pending_spanner_options.set(execute_options) try: return await super().execute_script(statement, *parameters, statement_config=statement_config, **kwargs) finally: - self._pending_execute_options = previous_options + _pending_spanner_options.reset(token) @overload def select_stream( @@ -1144,8 +1171,7 @@ def select_stream( native_only=native_only, **kwargs, ) - previous_options = self._pending_execute_options - self._pending_execute_options = execute_options + token = _pending_spanner_options.set(execute_options) try: return super().select_stream( statement, @@ -1157,7 +1183,7 @@ def select_stream( **kwargs, ) finally: - self._pending_execute_options = previous_options + _pending_spanner_options.reset(token) async def select_to_storage( self, @@ -1262,11 +1288,15 @@ async def load_from_arrow( msg = "Arrow import requires a Transaction context." raise SQLConversionError(msg) writer = cast("_SpannerAsyncWriteProtocol", conn) - for chunk in chunks: + for idx, chunk in enumerate(chunks): if conflict_keys and on_conflict == "nothing" and hasattr(writer, "insert_or_ignore"): writer.insert_or_ignore(table, columns, chunk) else: writer.insert_or_update(table, columns, chunk) + if self._owns_transaction and idx < len(chunks) - 1 and hasattr(writer, "commit"): + await writer.commit() + self._renew_finished_transaction() + writer = cast("_SpannerAsyncWriteProtocol", self.connection) self._check_pending_exception(exc_handler) telemetry_payload = self._ingest_telemetry(arrow_table) @@ -1390,9 +1420,8 @@ def _apply_mutation( raise create_mapped_exception(exc) from exc def _execute_kwargs(self, *, for_read: bool = False, for_batch: bool = False) -> dict[str, Any]: - return build_execute_kwargs( - self.driver_features, self._pending_execute_options, for_read=for_read, for_batch=for_batch - ) + options = _pending_spanner_options.get() or self._pending_execute_options + return build_execute_kwargs(self.driver_features, options, for_read=for_read, for_batch=for_batch) def _chunk_mutation_rows(self, columns: "list[str]", records: "list[tuple[Any, ...]]") -> "list[list[list[Any]]]": return chunk_mutation_rows(columns, records, self._coerce_params) @@ -1435,7 +1464,15 @@ def _transaction_driver(self, transaction: "SpannerAsyncConnection") -> "Spanner def _connection_in_transaction(self) -> bool: """Check if connection is in transaction.""" - return False + if not self._owns_transaction: + return True + conn = self.connection + return bool( + isinstance(conn, (SpannerTransaction, SpannerAsyncTransaction)) + and (conn._transaction_id is not None or bool(conn._mutations)) + and getattr(conn, "committed", None) is None + and not getattr(conn, "rolled_back", False) + ) def _coerce_params(self, params: "dict[str, Any] | list[Any] | tuple[Any, ...] | None") -> "dict[str, Any] | None": return coerce_params( diff --git a/sqlspec/adapters/spanner/events/store.py b/sqlspec/adapters/spanner/events/store.py index 77dbf4857..c514089a6 100644 --- a/sqlspec/adapters/spanner/events/store.py +++ b/sqlspec/adapters/spanner/events/store.py @@ -12,8 +12,10 @@ import math from typing import TYPE_CHECKING, Any +from sqlspec.adapters.spanner._typing import SpannerGoogleAPICallError from sqlspec.adapters.spanner.config import SpannerAsyncConfig, SpannerSyncConfig from sqlspec.adapters.spanner.core import ( + create_mapped_exception, execute_ddl_async, execute_ddl_sync, list_existing_table_names_async, @@ -55,12 +57,8 @@ def create_statements(self) -> "list[str]": return statements def drop_statements(self) -> "list[str]": - """Return drop statements in reverse dependency order. - - Spanner requires index to be dropped before the table. - The caller should handle errors for non-existent objects. - """ - return [f"DROP INDEX {self._index_name()}", f"DROP TABLE {self.table_name}"] + """Return drop statements in reverse dependency order.""" + return [f"DROP INDEX IF EXISTS {self._index_name()}", f"DROP TABLE IF EXISTS {self.table_name}"] def _column_types(self) -> "tuple[str, str, str]": """Return Spanner-specific column types.""" @@ -82,7 +80,18 @@ def _is_spangres(self) -> bool: """Return True when the statement dialect is Spangres (PostgreSQL on Spanner).""" statement_config = getattr(self._config, "statement_config", None) dialect = str(getattr(statement_config, "dialect", "spanner")).lower() - return dialect == "spangres" + if dialect in {"spangres", "postgres", "postgresql", "spanner_postgresql"}: + return True + conn_config = getattr(self._config, "connection_config", {}) + if isinstance(conn_config, dict): + db_dialect = str(conn_config.get("database_dialect", "")).upper() + if "POSTGRES" in db_dialect: + return True + return False + + def _resolve_ddl_timeout(self) -> "float | None": + timeout = self._extension_settings.get("ddl_timeout") + return float(timeout) if timeout is not None else None def _resolve_shard_count(self) -> int: """Validate and return the configured shard count.""" @@ -210,6 +219,13 @@ class SpannerSyncEventQueueStore(_SpannerEventStoreMixin, BaseEventQueueStore[Sp """ __slots__ = () + extension_config_options = BaseEventQueueStore.extension_config_options | frozenset({ + "ddl_timeout", + "index_options", + "retention_seconds", + "shard_count", + "table_options", + }) def create_table(self) -> None: """Create the event queue table and index through ``Database.update_ddl``. @@ -227,7 +243,10 @@ def create_table(self) -> None: return statements = self.create_statements() self._log_ddl("events.queue.create", len(statements)) - execute_ddl_sync(database, statements) + try: + execute_ddl_sync(database, statements, timeout=self._resolve_ddl_timeout()) + except SpannerGoogleAPICallError as exc: + raise create_mapped_exception(exc) from exc def drop_table(self) -> None: """Drop the event queue table and index through ``Database.update_ddl``. @@ -245,7 +264,10 @@ def drop_table(self) -> None: return statements = self.drop_statements() self._log_ddl("events.queue.drop", len(statements)) - execute_ddl_sync(database, statements) + try: + execute_ddl_sync(database, statements, timeout=self._resolve_ddl_timeout()) + except SpannerGoogleAPICallError as exc: + raise create_mapped_exception(exc) from exc class SpannerAsyncEventQueueStore(_SpannerEventStoreMixin, BaseEventQueueStore[SpannerAsyncConfig]): @@ -259,6 +281,13 @@ class SpannerAsyncEventQueueStore(_SpannerEventStoreMixin, BaseEventQueueStore[S """ __slots__ = () + extension_config_options = BaseEventQueueStore.extension_config_options | frozenset({ + "ddl_timeout", + "index_options", + "retention_seconds", + "shard_count", + "table_options", + }) async def create_table(self) -> None: """Create the event queue table and index through ``Database.update_ddl``. @@ -276,7 +305,10 @@ async def create_table(self) -> None: return statements = self.create_statements() self._log_ddl("events.queue.create", len(statements)) - await execute_ddl_async(database, statements) + try: + await execute_ddl_async(database, statements, timeout=self._resolve_ddl_timeout()) + except SpannerGoogleAPICallError as exc: + raise create_mapped_exception(exc) from exc async def drop_table(self) -> None: """Drop the event queue table and index through ``Database.update_ddl``. @@ -294,4 +326,7 @@ async def drop_table(self) -> None: return statements = self.drop_statements() self._log_ddl("events.queue.drop", len(statements)) - await execute_ddl_async(database, statements) + try: + await execute_ddl_async(database, statements, timeout=self._resolve_ddl_timeout()) + except SpannerGoogleAPICallError as exc: + raise create_mapped_exception(exc) from exc diff --git a/sqlspec/adapters/spanner/litestar/store.py b/sqlspec/adapters/spanner/litestar/store.py index 74de1b742..079f154ae 100644 --- a/sqlspec/adapters/spanner/litestar/store.py +++ b/sqlspec/adapters/spanner/litestar/store.py @@ -1,7 +1,7 @@ """Spanner session store for Litestar integration.""" from datetime import UTC, datetime, timedelta -from typing import TYPE_CHECKING, Any, NotRequired, Protocol, cast +from typing import TYPE_CHECKING, Any, Final, NotRequired, Protocol, cast from sqlspec.adapters.spanner._typing import SpannerGoogleAPICallError from sqlspec.adapters.spanner._typing import spanner_param_types as param_types @@ -9,6 +9,8 @@ create_mapped_exception, execute_ddl_async, execute_ddl_sync, + execute_partitioned_dml_async, + execute_partitioned_dml_sync, list_existing_table_names_async, list_existing_table_names_sync, ) @@ -33,11 +35,24 @@ async def run_in_transaction(self, func: "Callable[[AsyncTransaction], Any]") -> __all__ = ("SpannerAsyncStore", "SpannerLitestarConfig", "SpannerSyncStore") +_PARAM_TYPES_SESSION_ONLY: Final[dict[str, Any]] = {"session_id": param_types.STRING} +_PARAM_TYPES_EXPIRES_ONLY: Final[dict[str, Any]] = {"expires_at": param_types.TIMESTAMP} +_PARAM_TYPES_SESSION_EXPIRES: Final[dict[str, Any]] = { + "session_id": param_types.STRING, + "expires_at": param_types.TIMESTAMP, +} +_PARAM_TYPES_UPSERT: Final[dict[str, Any]] = { + "session_id": param_types.STRING, + "data": param_types.BYTES, + "expires_at": param_types.TIMESTAMP, +} + class SpannerLitestarConfig(LitestarConfig): """Spanner-specific Litestar settings. Use inside ``extension_config["litestar"]`` with this adapter's session store. + Note: Cloud Spanner Litestar store currently requires GoogleSQL dialect. """ shard_count: NotRequired[int] @@ -49,9 +64,15 @@ class SpannerLitestarConfig(LitestarConfig): index_options: NotRequired[str] """Index DDL options.""" + enable_row_deletion_policy: NotRequired[bool] + """Whether to enable Spanner native ROW DELETION POLICY for session expiration.""" + class _SpannerLitestarStoreCommonMixin: - """Shared SQL, DDL, and parameter helpers for Spanner Litestar session stores.""" + """Shared SQL, DDL, and parameter helpers for Spanner Litestar session stores. + + Note: Cloud Spanner Litestar store currently requires GoogleSQL dialect. + """ __slots__ = () @@ -59,6 +80,7 @@ class _SpannerLitestarStoreCommonMixin: _shard_count: int _table_options: "str | None" _index_options: "str | None" + _enable_row_deletion_policy: bool def _datetime_to_timestamp(self, dt: "datetime | None") -> "datetime | None": if dt is None: @@ -86,6 +108,14 @@ def _build_params( def _get_param_types( self, session_id: bool = True, expires_at: bool = False, data: bool = False ) -> "dict[str, Any]": + if session_id and data and expires_at: + return _PARAM_TYPES_UPSERT + if session_id and expires_at and not data: + return _PARAM_TYPES_SESSION_EXPIRES + if session_id and not expires_at and not data: + return _PARAM_TYPES_SESSION_ONLY + if not session_id and expires_at and not data: + return _PARAM_TYPES_EXPIRES_ONLY types: dict[str, Any] = {} if session_id: types["session_id"] = param_types.STRING @@ -121,21 +151,11 @@ def _build_renew_session_sql(self) -> str: update_sql = f"{update_sql} AND shard_id = MOD(FARM_FINGERPRINT(@session_id), {self._shard_count})" return update_sql - def _build_upsert_sql(self) -> "tuple[str, str]": - update_sql = f""" - UPDATE {self._table_name} - SET data = @data, - expires_at = @expires_at, - updated_at = PENDING_COMMIT_TIMESTAMP() - WHERE session_id = @session_id - """ - if self._shard_count > 1: - update_sql = f"{update_sql} AND shard_id = MOD(FARM_FINGERPRINT(@session_id), {self._shard_count})" - insert_sql = f""" - INSERT {self._table_name} (session_id, data, expires_at, created_at, updated_at) + def _build_upsert_sql(self) -> str: + return f""" + INSERT OR UPDATE INTO {self._table_name} (session_id, data, expires_at, created_at, updated_at) VALUES (@session_id, @data, @expires_at, PENDING_COMMIT_TIMESTAMP(), PENDING_COMMIT_TIMESTAMP()) """ - return update_sql, insert_sql def _build_delete_sql(self) -> str: sql = f"DELETE FROM {self._table_name} WHERE session_id = @session_id" @@ -186,14 +206,17 @@ def _table_ddl(self) -> str: options = "" if self._table_options: options = f"\nOPTIONS ({self._table_options})" + policy = "" + if self._enable_row_deletion_policy: + policy = ",\nROW DELETION POLICY (OLDER_THAN(expires_at, INTERVAL 0 DAY))" return f""" -CREATE TABLE {self._table_name} ( +CREATE TABLE IF NOT EXISTS {self._table_name} ( session_id STRING(128) NOT NULL, data BYTES(MAX) NOT NULL, expires_at TIMESTAMP, created_at TIMESTAMP NOT NULL OPTIONS (allow_commit_timestamp=true), updated_at TIMESTAMP NOT NULL OPTIONS (allow_commit_timestamp=true){shard_column} -) {pk}{options} +) {pk}{options}{policy} """ def _index_ddl(self) -> str: @@ -203,25 +226,32 @@ def _index_ddl(self) -> str: opts = "" if self._index_options: opts = f" OPTIONS ({self._index_options})" - return f"CREATE INDEX idx_{self._table_name}_expires_at ON {self._table_name}({leading}){opts}" + return f"CREATE INDEX IF NOT EXISTS idx_{self._table_name}_expires_at ON {self._table_name}({leading}){opts}" def _drop_table_sql(self) -> "list[str]": return [f"DROP INDEX idx_{self._table_name}_expires_at", f"DROP TABLE {self._table_name}"] class SpannerSyncStore(_SpannerLitestarStoreCommonMixin, BaseSQLSpecStore["SpannerSyncConfig"]): - """Spanner-backed Litestar session store using sync driver wrapped as async.""" + """Spanner-backed Litestar session store using sync driver wrapped as async. + + Note: Cloud Spanner Litestar store currently requires GoogleSQL dialect. + """ - __slots__ = ("_index_options", "_shard_count", "_table_options") + __slots__ = ("_enable_row_deletion_policy", "_index_options", "_shard_count", "_table_options") extension_config_options = BaseSQLSpecStore.extension_config_options | frozenset({ + "enable_row_deletion_policy", "index_options", + "row_deletion_policy", "shard_count", "table_options", }) def __init__(self, config: "SpannerSyncConfig") -> None: super().__init__(config) - self._shard_count, self._table_options, self._index_options = _extract_spanner_litestar_options(config) + (self._shard_count, self._table_options, self._index_options, self._enable_row_deletion_policy) = ( + _extract_spanner_litestar_options(config) + ) async def get(self, key: str, renew_for: "int | timedelta | None" = None) -> "bytes | None": return await async_(self._get)(key, renew_for) @@ -285,8 +315,8 @@ def _set(self, key: str, value: "str | bytes", expires_in: "int | timedelta | No expires_at = self._calculate_expires_at(expires_in) params = self._build_params(key, expires_at, data) types = self._get_param_types(session_id=True, expires_at=True, data=True) - update_sql, insert_sql = self._build_upsert_sql() - self._run_in_transaction(_SpannerSyncUpsertJob(update_sql, insert_sql, params, types)) + sql = self._build_upsert_sql() + self._run_in_transaction(_SpannerSyncUpsertJob(sql, params, types)) def _delete(self, key: str) -> None: sql = self._build_delete_sql() @@ -295,8 +325,8 @@ def _delete(self, key: str) -> None: self._run_in_transaction(_SpannerSyncExecuteUpdateJob(sql, params, types)) def _delete_all(self) -> None: - sql = self._build_delete_all_sql() - self._run_in_transaction(_SpannerSyncExecuteUpdateJob(sql)) + database = self._config.get_database() + execute_partitioned_dml_sync(database, self._build_delete_all_sql()) def _exists(self, key: str) -> bool: sql = self._build_exists_sql() @@ -317,9 +347,8 @@ def _expires_in(self, key: str) -> "int | None": return max(int(delta.total_seconds()), 0) def _delete_expired(self) -> int: - sql = self._build_delete_expired_sql() - result = self._run_in_transaction(_SpannerSyncExecuteUpdateCountJob(sql)) - return cast("int", result) + database = self._config.get_database() + return execute_partitioned_dml_sync(database, self._build_delete_expired_sql()) def _create_table(self) -> None: try: @@ -333,16 +362,20 @@ def _create_table(self) -> None: class SpannerAsyncStore(_SpannerLitestarStoreCommonMixin, BaseSQLSpecStore["SpannerAsyncConfig"]): """Spanner-backed Litestar session store using native async Spanner client.""" - __slots__ = ("_index_options", "_shard_count", "_table_options") + __slots__ = ("_enable_row_deletion_policy", "_index_options", "_shard_count", "_table_options") extension_config_options = BaseSQLSpecStore.extension_config_options | frozenset({ + "enable_row_deletion_policy", "index_options", + "row_deletion_policy", "shard_count", "table_options", }) def __init__(self, config: "SpannerAsyncConfig") -> None: super().__init__(config) - self._shard_count, self._table_options, self._index_options = _extract_spanner_litestar_options(config) + (self._shard_count, self._table_options, self._index_options, self._enable_row_deletion_policy) = ( + _extract_spanner_litestar_options(config) + ) async def get(self, key: str, renew_for: "int | timedelta | None" = None) -> "bytes | None": sql = self._build_select_session_sql() @@ -369,8 +402,8 @@ async def set(self, key: str, value: "str | bytes", expires_in: "int | timedelta expires_at = self._calculate_expires_at(expires_in) params = self._build_params(key, expires_at, data) types = self._get_param_types(session_id=True, expires_at=True, data=True) - update_sql, insert_sql = self._build_upsert_sql() - await self._run_in_transaction(_SpannerAsyncUpsertJob(update_sql, insert_sql, params, types)) + sql = self._build_upsert_sql() + await self._run_in_transaction(_SpannerAsyncUpsertJob(sql, params, types)) async def delete(self, key: str) -> None: sql = self._build_delete_sql() @@ -379,8 +412,8 @@ async def delete(self, key: str) -> None: await self._run_in_transaction(_SpannerAsyncExecuteUpdateJob(sql, params, types)) async def delete_all(self) -> None: - sql = self._build_delete_all_sql() - await self._run_in_transaction(_SpannerAsyncExecuteUpdateJob(sql)) + database = await self._database() + await execute_partitioned_dml_async(database, self._build_delete_all_sql()) async def exists(self, key: str) -> bool: sql = self._build_exists_sql() @@ -401,9 +434,8 @@ async def expires_in(self, key: str) -> "int | None": return max(int(delta.total_seconds()), 0) async def delete_expired(self) -> int: - sql = self._build_delete_expired_sql() - result = await self._run_in_transaction(_SpannerAsyncExecuteUpdateCountJob(sql)) - return int(result) + database = await self._database() + return await execute_partitioned_dml_async(database, self._build_delete_expired_sql()) async def create_table(self) -> None: if not self.create_schema_enabled: @@ -469,34 +501,28 @@ async def __call__(self, transaction: "AsyncTransaction") -> None: class _SpannerSyncUpsertJob: - __slots__ = ("_insert_sql", "_params", "_types", "_update_sql") + __slots__ = ("_params", "_sql", "_types") - def __init__(self, update_sql: str, insert_sql: str, params: "dict[str, Any]", types: "dict[str, Any]") -> None: - self._update_sql = update_sql - self._insert_sql = insert_sql + def __init__(self, sql: str, params: "dict[str, Any]", types: "dict[str, Any]") -> None: + self._sql = sql self._params = params self._types = types def __call__(self, transaction: "Transaction") -> None: writer = cast("_SpannerSyncTransactionProtocol", transaction) - row_ct = writer.execute_update(self._update_sql, params=self._params, param_types=self._types) - if row_ct == 0: - writer.execute_update(self._insert_sql, params=self._params, param_types=self._types) + writer.execute_update(self._sql, params=self._params, param_types=self._types) class _SpannerAsyncUpsertJob: - __slots__ = ("_insert_sql", "_params", "_types", "_update_sql") + __slots__ = ("_params", "_sql", "_types") - def __init__(self, update_sql: str, insert_sql: str, params: "dict[str, Any]", types: "dict[str, Any]") -> None: - self._update_sql = update_sql - self._insert_sql = insert_sql + def __init__(self, sql: str, params: "dict[str, Any]", types: "dict[str, Any]") -> None: + self._sql = sql self._params = params self._types = types async def __call__(self, transaction: "AsyncTransaction") -> None: - row_ct = await transaction.execute_update(self._update_sql, params=self._params, param_types=self._types) - if row_ct == 0: - await transaction.execute_update(self._insert_sql, params=self._params, param_types=self._types) + await transaction.execute_update(self._sql, params=self._params, param_types=self._types) class _SpannerSyncExecuteUpdateCountJob: @@ -520,9 +546,12 @@ async def __call__(self, transaction: "AsyncTransaction") -> int: return int(await transaction.execute_update(self._sql)) -def _extract_spanner_litestar_options(config: Any) -> "tuple[int, str | None, str | None]": +def _extract_spanner_litestar_options(config: Any) -> "tuple[int, str | None, str | None, bool]": litestar_cfg = cast("dict[str, Any]", config.extension_config.get("litestar", {})) shard_count = int(litestar_cfg.get("shard_count", 0)) if litestar_cfg.get("shard_count") else 0 table_options = cast("str | None", litestar_cfg.get("table_options")) index_options = cast("str | None", litestar_cfg.get("index_options")) - return shard_count, table_options, index_options + enable_row_deletion_policy = bool( + litestar_cfg.get("enable_row_deletion_policy", False) or litestar_cfg.get("row_deletion_policy", False) + ) + return shard_count, table_options, index_options, enable_row_deletion_policy diff --git a/sqlspec/adapters/spanner/type_converter.py b/sqlspec/adapters/spanner/type_converter.py index 1863d0b24..256555b4b 100644 --- a/sqlspec/adapters/spanner/type_converter.py +++ b/sqlspec/adapters/spanner/type_converter.py @@ -92,10 +92,10 @@ def bytes_to_spanner(value: "bytes | None") -> "bytes | None": def spanner_to_bytes(value: Any) -> "bytes | None": """Convert Spanner BYTES result to Python bytes. - Handles both raw bytes and base64-encoded bytes. + Handles both raw bytes and base64-encoded bytes or strings. Args: - value: Value from Spanner (bytes or None). + value: Value from Spanner (bytes, str, or None). Returns: Python bytes or None. @@ -103,7 +103,10 @@ def spanner_to_bytes(value: Any) -> "bytes | None": if value is None: return None if isinstance(value, (bytes, str)): - return base64.b64decode(value) + try: + return base64.b64decode(value, validate=True) + except Exception: + return value if isinstance(value, bytes) else value.encode("utf-8") return None @@ -189,56 +192,66 @@ def coerce_params_for_spanner( return None json_object_type = _get_json_object_type() - coerced: dict[str, Any] = {} - changed = False + coerced: dict[str, Any] | None = None for key, raw_value in params.items(): value = raw_value + val_changed = False + new_val = value if type(raw_value) is TypedParameter: value = raw_value.value - changed = True + new_val = value + val_changed = True declared = raw_value.semantic_name if declared is not None: normalized_declared = declared.strip().upper() if normalized_declared in {"ARRAY", "VECTOR"} and isinstance(value, (list, tuple)): - coerced[key] = [float(item) for item in value] + new_val = [float(item) for item in value] + if coerced is None: + coerced = dict(params) + coerced[key] = new_val continue if normalized_declared == "FLOAT32" and value is not None: - coerced[key] = float(value) + new_val = float(value) + if coerced is None: + coerced = dict(params) + coerced[key] = new_val continue if isinstance(value, _UUID_TYPES): if enable_uuid_conversion: - coerced[key] = str(value) - changed = True - else: - coerced[key] = value + new_val = str(value) + val_changed = True elif isinstance(value, bytes): - coerced[key] = bytes_to_spanner(value) - changed = True + new_val = bytes_to_spanner(value) + val_changed = True elif isinstance(value, datetime) and value.tzinfo is None: - coerced[key] = value.replace(tzinfo=UTC) - changed = True + new_val = value.replace(tzinfo=UTC) + val_changed = True elif isinstance(value, timedelta): - coerced[key] = Interval(days=value.days, nanos=(value.seconds * 1_000_000 + value.microseconds) * 1000) - changed = True + new_val = Interval(days=value.days, nanos=(value.seconds * 1_000_000 + value.microseconds) * 1000) + val_changed = True elif isinstance(value, json_object_type): - coerced[key] = value + pass elif isinstance(value, dict): - coerced[key] = spanner_json(value, json_serializer=json_serializer) - changed = True + new_val = spanner_json(value, json_serializer=json_serializer) + val_changed = True elif isinstance(value, (list, tuple)): if should_json_encode_sequence(value): - coerced[key] = spanner_json(list(value), json_serializer=json_serializer) - changed = True + new_val = spanner_json(list(value), json_serializer=json_serializer) + val_changed = True else: coerced_seq, seq_changed = _coerce_sequence_elements( value, enable_uuid_conversion=enable_uuid_conversion ) - coerced[key] = coerced_seq if seq_changed: - changed = True - else: - coerced[key] = value - return coerced if changed else params + new_val = coerced_seq + val_changed = True + + if val_changed: + if coerced is None: + coerced = dict(params) + coerced[key] = new_val + + return coerced if coerced is not None else params def infer_spanner_param_types(params: "dict[str, Any] | None") -> "dict[str, Any]": @@ -301,12 +314,15 @@ def _coerce_sequence_elements( value: "list[Any] | tuple[Any, ...]", *, enable_uuid_conversion: bool ) -> "tuple[list[Any] | tuple[Any, ...], bool]": """Coerce elements inside a 1-D sequence parameter when required by Spanner.""" - needs_element_coercion = any( - isinstance(item, (timedelta, bytes)) - or (enable_uuid_conversion and isinstance(item, _UUID_TYPES)) - or (isinstance(item, datetime) and item.tzinfo is None) - for item in value - ) + needs_element_coercion = False + for item in value: + if ( + isinstance(item, (timedelta, bytes)) + or (enable_uuid_conversion and isinstance(item, _UUID_TYPES)) + or (isinstance(item, datetime) and item.tzinfo is None) + ): + needs_element_coercion = True + break if not needs_element_coercion: if isinstance(value, tuple): return list(value), True @@ -375,7 +391,11 @@ def _infer_sequence_param_type(value: Any, param_types: Any, json_type: Any) -> return json_type if not value: return None - first = next((item for item in value if item is not None), None) + first = None + for item in value: + if item is not None: + first = item + break if first is None: return None if isinstance(first, bool): diff --git a/sqlspec/adapters/sqlanywhere/_typing.py b/sqlspec/adapters/sqlanywhere/_typing.py index d9420b415..5525e8df0 100644 --- a/sqlspec/adapters/sqlanywhere/_typing.py +++ b/sqlspec/adapters/sqlanywhere/_typing.py @@ -5,19 +5,13 @@ """ import contextlib +import importlib from typing import TYPE_CHECKING, Any +import sqlanydb as SQLANYDB_MODULE # noqa: N812 from mypy_extensions import mypyc_attr - -from sqlspec.utils.module_loader import import_optional_attr - -SQLANYDB_MODULE: Any = None -_DRIVER_CLS: Any = None - - -class SqlanywhereError(Exception): - """SAP SQL Anywhere driver exception base.""" - +from sqlanydb import Connection, Cursor +from sqlanydb import Error as SqlanywhereError if TYPE_CHECKING: from collections.abc import Callable @@ -27,13 +21,12 @@ class SqlanywhereError(Exception): from sqlspec.adapters.sqlanywhere.driver import SqlanywhereDriver from sqlspec.core import StatementConfig - SqlanywhereConnection: TypeAlias = Any - SqlanywhereRawCursor: TypeAlias = Any + SqlanywhereConnection: TypeAlias = Connection + SqlanywhereRawCursor: TypeAlias = Cursor SqlanywhereQueryParams: TypeAlias = Any - -if not TYPE_CHECKING: - SqlanywhereConnection = Any - SqlanywhereRawCursor = Any +else: + SqlanywhereConnection = Connection + SqlanywhereRawCursor = Cursor SqlanywhereQueryParams = Any __all__ = ( @@ -67,6 +60,9 @@ def __exit__(self, *_: Any) -> None: self.cursor.close() +_DRIVER_CLS: "type[SqlanywhereDriver] | None" = None + + @mypyc_attr(allow_interpreted_subclasses=True, native_class=False) class SqlanywhereSessionContext: """Sync context manager for SAP SQL Anywhere sessions.""" @@ -100,7 +96,9 @@ def __init__( def __enter__(self) -> "SqlanywhereDriver": global _DRIVER_CLS if _DRIVER_CLS is None: - _DRIVER_CLS = import_optional_attr("sqlspec.adapters.sqlanywhere.driver", "SqlanywhereDriver") + module = importlib.import_module("sqlspec.adapters.sqlanywhere.driver") + _DRIVER_CLS = module.SqlanywhereDriver + assert _DRIVER_CLS is not None self._connection = self._acquire_connection() self._driver = _DRIVER_CLS( connection=self._connection, statement_config=self._statement_config, driver_features=self._driver_features @@ -111,13 +109,10 @@ def __enter__(self) -> "SqlanywhereDriver": def __exit__( self, exc_type: "type[BaseException] | None", exc_val: "BaseException | None", exc_tb: "TracebackType | None" ) -> "bool | None": - if self._connection is None: - return None - try: - if self._driver is not None: - self._driver._release_open_work() - finally: - self._release_connection(self._connection) + if self._driver is not None: + self._driver._release_open_work() + if self._connection is not None: + self._release_connection(self._connection, exc_type=exc_type, exc_val=exc_val, exc_tb=exc_tb) self._connection = None self._driver = None return None diff --git a/sqlspec/adapters/sqlanywhere/adk/__init__.py b/sqlspec/adapters/sqlanywhere/adk/__init__.py index 5b4598be7..ca79ac88c 100644 --- a/sqlspec/adapters/sqlanywhere/adk/__init__.py +++ b/sqlspec/adapters/sqlanywhere/adk/__init__.py @@ -1,5 +1,5 @@ """SAP SQL Anywhere ADK store module.""" -from sqlspec.adapters.sqlanywhere.adk.store import SqlanywhereADKConfig, SqlanywhereADKMemoryStore, SqlanywhereADKStore +from sqlspec.adapters.sqlanywhere.adk.store import SqlanywhereADKMemoryStore, SqlanywhereADKStore -__all__ = ("SqlanywhereADKConfig", "SqlanywhereADKMemoryStore", "SqlanywhereADKStore") +__all__ = ("SqlanywhereADKMemoryStore", "SqlanywhereADKStore") diff --git a/sqlspec/adapters/sqlanywhere/adk/store.py b/sqlspec/adapters/sqlanywhere/adk/store.py index 6b94477ae..a41869198 100644 --- a/sqlspec/adapters/sqlanywhere/adk/store.py +++ b/sqlspec/adapters/sqlanywhere/adk/store.py @@ -10,7 +10,6 @@ resolve_rowcount, split_schema_and_table, ) -from sqlspec.config import ADKConfig from sqlspec.extensions.adk import BaseSyncADKStore, StoredEvent, StoredSession, normalize_session_list_options from sqlspec.extensions.adk.memory.store import BaseSyncADKMemoryStore from sqlspec.utils.serializers import from_json, to_json @@ -23,22 +22,20 @@ from sqlspec.extensions.adk import SessionOrderBy from sqlspec.extensions.adk.memory._types import StoredMemory -__all__ = ("SqlanywhereADKConfig", "SqlanywhereADKMemoryStore", "SqlanywhereADKStore") +__all__ = ("SqlanywhereADKMemoryStore", "SqlanywhereADKStore") SQLANYWHERE_TABLE_NOT_FOUND_SQLCODE: Final[int] = -141 JSON_COLUMN_TYPE: Final[str] = "LONG VARCHAR" -class SqlanywhereADKConfig(ADKConfig): - """SAP SQL Anywhere ADK extension settings.""" - - class SqlanywhereADKStore(BaseSyncADKStore["SqlanywhereConfig"]): """Synchronous SAP SQL Anywhere ADK session/event store.""" connector_name: ClassVar[str] = "sqlanywhere" __slots__ = () + _max_table_name_parts = 2 + def __init__(self, config: "SqlanywhereConfig") -> None: super().__init__(config) @@ -416,6 +413,8 @@ class SqlanywhereADKMemoryStore(BaseSyncADKMemoryStore["SqlanywhereConfig"]): __slots__ = () + _max_table_name_parts = 2 + def __init__(self, config: "SqlanywhereConfig") -> None: super().__init__(config) @@ -450,6 +449,7 @@ def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object content_json, content_text, metadata_json{owner_column} ) SELECT ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?{owner_value} + FROM SYS.DUMMY WHERE NOT EXISTS ( SELECT 1 FROM {_table_ref(self._memory_table)} WHERE event_id = ? @@ -508,7 +508,12 @@ def search_entries( ORDER BY timestamp DESC LIMIT ? """ - rows = self._execute_fetchall(sql, (*scope_params, f"%{escaped_query}%", limit_value)) + try: + rows = self._execute_fetchall(sql, (*scope_params, f"%{escaped_query}%", limit_value)) + except SqlanywhereError as exc: + if _is_sqlanywhere_table_missing(exc): + return [] + raise return [_memory_record_from_row(row) for row in rows] def delete_entries_by_session(self, session_id: str) -> int: diff --git a/sqlspec/adapters/sqlanywhere/core.py b/sqlspec/adapters/sqlanywhere/core.py index f36e7240d..7fbc5861d 100644 --- a/sqlspec/adapters/sqlanywhere/core.py +++ b/sqlspec/adapters/sqlanywhere/core.py @@ -142,13 +142,13 @@ def format_identifier(identifier: str) -> str: msg = "Table name must not be empty" raise SQLSpecError(msg) parts = split_qualified_identifier(cleaned, quote_chars='"', allow_bracket_quotes=True) - return ".".join(quote_double_identifier(part) for part in parts) + return ".".join([quote_double_identifier(part) for part in parts]) def build_insert_statement(table: str, columns: "list[str]") -> str: """Build a QMARK-parameterized INSERT statement for SQL Anywhere.""" - column_clause = ", ".join(quote_double_identifier(column) for column in columns) - placeholders = ", ".join("?" for _ in columns) + column_clause = ", ".join([quote_double_identifier(column) for column in columns]) + placeholders = ", ".join(["?"] * len(columns)) if columns else "" return f"INSERT INTO {format_identifier(table)} ({column_clause}) VALUES ({placeholders})" diff --git a/sqlspec/adapters/sqlanywhere/data_dictionary.py b/sqlspec/adapters/sqlanywhere/data_dictionary.py index 4743347cc..5c442d87a 100644 --- a/sqlspec/adapters/sqlanywhere/data_dictionary.py +++ b/sqlspec/adapters/sqlanywhere/data_dictionary.py @@ -69,7 +69,7 @@ def build_column_metadata(rows: "Iterable[dict[str, Any]]") -> "list[ColumnMetad table_name=str(r.get("table_name", "")), column_name=str(r.get("column_name", "")), data_type=str(r.get("data_type", "")), - is_nullable=bool(r.get("is_nullable", True)), + is_nullable=r.get("is_nullable") in (True, 1, "YES", "yes"), column_default=str(r["column_default"]) if r.get("column_default") is not None else None, ordinal_position=int(r.get("ordinal_position", 1)), max_length=int(r.get("max_length", 0)) if r.get("max_length") is not None else 0, @@ -128,8 +128,8 @@ def build_foreign_key_metadata(rows: "Iterable[dict[str, Any]]") -> "list[Foreig ForeignKeyMetadata( table_name=str(r.get("table_name", "")), column_name=str(r.get("column_name", "")), - referenced_table=str(r.get("referenced_table_name", "")), - referenced_column=str(r.get("referenced_column_name", "")), + referenced_table=str(r.get("referenced_table", r.get("referenced_table_name", ""))), + referenced_column=str(r.get("referenced_column", r.get("referenced_column_name", ""))), constraint_name=str(r.get("constraint_name", "")), schema=str(r["schema_name"]) if r.get("schema_name") else None, referenced_schema=str(r["referenced_schema"]) if r.get("referenced_schema") else None, @@ -144,9 +144,6 @@ class SqlanywhereSyncDataDictionary(SyncDataDictionaryBase): dialect: ClassVar[str] = "sqlanywhere" - def __init__(self) -> None: - super().__init__() - def get_metadata_capabilities( self, driver: "SqlanywhereDriver", domains: "Sequence[str] | None" = None ) -> MetadataCapabilityProfile: diff --git a/sqlspec/adapters/sqlanywhere/driver.py b/sqlspec/adapters/sqlanywhere/driver.py index 445f2e357..69fd3dec1 100644 --- a/sqlspec/adapters/sqlanywhere/driver.py +++ b/sqlspec/adapters/sqlanywhere/driver.py @@ -1,13 +1,13 @@ """SAP SQL Anywhere database driver implementation over sqlanydb.""" import contextlib +import sys from collections.abc import Sized from typing import TYPE_CHECKING, Any, Literal, cast from sqlspec.adapters.sqlanywhere._typing import ( SqlanywhereConnection, SqlanywhereCursor, - SqlanywhereError, SqlanywhereRawCursor, SqlanywhereSessionContext, ) @@ -34,7 +34,7 @@ rows_to_dicts, validate_savepoint_name, ) -from sqlspec.exceptions import SQLSpecError +from sqlspec.exceptions import SQLSpecError, TransactionError from sqlspec.utils.logging import get_logger if TYPE_CHECKING: @@ -66,7 +66,7 @@ class SqlanywhereExceptionHandler(BaseSyncExceptionHandler): def _handle_exception(self, exc_type: "type[BaseException] | None", exc_val: "BaseException") -> bool: if exc_type is None: return False - if isinstance(exc_val, (SqlanywhereError, Exception)) and not isinstance(exc_val, SQLSpecError): + if isinstance(exc_val, Exception) and not isinstance(exc_val, SQLSpecError): self.pending_exception = create_mapped_exception(exc_val, logger=logger) return True return False @@ -90,12 +90,19 @@ def start(self) -> None: try: cursor = cursor_manager.__enter__() handler = self._driver.handle_database_exceptions() - with handler: + if sys.exc_info()[1] is None: + with handler: + normalized_params = normalize_execute_parameters(self._parameters) + if normalized_params: + cursor.execute(self._sql, normalized_params) + else: + cursor.execute(self._sql) + else: normalized_params = normalize_execute_parameters(self._parameters) if normalized_params: - cursor.execute(self._sql, normalized_params) + self._driver._run_with_exception_handler(handler, cursor.execute, self._sql, normalized_params) else: - cursor.execute(self._sql) + self._driver._run_with_exception_handler(handler, cursor.execute, self._sql) self._driver._check_pending_exception(handler) except BaseException: with contextlib.suppress(Exception): @@ -110,8 +117,11 @@ def fetch_chunk(self) -> "list[dict[str, Any]]": cursor = cursor_manager.cursor handler = self._driver.handle_database_exceptions() rows: Any = [] - with handler: - rows = cursor.fetchmany(self._chunk_size) + if sys.exc_info()[1] is None: + with handler: + rows = cursor.fetchmany(self._chunk_size) + else: + rows = self._driver._run_with_exception_handler(handler, cursor.fetchmany, self._chunk_size) self._driver._check_pending_exception(handler) if not rows: return [] @@ -216,14 +226,14 @@ def commit(self) -> None: self._transaction_active = False except Exception as exc: msg = f"Failed to commit SQL Anywhere transaction: {exc}" - raise SQLSpecError(msg) from exc + raise TransactionError(msg) from exc def rollback(self) -> None: try: self.connection.rollback() except Exception as exc: msg = f"Failed to rollback SQL Anywhere transaction: {exc}" - raise SQLSpecError(msg) from exc + raise TransactionError(msg) from exc finally: self._transaction_active = False @@ -253,9 +263,19 @@ def has_schema(self, schema: str) -> bool: """Return whether the specified user/schema exists in SYS.SYSUSER.""" result = False exc_handler = self.handle_database_exceptions() - with exc_handler, self.with_cursor(self.connection) as cursor: - cursor.execute("SELECT 1 FROM SYS.SYSUSER WHERE user_name = ?", (schema,)) - result = cursor.fetchone() is not None + with self.with_cursor(self.connection) as cursor: + if sys.exc_info()[1] is None: + with exc_handler: + cursor.execute("SELECT 1 FROM SYS.SYSUSER WHERE user_name = ?", (schema,)) + result = cursor.fetchone() is not None + else: + + def _fetch() -> None: + nonlocal result + cursor.execute("SELECT 1 FROM SYS.SYSUSER WHERE user_name = ?", (schema,)) + result = cursor.fetchone() is not None + + self._run_with_exception_handler(exc_handler, _fetch) self._check_pending_exception(exc_handler) return result @@ -337,8 +357,12 @@ def load_from_arrow( if overwrite: statement = f"DELETE FROM {format_identifier(table)}" exc_handler = self.handle_database_exceptions() - with exc_handler, self.with_cursor(self.connection) as cursor: - cursor.execute(statement) + with self.with_cursor(self.connection) as cursor: + if sys.exc_info()[1] is None: + with exc_handler: + cursor.execute(statement) + else: + self._run_with_exception_handler(exc_handler, cursor.execute, statement) self._check_pending_exception(exc_handler) columns, records = self._arrow_table_to_rows(arrow_table) @@ -352,8 +376,14 @@ def load_from_arrow( ) normalized_records = normalize_execute_many_parameters(prepared_records) exc_handler = self.handle_database_exceptions() - with exc_handler, self.with_cursor(self.connection) as cursor: - cursor.executemany(insert_sql, cast("Sequence[QueryParams]", normalized_records)) + with self.with_cursor(self.connection) as cursor: + if sys.exc_info()[1] is None: + with exc_handler: + cursor.executemany(insert_sql, cast("Sequence[QueryParams]", normalized_records)) + else: + self._run_with_exception_handler( + exc_handler, cursor.executemany, insert_sql, cast("Sequence[QueryParams]", normalized_records) + ) self._check_pending_exception(exc_handler) telemetry_payload = self._ingest_telemetry(arrow_table) diff --git a/sqlspec/adapters/sqlanywhere/events/__init__.py b/sqlspec/adapters/sqlanywhere/events/__init__.py index 70252c8ad..d40657a0a 100644 --- a/sqlspec/adapters/sqlanywhere/events/__init__.py +++ b/sqlspec/adapters/sqlanywhere/events/__init__.py @@ -1,5 +1,5 @@ """Events integration for SAP SQL Anywhere adapter.""" -from sqlspec.adapters.sqlanywhere.events.store import SqlanywhereEventQueueStore, SqlanywhereSyncEventQueueStore +from sqlspec.adapters.sqlanywhere.events.store import SqlanywhereEventQueueStore -__all__ = ("SqlanywhereEventQueueStore", "SqlanywhereSyncEventQueueStore") +__all__ = ("SqlanywhereEventQueueStore",) diff --git a/sqlspec/adapters/sqlanywhere/events/store.py b/sqlspec/adapters/sqlanywhere/events/store.py index 9a4e23b35..8c945dd43 100644 --- a/sqlspec/adapters/sqlanywhere/events/store.py +++ b/sqlspec/adapters/sqlanywhere/events/store.py @@ -3,12 +3,12 @@ from sqlspec.adapters.sqlanywhere.config import SqlanywhereConfig from sqlspec.extensions.events import BaseEventQueueStore -__all__ = ("SqlanywhereEventQueueStore", "SqlanywhereSyncEventQueueStore") +__all__ = ("SqlanywhereEventQueueStore",) _VARCHAR_TEXT_THRESHOLD = 4000 -class SqlanywhereSyncEventQueueStore(BaseEventQueueStore[SqlanywhereConfig]): +class SqlanywhereEventQueueStore(BaseEventQueueStore[SqlanywhereConfig]): """SAP SQL Anywhere DDL hooks for the event queue store.""" __slots__ = () @@ -26,6 +26,3 @@ def _integer_type(self) -> str: def _timestamp_default(self) -> str: return "CURRENT UTC TIMESTAMP" - - -SqlanywhereEventQueueStore = SqlanywhereSyncEventQueueStore diff --git a/sqlspec/adapters/sqlanywhere/litestar/__init__.py b/sqlspec/adapters/sqlanywhere/litestar/__init__.py index d4b7031a6..8c56c93a6 100644 --- a/sqlspec/adapters/sqlanywhere/litestar/__init__.py +++ b/sqlspec/adapters/sqlanywhere/litestar/__init__.py @@ -1,5 +1,5 @@ """Litestar integration for SAP SQL Anywhere adapter.""" -from sqlspec.adapters.sqlanywhere.litestar.store import SqlanywhereStore +from sqlspec.adapters.sqlanywhere.litestar.store import SqlanywhereLitestarConfig, SqlanywhereStore -__all__ = ("SqlanywhereStore",) +__all__ = ("SqlanywhereLitestarConfig", "SqlanywhereStore") diff --git a/sqlspec/adapters/sqlanywhere/litestar/store.py b/sqlspec/adapters/sqlanywhere/litestar/store.py index e2b325534..d3e16d050 100644 --- a/sqlspec/adapters/sqlanywhere/litestar/store.py +++ b/sqlspec/adapters/sqlanywhere/litestar/store.py @@ -4,14 +4,19 @@ from typing import TYPE_CHECKING, Any from sqlspec.adapters.sqlanywhere._typing import SqlanywhereCursor -from sqlspec.adapters.sqlanywhere.core import quote_double_identifier, split_schema_and_table +from sqlspec.adapters.sqlanywhere.core import quote_double_identifier, resolve_rowcount, split_schema_and_table +from sqlspec.config import LitestarConfig from sqlspec.extensions.litestar.store import BaseSQLSpecStore from sqlspec.utils.sync_tools import async_ if TYPE_CHECKING: from sqlspec.adapters.sqlanywhere.config import SqlanywhereConfig -__all__ = ("SqlanywhereStore",) +__all__ = ("SqlanywhereLitestarConfig", "SqlanywhereStore") + + +class SqlanywhereLitestarConfig(LitestarConfig): + """Configuration for SAP SQL Anywhere Litestar session store.""" class SqlanywhereStore(BaseSQLSpecStore["SqlanywhereConfig"]): @@ -182,7 +187,7 @@ def _delete_expired(self) -> int: with self._config.provide_connection() as conn, SqlanywhereCursor(conn) as cursor: cursor.execute(sql) conn.commit() - count = cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0 + count = resolve_rowcount(cursor) if count > 0: self._log_delete_expired(count) return count diff --git a/sqlspec/adapters/sqlanywhere/pool.py b/sqlspec/adapters/sqlanywhere/pool.py index e10f930f8..1f740f674 100644 --- a/sqlspec/adapters/sqlanywhere/pool.py +++ b/sqlspec/adapters/sqlanywhere/pool.py @@ -5,11 +5,10 @@ import threading import time from contextlib import contextmanager -from importlib import import_module from typing import TYPE_CHECKING, Any, cast -from sqlspec.adapters.sqlanywhere._typing import SqlanywhereConnection -from sqlspec.exceptions import MissingDependencyError +from sqlspec.adapters.sqlanywhere._typing import SQLANYDB_MODULE, SqlanywhereConnection +from sqlspec.exceptions import OperationalError from sqlspec.utils.logging import POOL_LOGGER_NAME, get_logger, log_with_context from sqlspec.utils.uuids import uuid4 @@ -73,11 +72,10 @@ def _database_name(self) -> str: def new_connection(self) -> SqlanywhereConnection: """Open a standalone SAP SQL Anywhere connection via ``sqlanydb.connect``.""" - sqlanydb = _load_sqlanydb_module() - connection = sqlanydb.connect(**self._connection_parameters) + connection = SQLANYDB_MODULE.connect(**self._connection_parameters) if self._on_connection_create is not None: self._on_connection_create(connection) - return cast("SqlanywhereConnection", connection) + return connection @contextmanager def get_connection(self) -> "Generator[SqlanywhereConnection, None, None]": @@ -142,8 +140,17 @@ def _prune_dead_thread_connections(self) -> None: def _create_connection(self) -> SqlanywhereConnection: self._prune_dead_thread_connections() + with self._registry_lock: + if self._max_size > 0 and len(self._connection_registry) >= self._max_size: + msg = f"Connection pool exhausted (max_size={self._max_size})" + raise OperationalError(msg) connection = self.new_connection() with self._registry_lock: + if self._max_size > 0 and len(self._connection_registry) >= self._max_size: + with contextlib.suppress(Exception): + connection.close() + msg = f"Connection pool exhausted (max_size={self._max_size})" + raise OperationalError(msg) self._connection_registry.add(connection) self._connection_threads[connection] = threading.current_thread() return connection @@ -227,11 +234,3 @@ def _close_thread_connection(self) -> None: self._thread_local.created_at = 0.0 self._thread_local.last_used = 0.0 self._thread_local.generation = -1 - - -def _load_sqlanydb_module() -> Any: - """Import ``sqlanydb`` lazily or raise ``MissingDependencyError``.""" - try: - return import_module("sqlanydb") - except ImportError as exc: - raise MissingDependencyError(package="sqlanydb", install_package="sqlanywhere") from exc diff --git a/sqlspec/adapters/sqlite/_typing.py b/sqlspec/adapters/sqlite/_typing.py index e2661eff7..4b3e87a4e 100644 --- a/sqlspec/adapters/sqlite/_typing.py +++ b/sqlspec/adapters/sqlite/_typing.py @@ -6,7 +6,7 @@ import contextlib import sqlite3 as sqlite_module -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, ClassVar _SqliteConnection = sqlite_module.Connection @@ -90,12 +90,15 @@ class SqliteSessionContext: "_acquire_connection", "_connection", "_driver", + "_driver_class", "_driver_features", "_prepare_driver", "_release_connection", "_statement_config", ) + _default_driver_class: ClassVar["type[SqliteDriver] | None"] = None + def __init__( self, acquire_connection: "Callable[[], Any]", @@ -103,20 +106,26 @@ def __init__( statement_config: "StatementConfig", driver_features: "dict[str, Any]", prepare_driver: "Callable[[SqliteDriver], SqliteDriver]", + driver_class: "type[SqliteDriver] | None" = None, ) -> None: self._acquire_connection = acquire_connection self._release_connection = release_connection self._statement_config = statement_config self._driver_features = driver_features self._prepare_driver = prepare_driver + self._driver_class = driver_class or self._default_driver_class self._connection: Any = None self._driver: SqliteDriver | None = None def __enter__(self) -> "SqliteDriver": - from sqlspec.adapters.sqlite.driver import SqliteDriver + driver_cls = self._driver_class + if driver_cls is None: + from sqlspec.adapters.sqlite.driver import SqliteDriver + + driver_cls = SqliteDriver self._connection = self._acquire_connection() - self._driver = SqliteDriver( + self._driver = driver_cls( connection=self._connection, statement_config=self._statement_config, driver_features=self._driver_features ) return self._prepare_driver(self._driver) diff --git a/sqlspec/adapters/sqlite/adk/store.py b/sqlspec/adapters/sqlite/adk/store.py index a8e563ea6..6a919cfbe 100644 --- a/sqlspec/adapters/sqlite/adk/store.py +++ b/sqlspec/adapters/sqlite/adk/store.py @@ -1,6 +1,8 @@ """SQLite sync ADK store for Google Agent Development Kit session/event storage.""" +import array import contextlib +import math import re from collections.abc import Mapping from datetime import UTC, datetime, timedelta @@ -478,7 +480,11 @@ def get_events( params.append(_datetime_to_julian(after_timestamp)) where_clause = " AND ".join(where_clauses) - limit_clause = f" LIMIT {limit}" if limit else "" + if limit is not None: + limit_clause = " LIMIT ?" + params.append(int(limit)) + else: + limit_clause = "" sql = f""" SELECT id, app_name, user_id, session_id, invocation_id, timestamp, event_data @@ -887,8 +893,8 @@ def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object INSERT OR IGNORE INTO {self._memory_table} (id, session_id, app_name, user_id, scope, event_id, author, {self._owner_id_column_name}, timestamp, content_json, - content_text, metadata_json, inserted_at) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + content_text, metadata_json, inserted_at, embedding) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """ for entry in entries: timestamp_julian = _datetime_to_julian(entry["timestamp"]) @@ -911,13 +917,14 @@ def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object entry["content_text"], metadata_json_str, inserted_at_julian, + _pack_embedding(entry.get("embedding")), )) else: sql = f""" INSERT OR IGNORE INTO {self._memory_table} (id, session_id, app_name, user_id, scope, event_id, author, - timestamp, content_json, content_text, metadata_json, inserted_at) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + timestamp, content_json, content_text, metadata_json, inserted_at, embedding) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """ for entry in entries: timestamp_julian = _datetime_to_julian(entry["timestamp"]) @@ -939,6 +946,7 @@ def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object entry["content_text"], metadata_json_str, inserted_at_julian, + _pack_embedding(entry.get("embedding")), )) cursor = conn.executemany(sql, params_list) @@ -959,19 +967,22 @@ def search_entries( scope_filter: Literal["all", "user", "app"] = "all", embedding: "Sequence[float] | None" = None, ) -> "list[StoredMemory]": - """Search memory entries by text query.""" + """Search memory entries by text query or vector embedding.""" if not self._enabled: msg = "Memory store is disabled" raise RuntimeError(msg) - if not query or not query.strip(): + has_query = bool(query and query.strip()) + if not has_query and embedding is None: return [] effective_limit = limit if limit is not None else self._max_results + fetch_limit = effective_limit if embedding is None else max(effective_limit * 5, 50) - if self._use_fts: + records: list[StoredMemory] | None = None + if has_query and self._use_fts: try: - return self._search_entries_fts(query, app_name, user_id, effective_limit, scope_filter) + records = self._search_entries_fts(query, app_name, user_id, fetch_limit, scope_filter) except sqlite3.OperationalError as exc: error_text = str(exc) if SQLITE_TABLE_NOT_FOUND_ERROR in error_text and f"{self._memory_table}_fts" not in error_text: @@ -979,12 +990,29 @@ def search_entries( logger.warning("FTS search failed; falling back to simple search: %s", exc) except Exception as exc: logger.warning("FTS search failed; falling back to simple search: %s", exc) - try: - return self._search_entries_simple(query, app_name, user_id, effective_limit, scope_filter) - except sqlite3.OperationalError as exc: - if SQLITE_TABLE_NOT_FOUND_ERROR in str(exc): - return [] - raise + + if records is None: + try: + if has_query: + records = self._search_entries_simple(query, app_name, user_id, fetch_limit, scope_filter) + else: + records = self._search_entries_vector(app_name, user_id, fetch_limit, scope_filter) + except sqlite3.OperationalError as exc: + if SQLITE_TABLE_NOT_FOUND_ERROR in str(exc): + return [] + raise + + if embedding is not None and records: + query_vec = [float(x) for x in embedding] + + def _score(r: "StoredMemory") -> float: + emb = r.get("embedding") + return _cosine_similarity(query_vec, emb) if emb is not None else -1.0 + + records.sort(key=_score, reverse=True) + records = records[:effective_limit] + + return records def delete_entries_by_session(self, session_id: str) -> int: """Delete all memory entries for a specific session.""" @@ -1092,7 +1120,8 @@ def _memory_table_ddl(self) -> str: content_json TEXT NOT NULL, content_text TEXT NOT NULL, metadata_json TEXT, - inserted_at REAL NOT NULL + inserted_at REAL NOT NULL, + embedding BLOB ); CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_app_scope_user_time @@ -1169,6 +1198,21 @@ def _search_entries_simple( params = (*scope_params, pattern, limit) return self._fetch_records(sql, params) + def _search_entries_vector( + self, app_name: str, user_id: str, limit: int, scope_filter: Literal["all", "user", "app"] = "all" + ) -> "list[StoredMemory]": + where_scope, scope_params = _build_sqlite_scope_clause("", app_name, user_id, scope_filter) + sql = f""" + SELECT * + FROM {self._memory_table} + WHERE {where_scope} + AND embedding IS NOT NULL + ORDER BY timestamp DESC + LIMIT ? + """ + params = (*scope_params, limit) + return self._fetch_records(sql, params) + def _fetch_records(self, sql: str, params: "tuple[Any, ...]") -> "list[StoredMemory]": with self._config.provide_connection() as conn: self._ensure_foreign_keys(conn) @@ -1205,7 +1249,7 @@ def _fetch_records(self, sql: str, params: "tuple[Any, ...]") -> "list[StoredMem raw["inserted_at"] = _julian_to_datetime(raw["inserted_at"]) raw["content_json"] = from_json(raw["content_json"]) if raw["content_json"] else {} raw["metadata_json"] = from_json(raw["metadata_json"]) if raw["metadata_json"] else None - raw["embedding"] = None + raw["embedding"] = _unpack_embedding(raw.get("embedding")) records.append(cast("StoredMemory", raw)) return records @@ -1303,6 +1347,34 @@ def _build_sqlite_scope_clause( return f"{prefix}app_name = ? AND {prefix}scope = 'app'", (app_name,) +def _pack_embedding(raw: Any) -> bytes | None: + """Pack float sequence into float32 binary BLOB.""" + if raw is None: + return None + return array.array("f", [float(x) for x in raw]).tobytes() + + +def _unpack_embedding(blob: Any) -> list[float] | None: + """Unpack float32 binary BLOB into float list.""" + if blob is None: + return None + return list(array.array("f", blob)) + + +def _cosine_similarity(v1: "Sequence[float]", v2: "Sequence[float]") -> float: + """Compute cosine similarity between two float vectors.""" + dot = 0.0 + norm1 = 0.0 + norm2 = 0.0 + for a, b in zip(v1, v2, strict=False): + dot += a * b + norm1 += a * a + norm2 += b * b + if norm1 == 0.0 or norm2 == 0.0: + return 0.0 + return float(dot / (math.sqrt(norm1) * math.sqrt(norm2))) + + def _session_list_query( session_table: str, app_name: str, diff --git a/sqlspec/adapters/sqlite/config.py b/sqlspec/adapters/sqlite/config.py index b17cb99fe..b7a30de6a 100644 --- a/sqlspec/adapters/sqlite/config.py +++ b/sqlspec/adapters/sqlite/config.py @@ -76,7 +76,7 @@ class SqliteConnectionParams(TypedDict): enable_optimizations: NotRequired[bool] """Apply recommended SQLite ``PRAGMA`` settings (WAL journal mode, synchronous=NORMAL, busy_timeout) on connect.""" enable_foreign_keys: NotRequired[bool] - """Execute ``PRAGMA foreign_keys = ON`` when initializing each connection (defaults to ``True``).""" + """Execute ``PRAGMA foreign_keys = ON`` when initializing each connection (defaults to ``False``).""" extra: NotRequired[dict[str, Any]] """Additional keyword arguments forwarded to ``sqlite3.connect``.""" @@ -251,7 +251,7 @@ def __init__( if database is not None: config_dict["database"] = database if "database" not in config_dict or config_dict["database"] == ":memory:": - config_dict["database"] = f"file:memory_{uuid4().hex}?mode=memory&cache=private" + config_dict["database"] = f"file:memory_{uuid4().hex}?mode=memory&cache=shared" config_dict["uri"] = True elif "database" in config_dict: database_path = str(config_dict["database"]) diff --git a/sqlspec/adapters/sqlite/core.py b/sqlspec/adapters/sqlite/core.py index fe90822bf..aecdbb868 100644 --- a/sqlspec/adapters/sqlite/core.py +++ b/sqlspec/adapters/sqlite/core.py @@ -6,7 +6,7 @@ from collections.abc import Mapping from datetime import UTC, date, datetime, timedelta from decimal import Decimal -from typing import TYPE_CHECKING, Any, cast +from typing import TYPE_CHECKING, Any, Final, cast from sqlglot import exp @@ -45,6 +45,7 @@ "SECONDS_PER_DAY", "SQLITE_CONNECT_SUPPORTS_AUTOCOMMIT", "SqliteStreamSource", + "_filter_connection_params", "apply_driver_features", "apply_extension_pragmas", "build_connection_config", @@ -75,14 +76,15 @@ SQLITE_CONSTRAINT_PRIMARYKEY_CODE = 1555 SQLITE_CONSTRAINT_ROWID_CODE = 2579 SQLITE_CONSTRAINT_FOREIGNKEY_CODE = 787 -SQLITE_CONSTRAINT_NOTNULL_CODE = 1811 -SQLITE_CONSTRAINT_CHECK_CODE = 531 +SQLITE_CONSTRAINT_NOTNULL_CODE = 1299 +SQLITE_CONSTRAINT_CHECK_CODE = 275 SQLITE_CONSTRAINT_DATATYPE_CODE = 3091 SQLITE_CONSTRAINT_CODE = 19 SQLITE_CANTOPEN_CODE = 14 SQLITE_IOERR_CODE = 10 SQLITE_CORRUPT_CODE = 11 SQLITE_MISMATCH_CODE = 20 +SQLITE_MISUSE_CODE = 21 SQLITE_BUSY_CODE = 5 SQLITE_LOCKED_CODE = 6 SQLITE_INTERRUPT_CODE = 9 @@ -101,6 +103,8 @@ "PRAGMA mmap_size = 30000000", "PRAGMA journal_size_limit = 67108864", ) +_UNIX_EPOCH: Final[datetime] = datetime(1970, 1, 1, tzinfo=UTC) +_DEFAULT_SCHEMAS: Final[frozenset[str]] = frozenset({"temp", "main"}) def datetime_to_julian(dt: "datetime | None") -> "float | None": @@ -109,8 +113,7 @@ def datetime_to_julian(dt: "datetime | None") -> "float | None": return None if dt.tzinfo is None: dt = dt.replace(tzinfo=UTC) - epoch = datetime(1970, 1, 1, tzinfo=UTC) - delta_days = (dt - epoch).total_seconds() / SECONDS_PER_DAY + delta_days = (dt - _UNIX_EPOCH).total_seconds() / SECONDS_PER_DAY return JULIAN_EPOCH + delta_days @@ -120,7 +123,7 @@ def julian_to_datetime(julian: "float | None") -> "datetime | None": return None delta_days = julian - JULIAN_EPOCH timestamp = delta_days * SECONDS_PER_DAY - return datetime(1970, 1, 1, tzinfo=UTC) + timedelta(seconds=timestamp) + return _UNIX_EPOCH + timedelta(seconds=timestamp) def end_transaction( @@ -404,6 +407,32 @@ def normalize_execute_many_parameters(parameters: Any) -> Any: return parameters +_EXCLUDED_CONNECTION_KEYS: Final[frozenset[str]] = frozenset({ + "enable_foreign_keys", + "enable_optimizations", + "health_check_interval", + "pool_min_size", + "pool_max_size", + "pool_timeout", + "pool_recycle_seconds", + "extra", + "db", + "path", + "file", +}) + + +def _filter_connection_params(mapping: Mapping[str, Any]) -> dict[str, Any]: + """Filter connection parameters for SQLite.""" + return { + key: value + for key, value in mapping.items() + if key not in _EXCLUDED_CONNECTION_KEYS + and (value is not None or key == "isolation_level") + and (key != "autocommit" or SQLITE_CONNECT_SUPPORTS_AUTOCOMMIT) + } + + def build_connection_config(connection_config: "Mapping[str, Any]") -> "dict[str, Any]": """Build connection configuration for pool creation. @@ -413,30 +442,7 @@ def build_connection_config(connection_config: "Mapping[str, Any]") -> "dict[str Returns: Dictionary with connection parameters. """ - excluded_keys = { - "enable_foreign_keys", - "enable_optimizations", - "health_check_interval", - "pool_min_size", - "pool_max_size", - "pool_timeout", - "pool_recycle_seconds", - "extra", - "db", - "path", - "file", - } - - def _filter_params(mapping: Mapping[str, Any]) -> dict[str, Any]: - return { - key: value - for key, value in mapping.items() - if key not in excluded_keys - and (value is not None or key == "isolation_level") - and (key != "autocommit" or SQLITE_CONNECT_SUPPORTS_AUTOCOMMIT) - } - - connection_parameters = _filter_params(connection_config) + connection_parameters = _filter_connection_params(connection_config) if "database" not in connection_parameters: database = ( @@ -450,7 +456,7 @@ def _filter_params(mapping: Mapping[str, Any]) -> dict[str, Any]: extra = connection_config.get("extra") if isinstance(extra, Mapping): - connection_parameters.update(_filter_params(extra)) + connection_parameters.update(_filter_connection_params(extra)) return connection_parameters @@ -475,7 +481,7 @@ def create_mapped_exception(error: BaseException, *, logger: Any | None = None) Returns: A SQLSpec exception that wraps the original error """ - del logger + _ = logger if has_sqlite_error(error): error_code = error.sqlite_errorcode error_name = error.sqlite_errorname @@ -493,13 +499,19 @@ def create_mapped_exception(error: BaseException, *, logger: Any | None = None) return _create_sqlite_error(error, error_code, DeadlockError, "database busy") if primary_code == SQLITE_LOCKED_CODE or (isinstance(error_name, str) and error_name.startswith("SQLITE_LOCKED")): return _create_sqlite_error(error, error_code, DeadlockError, "database locked") - if "locked" in error_msg or "busy" in error_msg: - return _create_sqlite_error(error, error_code or 0, DeadlockError, "database locked") + if error_code is None and ( + "database is locked" in error_msg + or "database table is locked" in error_msg + or "database locked" in error_msg + or "database is busy" in error_msg + or "database busy" in error_msg + ): + return _create_sqlite_error(error, 0, DeadlockError, "database locked") if primary_code == SQLITE_INTERRUPT_CODE or error_name == "SQLITE_INTERRUPT": return _create_sqlite_error(error, error_code, OperationCancelledError, "query interrupted") - if "interrupt" in error_msg: - return _create_sqlite_error(error, error_code or 0, OperationCancelledError, "query interrupted") + if error_code is None and ("query was interrupted" in error_msg or "interrupted" in error_msg): + return _create_sqlite_error(error, 0, OperationCancelledError, "query interrupted") if primary_code == SQLITE_PERM_CODE or error_name == "SQLITE_PERM": return _create_sqlite_error(error, error_code, PermissionDeniedError, "permission denied") @@ -507,8 +519,13 @@ def create_mapped_exception(error: BaseException, *, logger: Any | None = None) isinstance(error_name, str) and error_name.startswith("SQLITE_READONLY") ): return _create_sqlite_error(error, error_code, PermissionDeniedError, "database is read-only") - if "permission denied" in error_msg or "readonly" in error_msg: - return _create_sqlite_error(error, error_code or 0, PermissionDeniedError, "permission denied") + if error_code is None and ( + "attempt to write a readonly database" in error_msg + or "readonly database" in error_msg + or "permission denied" in error_msg + or "access permission denied" in error_msg + ): + return _create_sqlite_error(error, 0, PermissionDeniedError, "permission denied") if primary_code == SQLITE_CANTOPEN_CODE or ( isinstance(error_name, str) and error_name.startswith("SQLITE_CANTOPEN") @@ -522,9 +539,16 @@ def create_mapped_exception(error: BaseException, *, logger: Any | None = None) if primary_code == SQLITE_MISMATCH_CODE or error_name == "SQLITE_MISMATCH": return _create_sqlite_error(error, error_code, DataError, "data error") + if primary_code == SQLITE_MISUSE_CODE or error_name == "SQLITE_MISUSE": + return _create_sqlite_error(error, error_code, DataError, "programming error") + if "syntax" in error_msg: return _create_sqlite_error(error, error_code, SQLParsingError, "SQL syntax error") + if isinstance(error, sqlite3.ProgrammingError): + return _create_sqlite_error(error, error_code, DataError, "programming error") + if isinstance(error, sqlite3.InterfaceError): + return _create_sqlite_error(error, error_code, OperationalError, "interface error") if primary_code == 1 or isinstance(error, sqlite3.OperationalError): return _create_sqlite_error(error, error_code, OperationalError, "operational error") if isinstance(error, sqlite3.DataError): @@ -644,7 +668,7 @@ def _target_supports_rowid(connection: Any, target: "tuple[str | None, str]") -> return _target_supports_rowid_legacy(connection, target) schema_order = ["temp", "main"] - if not any(row[0].casefold() in {"temp", "main"} for row in candidates): + if not any(row[0].casefold() in _DEFAULT_SCHEMAS for row in candidates): try: database_cursor = connection.execute("PRAGMA database_list") try: diff --git a/sqlspec/adapters/sqlite/data_dictionary.py b/sqlspec/adapters/sqlite/data_dictionary.py index 89a43a65a..6267502d8 100644 --- a/sqlspec/adapters/sqlite/data_dictionary.py +++ b/sqlspec/adapters/sqlite/data_dictionary.py @@ -26,6 +26,7 @@ sqlite_schema_ddl_results, ) from sqlspec.driver import SyncDataDictionaryBase +from sqlspec.exceptions import SQLFileNotFoundError if TYPE_CHECKING: from sqlspec.adapters.sqlite.driver import SqliteDriver @@ -138,10 +139,13 @@ def get_columns( rows = driver.select(query_text, table_name=table, schema_name=schema_name) if any(str(row.get("extra") or "").endswith("generated") for row in rows): table_sql = self._get_table_sql(driver, schema_name, table) + parsed_expressions: dict[str, dict[str, str]] = {} for row in rows: if str(row.get("extra") or "").endswith("generated"): - expressions = sqlite_generated_expressions(table_sql.get(str(row.get("table_name")))) - expression = expressions.get(str(row.get("column_name"))) + table_name_str = str(row.get("table_name")) + if table_name_str not in parsed_expressions: + parsed_expressions[table_name_str] = sqlite_generated_expressions(table_sql.get(table_name_str)) + expression = parsed_expressions[table_name_str].get(str(row.get("column_name"))) if expression is not None: row["generation_expression"] = expression return cast("list[ColumnMetadata]", rows) @@ -270,7 +274,7 @@ def _get_domain_query_text(self, domain: str, query_name: str) -> str: query_text = get_data_dictionary_loader().get_domain_query_text(type(self).dialect, domain, query_name) if query_text is None: msg = f"Missing SQLite data-dictionary query: {domain}/{query_name}" - raise RuntimeError(msg) + raise SQLFileNotFoundError(msg) return query_text diff --git a/sqlspec/adapters/sqlite/driver.py b/sqlspec/adapters/sqlite/driver.py index 1729a5138..4f32afe2b 100644 --- a/sqlspec/adapters/sqlite/driver.py +++ b/sqlspec/adapters/sqlite/driver.py @@ -1,7 +1,7 @@ """SQLite driver implementation.""" import contextlib -from typing import TYPE_CHECKING, Any, Literal, cast +from typing import TYPE_CHECKING, Any, Final, Literal, cast from mypy_extensions import mypyc_attr @@ -15,6 +15,7 @@ create_mapped_exception, default_statement_config, driver_profile, + end_transaction, format_identifier, normalize_execute_many_parameters, normalize_execute_parameters, @@ -25,9 +26,11 @@ from sqlspec.core.cache import get_cache_config from sqlspec.core.parameters._registry import register_driver_profile from sqlspec.core.result._base import ArrowResult, DMLResult +from sqlspec.driver._common import validate_savepoint_name from sqlspec.driver._exception_handler import BaseSyncExceptionHandler from sqlspec.driver._stream import SyncRowStream from sqlspec.driver._sync import SyncDriverAdapterBase +from sqlspec.exceptions import ImproperConfigurationError from sqlspec.utils.type_guards import resolve_row_format if TYPE_CHECKING: @@ -44,6 +47,8 @@ __all__ = ("SqliteCursor", "SqliteDriver", "SqliteExceptionHandler", "SqliteSessionContext") +_ROWID_SAFE_OPERATIONS: Final[frozenset[str]] = frozenset({"SELECT", "INSERT", "UPDATE", "DELETE"}) + @mypyc_attr(allow_interpreted_subclasses=True) class SqliteExceptionHandler(BaseSyncExceptionHandler): @@ -175,28 +180,12 @@ def dispatch_execute_many(self, cursor: Any, statement: "SQL") -> "ExecutionResu return self.create_execution_result(cursor, rowcount_override=affected_rows, is_many_result=True) def dispatch_execute_script(self, cursor: Any, statement: "SQL") -> "ExecutionResult": - """Execute SQL script with statement splitting and parameter handling. - - Args: - cursor: SQLite cursor object - statement: SQL statement containing multiple statements - - Returns: - ExecutionResult with script execution details - """ + """Execute SQL script with statement splitting and parameter handling.""" self._rowid_target_cache.clear() - sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) - statements = self.split_script_statements(sql, statement.statement_config, strip_trailing_semicolon=True) - - successful_count = 0 - - for stmt in statements: - cursor.execute(stmt, normalize_execute_parameters(prepared_parameters)) - successful_count += 1 - - return self.create_execution_result( - cursor, statement_count=len(statements), successful_statements=successful_count, is_script_result=True - ) + try: + return super().dispatch_execute_script(cursor, statement) + finally: + self._rowid_target_cache.clear() def begin(self, mode: "Literal['DEFERRED', 'IMMEDIATE', 'EXCLUSIVE'] | None" = None) -> None: """Begin a database transaction. @@ -211,13 +200,16 @@ def begin(self, mode: "Literal['DEFERRED', 'IMMEDIATE', 'EXCLUSIVE'] | None" = N transaction_mode = mode if mode is not None else self.driver_features.get("default_transaction_mode") if transaction_mode is not None and transaction_mode not in {"DEFERRED", "IMMEDIATE", "EXCLUSIVE"}: msg = "Transaction mode must be DEFERRED, IMMEDIATE, or EXCLUSIVE" - raise ValueError(msg) + raise ImproperConfigurationError(msg) try: if not self.connection.in_transaction: stmt = f"BEGIN {transaction_mode}" if transaction_mode else "BEGIN" cursor = self.connection.execute(stmt) - with contextlib.suppress(Exception): - cursor.close() + with contextlib.suppress(AttributeError, TypeError): + self.connection.in_transaction = True + if cursor is not None: + with contextlib.suppress(Exception): + cursor.close() except sqlite3.Error as e: raise create_mapped_exception(e) from e @@ -228,13 +220,9 @@ def commit(self) -> None: SQLSpecError: If transaction cannot be committed """ try: - if self._in_autocommit_mode(): - if self.connection.in_transaction: - cursor = self.connection.execute("COMMIT") - with contextlib.suppress(Exception): - cursor.close() - return - self.connection.commit() + end_transaction(self.connection, commit=True, supports_autocommit=self._in_autocommit_mode()) + with contextlib.suppress(AttributeError, TypeError): + self.connection.in_transaction = False except sqlite3.Error as e: raise create_mapped_exception(e) from e @@ -245,19 +233,36 @@ def rollback(self) -> None: SQLSpecError: If transaction cannot be rolled back """ try: - if self._in_autocommit_mode(): - if self.connection.in_transaction: - cursor = self.connection.execute("ROLLBACK") - with contextlib.suppress(Exception): - cursor.close() - self._rowid_target_cache.clear() - return - self.connection.rollback() + end_transaction(self.connection, commit=False, supports_autocommit=self._in_autocommit_mode()) + with contextlib.suppress(AttributeError, TypeError): + self.connection.in_transaction = False self._rowid_target_cache.clear() except sqlite3.Error as e: self._rowid_target_cache.clear() raise create_mapped_exception(e) from e + def create_savepoint(self, name: str) -> None: + """Create a savepoint within the current transaction without script splitting.""" + safe_name = validate_savepoint_name(name) + try: + cursor = self.connection.execute(f"SAVEPOINT {safe_name}") + if cursor is not None: + with contextlib.suppress(Exception): + cursor.close() + except sqlite3.Error as exc: + raise create_mapped_exception(exc) from exc + + def release_savepoint(self, name: str) -> None: + """Release a savepoint without clearing rowid target cache or splitting script.""" + safe_name = validate_savepoint_name(name) + try: + cursor = self.connection.execute(f"RELEASE SAVEPOINT {safe_name}") + if cursor is not None: + with contextlib.suppress(Exception): + cursor.close() + except sqlite3.Error as exc: + raise create_mapped_exception(exc) from exc + def with_cursor(self, connection: "SqliteConnection") -> "SqliteCursor": """Create context manager for SQLite cursor. @@ -351,7 +356,7 @@ def load_from_arrow( insert_sql = build_insert_statement(table, columns) needs_prep = self._arrow_rows_need_preparation(arrow_table) - owns_transaction = not self.connection.in_transaction + owns_transaction = not self._connection_in_transaction() try: if owns_transaction: transaction_mode = cast( @@ -359,20 +364,19 @@ def load_from_arrow( self.driver_features.get("default_transaction_mode") or "IMMEDIATE", ) self.begin(mode=transaction_mode) - if overwrite: - statement = f"DELETE FROM {format_identifier(table)}" - with self.with_cursor(self.connection) as cursor: + with self.with_cursor(self.connection) as cursor: + if overwrite: + statement = f"DELETE FROM {format_identifier(table)}" cursor.execute(statement) - for batch in arrow_table.to_batches(max_chunksize=batch_size): - pydict = batch.to_pydict() - records = list(zip(*[pydict[col] for col in columns], strict=False)) - if records: - prepared_records = ( - self.prepare_driver_parameters(records, self.statement_config, is_many=True) - if needs_prep - else records - ) - with self.with_cursor(self.connection) as cursor: + for batch in arrow_table.to_batches(max_chunksize=batch_size): + pydict = batch.to_pydict() + records = list(zip(*[pydict[col] for col in columns], strict=False)) + if records: + prepared_records = ( + self.prepare_driver_parameters(records, self.statement_config, is_many=True) + if needs_prep + else records + ) cursor.executemany(insert_sql, cast("Any", prepared_records)) if owns_transaction: self.commit() @@ -388,22 +392,6 @@ def load_from_arrow( self._attach_partition_telemetry(telemetry_payload, partitioner) return self._storage_job(telemetry_payload, telemetry) - def load_from_storage( - self, - table: str, - source: "StorageDestination", - *, - file_format: "StorageFormat", - partitioner: "dict[str, object] | None" = None, - overwrite: bool = False, - **kwargs: Any, - ) -> "StorageBridgeJob": - """Load staged artifacts from storage into SQLite.""" - - return super().load_from_storage( - table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs - ) - def collect_rows(self, cursor: Any, fetched: "list[Any]") -> "tuple[list[Any], list[str], int]": """Collect SQLite rows for the direct execution path.""" return collect_rows(fetched, cursor.description) @@ -431,10 +419,11 @@ def _execute_cache_hit( """ cursor = None returns_rows = cached.operation_profile.returns_rows + normalized_parameters = normalize_execute_parameters(params) self._invalidate_rowid_target_cache(cached.operation_type) try: try: - cursor = self.connection.execute(cached.compiled_sql, params) + cursor = self.connection.execute(cached.compiled_sql, normalized_parameters) except sqlite3.Error as exc: raise create_mapped_exception(exc) from exc @@ -483,7 +472,7 @@ def _execute_cache_hit( cursor.close() def _invalidate_rowid_target_cache(self, operation_type: "OperationType") -> None: - if operation_type not in {"SELECT", "INSERT", "UPDATE", "DELETE"}: + if operation_type not in _ROWID_SAFE_OPERATIONS: self._rowid_target_cache.clear() def _connection_in_transaction(self) -> bool: diff --git a/sqlspec/adapters/sqlite/litestar/__init__.py b/sqlspec/adapters/sqlite/litestar/__init__.py index 99aeee476..4165b2867 100644 --- a/sqlspec/adapters/sqlite/litestar/__init__.py +++ b/sqlspec/adapters/sqlite/litestar/__init__.py @@ -1,5 +1,5 @@ """Litestar integration for SQLite adapter.""" -from sqlspec.adapters.sqlite.litestar.store import SqliteLitestarConfig, SQLiteStore +from sqlspec.adapters.sqlite.litestar.store import SqliteLitestarConfig, SqliteStore -__all__ = ("SQLiteStore", "SqliteLitestarConfig") +__all__ = ("SqliteLitestarConfig", "SqliteStore") diff --git a/sqlspec/adapters/sqlite/litestar/store.py b/sqlspec/adapters/sqlite/litestar/store.py index 01d3bee41..29d977cde 100644 --- a/sqlspec/adapters/sqlite/litestar/store.py +++ b/sqlspec/adapters/sqlite/litestar/store.py @@ -20,7 +20,7 @@ if TYPE_CHECKING: from sqlspec.adapters.sqlite.config import SqliteConfig -__all__ = ("JULIAN_EPOCH", "SECONDS_PER_DAY", "SQLiteStore", "SqliteLitestarConfig") +__all__ = ("JULIAN_EPOCH", "SECONDS_PER_DAY", "SqliteLitestarConfig", "SqliteStore") class SqliteLitestarConfig(LitestarConfig): @@ -36,7 +36,7 @@ class SqliteLitestarConfig(LitestarConfig): """Validated PRAGMA overrides applied during schema preparation.""" -class SQLiteStore(BaseSQLSpecStore["SqliteConfig"]): +class SqliteStore(BaseSQLSpecStore["SqliteConfig"]): """SQLite session store using synchronous SQLite driver. Implements server-side session storage for Litestar using SQLite diff --git a/sqlspec/adapters/sqlite/pool.py b/sqlspec/adapters/sqlite/pool.py index a2de6cbf0..b69711d59 100644 --- a/sqlspec/adapters/sqlite/pool.py +++ b/sqlspec/adapters/sqlite/pool.py @@ -39,6 +39,8 @@ class SqliteConnectionPool: __slots__ = ( "_anchor_connection", + "_checked_out_count", + "_connection_depths", "_connection_parameters", "_connection_registry", "_enable_foreign_keys", @@ -52,6 +54,7 @@ class SqliteConnectionPool: "_recycle_seconds", "_registry_lock", "_runtime_setup", + "_thread_connections", "_thread_local", ) @@ -86,6 +89,9 @@ def __init__( self._anchor_connection: SqliteConnection | None = None self._thread_local = threading.local() self._connection_registry: set[SqliteConnection] = set() + self._thread_connections: dict[threading.Thread, SqliteConnection] = {} + self._checked_out_count = 0 + self._connection_depths: dict[SqliteConnection, int] = {} self._generation = 0 self._registry_lock = threading.Lock() self._enable_optimizations = enable_optimizations @@ -122,8 +128,6 @@ def new_connection(self) -> SqliteConnection: _execute_and_close(connection, "PRAGMA synchronous = OFF") _execute_and_close(connection, "PRAGMA temp_store = MEMORY") _execute_and_close(connection, f"PRAGMA cache_size = {SQLITE_MEMORY_CACHE_SIZE}") - if self._is_shared_memory: - _execute_and_close(connection, "PRAGMA read_uncommitted = ON") else: _enable_wal(connection) _execute_and_close(connection, "PRAGMA synchronous = NORMAL") @@ -153,19 +157,36 @@ def get_connection(self) -> "Generator[SqliteConnection, None, None]": SqliteConnection: A thread-local connection. """ connection = self._get_thread_connection() + with self._registry_lock: + depth = self._connection_depths.get(connection, 0) + self._connection_depths[connection] = depth + 1 + if depth == 0: + self._checked_out_count += 1 + outermost = depth == 0 try: yield connection except BaseException: - with contextlib.suppress(Exception): - end_transaction(connection, commit=False) - raise - else: - try: - end_transaction(connection, commit=True) - except BaseException: + if outermost: with contextlib.suppress(Exception): end_transaction(connection, commit=False) - raise + raise + else: + if outermost: + try: + end_transaction(connection, commit=True) + except BaseException: + with contextlib.suppress(Exception): + end_transaction(connection, commit=False) + raise + finally: + with self._registry_lock: + remaining = self._connection_depths.get(connection, 1) - 1 + if remaining > 0: + self._connection_depths[connection] = remaining + else: + self._connection_depths.pop(connection, None) + if outermost and self._checked_out_count > 0: + self._checked_out_count -= 1 def close(self) -> None: """Close every connection this pool opened, on any thread.""" @@ -173,6 +194,9 @@ def close(self) -> None: with self._registry_lock: orphaned = list(self._connection_registry) self._connection_registry.clear() + self._thread_connections.clear() + self._connection_depths.clear() + self._checked_out_count = 0 anchor = self._anchor_connection self._anchor_connection = None self._generation += 1 @@ -200,14 +224,27 @@ def release(self, connection: SqliteConnection) -> None: with contextlib.suppress(Exception): end_transaction(connection, commit=False) + def _prune_dead_threads(self) -> None: + """Prune connections owned by dead threads.""" + with self._registry_lock: + dead_threads = [t for t in self._thread_connections if not t.is_alive()] + for t in dead_threads: + conn = self._thread_connections.pop(t, None) + if conn is not None: + self._connection_registry.discard(conn) + with contextlib.suppress(Exception): + conn.close() + def size(self) -> int: """Get total number of active thread-local connections in pool.""" + self._prune_dead_threads() with self._registry_lock: return len(self._connection_registry) def checked_out(self) -> int: - """Get number of checked out connections (always 0).""" - return 0 + """Get number of checked out connections.""" + with self._registry_lock: + return self._checked_out_count def _create_connection(self) -> SqliteConnection: """Create a pool-owned connection and record it for shutdown.""" @@ -257,6 +294,8 @@ def _get_thread_connection(self) -> SqliteConnection: self._thread_local.connection = conn self._thread_local.created_at = now self._thread_local.last_used = now + with self._registry_lock: + self._thread_connections[threading.current_thread()] = conn return conn created_at = getattr(self._thread_local, "created_at", 0.0) @@ -276,6 +315,8 @@ def _get_thread_connection(self) -> SqliteConnection: self._thread_local.connection = conn self._thread_local.created_at = now self._thread_local.last_used = now + with self._registry_lock: + self._thread_connections[threading.current_thread()] = conn return conn last_used = getattr(self._thread_local, "last_used", 0.0) @@ -296,15 +337,22 @@ def _get_thread_connection(self) -> SqliteConnection: self._thread_local.connection = conn self._thread_local.created_at = now self._thread_local.last_used = now + with self._registry_lock: + self._thread_connections[threading.current_thread()] = conn return conn self._thread_local.last_used = now + with self._registry_lock: + self._thread_connections[threading.current_thread()] = conn return cast("SqliteConnection", conn) def _retire_connection(self, connection: SqliteConnection) -> None: """Close a pool-owned connection and drop it from the shutdown registry.""" with self._registry_lock: self._connection_registry.discard(connection) + dead_threads = [t for t, c in self._thread_connections.items() if c is connection] + for t in dead_threads: + self._thread_connections.pop(t, None) with contextlib.suppress(Exception): connection.close() diff --git a/sqlspec/adapters/sqlite/type_converter.py b/sqlspec/adapters/sqlite/type_converter.py index 74f1c7d91..d47069d4a 100644 --- a/sqlspec/adapters/sqlite/type_converter.py +++ b/sqlspec/adapters/sqlite/type_converter.py @@ -13,15 +13,12 @@ from typing import TYPE_CHECKING, Any from sqlspec.adapters.sqlite._typing import sqlite_module as sqlite3 -from sqlspec.utils.logging import get_logger if TYPE_CHECKING: from collections.abc import Callable __all__ = ("json_adapter", "json_converter", "register_type_handlers", "unregister_type_handlers") -logger = get_logger(__name__) - DEFAULT_JSON_TYPE = "JSON" diff --git a/sqlspec/adapters/sybase/_typing.py b/sqlspec/adapters/sybase/_typing.py index 7b8826d9d..fd5a2c1ab 100644 --- a/sqlspec/adapters/sybase/_typing.py +++ b/sqlspec/adapters/sybase/_typing.py @@ -5,39 +5,30 @@ """ import contextlib +import importlib from typing import TYPE_CHECKING, Any +import pymssql as PYMSSQL_MODULE # noqa: N812 from mypy_extensions import mypyc_attr - -from sqlspec.utils.module_loader import import_optional_attr - -PYMSSQL_MODULE: Any = None -_DRIVER_CLS: Any = None - - -class SybaseError(Exception): - """Sybase ASE driver exception base.""" - +from pymssql import Connection, Cursor +from pymssql import Error as SybaseError if TYPE_CHECKING: from collections.abc import Callable from types import TracebackType from typing import TypeAlias - from pymssql import Connection as _SybaseConnection - from pymssql import Cursor as _SybaseRawCursor from pymssql._pymssql import QueryParams as SybaseQueryParams from sqlspec.adapters.sybase.driver import SybaseDriver from sqlspec.core import StatementConfig - SybaseConnection: TypeAlias = _SybaseConnection - SybaseRawCursor: TypeAlias = _SybaseRawCursor - -if not TYPE_CHECKING: + SybaseConnection: TypeAlias = Connection + SybaseRawCursor: TypeAlias = Cursor +else: SybaseQueryParams = Any - SybaseConnection = Any - SybaseRawCursor = Any + SybaseConnection = Connection + SybaseRawCursor = Cursor __all__ = ( "PYMSSQL_MODULE", @@ -80,6 +71,9 @@ def __exit__(self, *_: Any) -> None: self.cursor.close() +_DRIVER_CLS: "type[SybaseDriver] | None" = None + + @mypyc_attr(allow_interpreted_subclasses=True, native_class=False) class SybaseSessionContext: """Sync context manager for Sybase ASE sessions.""" @@ -113,7 +107,9 @@ def __init__( def __enter__(self) -> "SybaseDriver": global _DRIVER_CLS if _DRIVER_CLS is None: - _DRIVER_CLS = import_optional_attr("sqlspec.adapters.sybase.driver", "SybaseDriver") + module = importlib.import_module("sqlspec.adapters.sybase.driver") + _DRIVER_CLS = module.SybaseDriver + assert _DRIVER_CLS is not None self._connection = self._acquire_connection() self._driver = _DRIVER_CLS( connection=self._connection, statement_config=self._statement_config, driver_features=self._driver_features diff --git a/sqlspec/adapters/sybase/adk/__init__.py b/sqlspec/adapters/sybase/adk/__init__.py index c134a08d0..696f9b01b 100644 --- a/sqlspec/adapters/sybase/adk/__init__.py +++ b/sqlspec/adapters/sybase/adk/__init__.py @@ -1,5 +1,5 @@ """Sybase ASE ADK extension.""" -from sqlspec.adapters.sybase.adk.store import SybaseADKConfig, SybaseADKMemoryStore, SybaseADKStore +from sqlspec.adapters.sybase.adk.store import SybaseADKMemoryStore, SybaseADKStore -__all__ = ("SybaseADKConfig", "SybaseADKMemoryStore", "SybaseADKStore") +__all__ = ("SybaseADKMemoryStore", "SybaseADKStore") diff --git a/sqlspec/adapters/sybase/adk/store.py b/sqlspec/adapters/sybase/adk/store.py index d009c1f78..b9d7110c7 100644 --- a/sqlspec/adapters/sybase/adk/store.py +++ b/sqlspec/adapters/sybase/adk/store.py @@ -1,17 +1,19 @@ """Sybase ASE ADK stores for Google Agent Development Kit session storage.""" +import contextlib from datetime import datetime from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, cast from sqlspec.adapters.sybase._typing import SybaseCursor, SybaseError from sqlspec.adapters.sybase.core import ( + create_mapped_exception, escape_sql_literal, extract_error_number, quote_bracket_identifier, resolve_rowcount, split_schema_and_table, ) -from sqlspec.config import ADKConfig +from sqlspec.exceptions import NotFoundError, SQLSpecError from sqlspec.extensions.adk import BaseSyncADKStore, StoredEvent, StoredSession, normalize_session_list_options from sqlspec.extensions.adk.memory.store import BaseSyncADKMemoryStore from sqlspec.utils.serializers import from_json, to_json @@ -25,23 +27,21 @@ from sqlspec.extensions.adk import SessionOrderBy from sqlspec.extensions.adk.memory._types import StoredMemory -__all__ = ("SybaseADKConfig", "SybaseADKMemoryStore", "SybaseADKStore") +__all__ = ("SybaseADKMemoryStore", "SybaseADKStore") SYBASE_TABLE_NOT_FOUND_ERROR: Final[int] = 208 SYBASE_SCHEMA: Final[str] = "dbo" JSON_COLUMN_TYPE: Final[str] = "TEXT" -class SybaseADKConfig(ADKConfig): - """Sybase ASE ADK extension settings.""" - - class SybaseADKStore(BaseSyncADKStore["SybaseConfig"]): """Synchronous Sybase ASE ADK session/event store.""" connector_name: ClassVar[str] = "sybase" __slots__ = () + _max_table_name_parts = 2 + def __init__(self, config: "SybaseConfig") -> None: super().__init__(config) @@ -117,7 +117,7 @@ def get_session( """, (app_name, user_id, session_id), ) - except SybaseError as exc: + except (SybaseError, SQLSpecError) as exc: if _is_sybase_table_missing(exc): return None raise @@ -155,7 +155,7 @@ def list_sessions( ) try: rows = self._execute_fetchall(sql, params) - except SybaseError as exc: + except (SybaseError, SQLSpecError) as exc: if _is_sybase_table_missing(exc): return [] raise @@ -237,7 +237,7 @@ def get_events( sql, params = self._events_query(app_name, user_id, session_id, after_timestamp, limit) try: rows = self._execute_fetchall(sql, params) - except SybaseError as exc: + except (SybaseError, SQLSpecError) as exc: if _is_sybase_table_missing(exc): return [] raise @@ -252,7 +252,7 @@ def delete_expired_events(self, before: datetime, app_name: "str | None" = None) params.append(app_name) try: return self._execute(sql, tuple(params), commit=True) - except SybaseError as exc: + except (SybaseError, SQLSpecError) as exc: if _is_sybase_table_missing(exc): return 0 raise @@ -266,7 +266,7 @@ def delete_idle_sessions(self, updated_before: datetime, app_name: "str | None" params.append(app_name) try: return self._execute(sql, tuple(params), commit=True) - except SybaseError as exc: + except (SybaseError, SQLSpecError) as exc: if _is_sybase_table_missing(exc): return 0 raise @@ -280,7 +280,7 @@ def delete_idle_user_states(self, updated_before: datetime, app_name: "str | Non params.append(app_name) try: return self._execute(sql, tuple(params), commit=True) - except SybaseError as exc: + except (SybaseError, SQLSpecError) as exc: if _is_sybase_table_missing(exc): return 0 raise @@ -291,7 +291,7 @@ def get_app_state(self, app_name: str) -> "dict[str, Any] | None": row = self._execute_fetchone( f"SELECT TOP 1 state FROM {_table_ref(self._app_state_table)} WHERE app_name = %s", (app_name,) ) - except SybaseError as exc: + except (SybaseError, SQLSpecError) as exc: if _is_sybase_table_missing(exc): return None raise @@ -308,7 +308,7 @@ def get_user_state(self, app_name: str, user_id: str) -> "dict[str, Any] | None" """, (app_name, user_id), ) - except SybaseError as exc: + except (SybaseError, SQLSpecError) as exc: if _is_sybase_table_missing(exc): return None raise @@ -334,7 +334,7 @@ def get_metadata(self, key: str) -> "str | None": row = self._execute_fetchone( f"SELECT TOP 1 value FROM {_table_ref(self._metadata_table)} WHERE [key] = %s", (key,) ) - except SybaseError as exc: + except (SybaseError, SQLSpecError) as exc: if _is_sybase_table_missing(exc): return None raise @@ -409,23 +409,38 @@ def _events_query( def _execute_fetchone(self, sql: str, params: "tuple[Any, ...]" = (), *, commit: bool = False) -> "Any | None": with self._config.provide_connection() as conn, SybaseCursor(conn, as_dict=False) as cursor: - cursor.execute(sql, params) - row = cursor.fetchone() - if commit: - conn.commit() + try: + cursor.execute(sql, params) + row = cursor.fetchone() + if commit: + conn.commit() + except Exception as exc: + if commit: + with contextlib.suppress(Exception): + conn.rollback() + raise create_mapped_exception(exc) from exc return row def _execute_fetchall(self, sql: str, params: "tuple[Any, ...]" = ()) -> "list[Any]": with self._config.provide_connection() as conn, SybaseCursor(conn, as_dict=False) as cursor: - cursor.execute(sql, params) - return list(cursor.fetchall()) + try: + cursor.execute(sql, params) + return list(cursor.fetchall()) + except Exception as exc: + raise create_mapped_exception(exc) from exc def _execute(self, sql: str, params: "tuple[Any, ...]" = (), *, commit: bool = False) -> int: with self._config.provide_connection() as conn, SybaseCursor(conn, as_dict=False) as cursor: - cursor.execute(sql, params) - rowcount = resolve_rowcount(cursor) - if commit: - conn.commit() + try: + cursor.execute(sql, params) + rowcount = resolve_rowcount(cursor) + if commit: + conn.commit() + except Exception as exc: + if commit: + with contextlib.suppress(Exception): + conn.rollback() + raise create_mapped_exception(exc) from exc return rowcount @@ -434,6 +449,8 @@ class SybaseADKMemoryStore(BaseSyncADKMemoryStore["SybaseConfig"]): __slots__ = () + _max_table_name_parts = 2 + def __init__(self, config: "SybaseConfig") -> None: super().__init__(config) @@ -498,9 +515,10 @@ def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object params.append(entry["event_id"]) cursor.execute(sql, tuple(params)) inserted += resolve_rowcount(cursor) - except Exception: - conn.rollback() - raise + except Exception as exc: + with contextlib.suppress(Exception): + conn.rollback() + raise create_mapped_exception(exc) from exc conn.commit() return inserted @@ -594,15 +612,24 @@ def _drop_memory_table_sql(self) -> "list[str]": def _execute_fetchall(self, sql: str, params: "tuple[Any, ...]" = ()) -> "list[Any]": with self._config.provide_connection() as conn, SybaseCursor(conn, as_dict=False) as cursor: - cursor.execute(sql, params) - return list(cursor.fetchall()) + try: + cursor.execute(sql, params) + return list(cursor.fetchall()) + except Exception as exc: + raise create_mapped_exception(exc) from exc def _execute(self, sql: str, params: "tuple[Any, ...]" = (), *, commit: bool = False) -> int: with self._config.provide_connection() as conn, SybaseCursor(conn, as_dict=False) as cursor: - cursor.execute(sql, params) - rowcount = resolve_rowcount(cursor) - if commit: - conn.commit() + try: + cursor.execute(sql, params) + rowcount = resolve_rowcount(cursor) + if commit: + conn.commit() + except Exception as exc: + if commit: + with contextlib.suppress(Exception): + conn.rollback() + raise create_mapped_exception(exc) from exc return rowcount @@ -863,11 +890,15 @@ def _json_dict(value: Any) -> "dict[str, Any]": def _is_sybase_table_missing(exc: BaseException) -> bool: - text = str(exc).lower() + target = getattr(exc, "__cause__", None) or exc + text = f"{exc!s} {target!s}".lower() return ( "not found" in text or "invalid object name" in text or extract_error_number(exc) == SYBASE_TABLE_NOT_FOUND_ERROR + or extract_error_number(target) == SYBASE_TABLE_NOT_FOUND_ERROR + or isinstance(exc, NotFoundError) + or isinstance(target, NotFoundError) ) diff --git a/sqlspec/adapters/sybase/core.py b/sqlspec/adapters/sybase/core.py index 37d946fcf..fe597834b 100644 --- a/sqlspec/adapters/sybase/core.py +++ b/sqlspec/adapters/sybase/core.py @@ -140,13 +140,13 @@ def format_identifier(identifier: str) -> str: msg = "Table name must not be empty" raise SQLSpecError(msg) parts = split_qualified_identifier(cleaned, quote_chars='"', allow_bracket_quotes=True) - return ".".join(quote_bracket_identifier(part) for part in parts) + return ".".join([quote_bracket_identifier(part) for part in parts]) def build_insert_statement(table: str, columns: "list[str]") -> str: """Build a pymssql-compatible INSERT statement for Sybase ASE.""" - column_clause = ", ".join(quote_bracket_identifier(column) for column in columns) - placeholders = ", ".join("%s" for _ in columns) + column_clause = ", ".join([quote_bracket_identifier(column) for column in columns]) + placeholders = ", ".join(["%s"] * len(columns)) if columns else "" return f"INSERT INTO {format_identifier(table)} ({column_clause}) VALUES ({placeholders})" @@ -155,7 +155,7 @@ def escape_literal_percent(sql: str, parameters: Any, validator: "ParameterValid if not parameters or "%" not in sql: return sql keep = {info.position for info in validator.extract_parameters(sql)} - return "".join("%%" if char == "%" and index not in keep else char for index, char in enumerate(sql)) + return "".join(["%%" if char == "%" and index not in keep else char for index, char in enumerate(sql)]) def normalize_execute_parameters(parameters: Any) -> Any: diff --git a/sqlspec/adapters/sybase/driver.py b/sqlspec/adapters/sybase/driver.py index 8b8b1e17e..68e0fc5ee 100644 --- a/sqlspec/adapters/sybase/driver.py +++ b/sqlspec/adapters/sybase/driver.py @@ -1,16 +1,11 @@ """Sybase ASE database driver implementation over pymssql (FreeTDS TDS 5.0).""" import contextlib +import sys from collections.abc import Sized from typing import TYPE_CHECKING, Any, Literal, cast -from sqlspec.adapters.sybase._typing import ( - SybaseConnection, - SybaseCursor, - SybaseError, - SybaseRawCursor, - SybaseSessionContext, -) +from sqlspec.adapters.sybase._typing import SybaseConnection, SybaseCursor, SybaseRawCursor, SybaseSessionContext from sqlspec.adapters.sybase.core import ( build_insert_statement, collect_rows, @@ -61,7 +56,7 @@ class SybaseExceptionHandler(BaseSyncExceptionHandler): def _handle_exception(self, exc_type: "type[BaseException] | None", exc_val: "BaseException") -> bool: if exc_type is None: return False - if isinstance(exc_val, (SybaseError, Exception)) and not isinstance(exc_val, SQLSpecError): + if isinstance(exc_val, Exception) and not isinstance(exc_val, SQLSpecError): self.pending_exception = create_mapped_exception(exc_val, logger=logger) return True return False @@ -85,12 +80,15 @@ def start(self) -> None: try: cursor = cursor_manager.__enter__() handler = self._driver.handle_database_exceptions() - with handler: - normalized_params = normalize_execute_parameters(self._parameters) - sql = escape_literal_percent( - self._sql, normalized_params, self._driver.statement_config.parameter_validator - ) - cursor.execute(sql, normalized_params) + normalized_params = normalize_execute_parameters(self._parameters) + sql = escape_literal_percent( + self._sql, normalized_params, self._driver.statement_config.parameter_validator + ) + if sys.exc_info()[1] is None: + with handler: + cursor.execute(sql, normalized_params) + else: + self._driver._run_with_exception_handler(handler, cursor.execute, sql, normalized_params) self._driver._check_pending_exception(handler) except BaseException: with contextlib.suppress(Exception): @@ -105,8 +103,11 @@ def fetch_chunk(self) -> "list[dict[str, Any]]": cursor = cursor_manager.cursor handler = self._driver.handle_database_exceptions() rows: Any = [] - with handler: - rows = cursor.fetchmany(self._chunk_size) + if sys.exc_info()[1] is None: + with handler: + rows = cursor.fetchmany(self._chunk_size) + else: + rows = self._driver._run_with_exception_handler(handler, cursor.fetchmany, self._chunk_size) self._driver._check_pending_exception(handler) if not rows: return [] @@ -209,9 +210,10 @@ def begin(self) -> None: cursor.execute("BEGIN TRANSACTION") self._explicit_transaction = explicit self._transaction_active = True - except SybaseError as exc: + except Exception as exc: + mapped = create_mapped_exception(exc, logger=logger) msg = f"Failed to begin Sybase ASE transaction: {exc}" - raise SQLSpecError(msg) from exc + raise type(mapped)(msg) from exc def commit(self) -> None: try: @@ -222,9 +224,10 @@ def commit(self) -> None: self.connection.commit() self._explicit_transaction = False self._transaction_active = False - except SybaseError as exc: + except Exception as exc: + mapped = create_mapped_exception(exc, logger=logger) msg = f"Failed to commit Sybase ASE transaction: {exc}" - raise SQLSpecError(msg) from exc + raise type(mapped)(msg) from exc def rollback(self) -> None: try: @@ -233,9 +236,10 @@ def rollback(self) -> None: cursor.execute("IF @@TRANCOUNT > 0 ROLLBACK TRANSACTION") else: self.connection.rollback() - except SybaseError as exc: + except Exception as exc: + mapped = create_mapped_exception(exc, logger=logger) msg = f"Failed to rollback Sybase ASE transaction: {exc}" - raise SQLSpecError(msg) from exc + raise type(mapped)(msg) from exc finally: self._explicit_transaction = False self._transaction_active = False @@ -266,9 +270,19 @@ def has_schema(self, schema: str) -> bool: """Return whether the specified user/schema exists in dbo.sysusers.""" result = False exc_handler = self.handle_database_exceptions() - with exc_handler, self.with_cursor(self.connection) as cursor: - cursor.execute("SELECT 1 FROM dbo.sysusers WHERE name = %s", (schema,)) - result = cursor.fetchone() is not None + with self.with_cursor(self.connection) as cursor: + if sys.exc_info()[1] is None: + with exc_handler: + cursor.execute("SELECT 1 FROM dbo.sysusers WHERE name = %s", (schema,)) + result = cursor.fetchone() is not None + else: + + def _fetch() -> None: + nonlocal result + cursor.execute("SELECT 1 FROM dbo.sysusers WHERE name = %s", (schema,)) + result = cursor.fetchone() is not None + + self._run_with_exception_handler(exc_handler, _fetch) self._check_pending_exception(exc_handler) return result @@ -350,8 +364,12 @@ def load_from_arrow( if overwrite: statement = f"DELETE FROM {format_identifier(table)}" exc_handler = self.handle_database_exceptions() - with exc_handler, self.with_cursor(self.connection) as cursor: - cursor.execute(statement) + with self.with_cursor(self.connection) as cursor: + if sys.exc_info()[1] is None: + with exc_handler: + cursor.execute(statement) + else: + self._run_with_exception_handler(exc_handler, cursor.execute, statement) self._check_pending_exception(exc_handler) columns, records = self._arrow_table_to_rows(arrow_table) @@ -367,8 +385,14 @@ def load_from_arrow( ) normalized_records = normalize_execute_many_parameters(prepared_records) exc_handler = self.handle_database_exceptions() - with exc_handler, self.with_cursor(self.connection) as cursor: - cursor.executemany(insert_sql, cast("Sequence[QueryParams]", normalized_records)) + with self.with_cursor(self.connection) as cursor: + if sys.exc_info()[1] is None: + with exc_handler: + cursor.executemany(insert_sql, cast("Sequence[QueryParams]", normalized_records)) + else: + self._run_with_exception_handler( + exc_handler, cursor.executemany, insert_sql, cast("Sequence[QueryParams]", normalized_records) + ) self._check_pending_exception(exc_handler) telemetry_payload = self._ingest_telemetry(arrow_table) diff --git a/sqlspec/adapters/sybase/events/__init__.py b/sqlspec/adapters/sybase/events/__init__.py index 7bf12a0d3..8652828d1 100644 --- a/sqlspec/adapters/sybase/events/__init__.py +++ b/sqlspec/adapters/sybase/events/__init__.py @@ -1,5 +1,5 @@ """Sybase ASE event extension.""" -from sqlspec.adapters.sybase.events.store import SybaseEventQueueStore, SybaseSyncEventQueueStore +from sqlspec.adapters.sybase.events.store import SybaseEventQueueStore -__all__ = ("SybaseEventQueueStore", "SybaseSyncEventQueueStore") +__all__ = ("SybaseEventQueueStore",) diff --git a/sqlspec/adapters/sybase/events/store.py b/sqlspec/adapters/sybase/events/store.py index d5871ac5d..02cd6d245 100644 --- a/sqlspec/adapters/sybase/events/store.py +++ b/sqlspec/adapters/sybase/events/store.py @@ -4,12 +4,12 @@ from sqlspec.adapters.sybase.core import escape_sql_literal, quote_bracket_identifier, split_schema_and_table from sqlspec.extensions.events import BaseEventQueueStore -__all__ = ("SybaseEventQueueStore", "SybaseSyncEventQueueStore") +__all__ = ("SybaseEventQueueStore",) _VARCHAR_TEXT_THRESHOLD = 4000 -class SybaseSyncEventQueueStore(BaseEventQueueStore[SybaseConfig]): +class SybaseEventQueueStore(BaseEventQueueStore[SybaseConfig]): """Sybase ASE DDL hooks for the event queue store.""" __slots__ = () @@ -75,9 +75,6 @@ def _wrap_drop_statement(self, statement: str) -> str: return f"IF OBJECT_ID('{object_name}', 'U') IS NOT NULL EXECUTE('{escaped_statement}');" -SybaseEventQueueStore = SybaseSyncEventQueueStore - - def _object_name(table_name: str) -> str: schema_name, bare_table_name = split_schema_and_table(table_name) return f"{quote_bracket_identifier(schema_name)}.{quote_bracket_identifier(bare_table_name)}" diff --git a/sqlspec/adapters/sybase/litestar/store.py b/sqlspec/adapters/sybase/litestar/store.py index e81a4fe0d..3a6e8ae15 100644 --- a/sqlspec/adapters/sybase/litestar/store.py +++ b/sqlspec/adapters/sybase/litestar/store.py @@ -4,7 +4,13 @@ from typing import TYPE_CHECKING, Any from sqlspec.adapters.sybase._typing import SybaseCursor -from sqlspec.adapters.sybase.core import escape_sql_literal, quote_bracket_identifier, split_schema_and_table +from sqlspec.adapters.sybase.core import ( + create_mapped_exception, + escape_sql_literal, + quote_bracket_identifier, + resolve_rowcount, + split_schema_and_table, +) from sqlspec.extensions.litestar.store import BaseSQLSpecStore from sqlspec.utils.sync_tools import async_ @@ -101,30 +107,33 @@ def _get(self, key: str, renew_for: "int | timedelta | None" = None) -> "bytes | WHERE session_id = %s AND (expires_at IS NULL OR expires_at > GETUTCDATE()) """ - with self._config.provide_connection() as conn: - with SybaseCursor(conn) as cursor: - cursor.execute(sql, (key,)) - row = cursor.fetchone() - - if row is None: - return None - - expires_at = _normalize_utc(_row_value(row, "expires_at", 1)) - if renew_for is not None and expires_at is not None: - new_expires_at = self._calculate_expires_at(renew_for) - if new_expires_at is not None: - with SybaseCursor(conn) as update_cursor: - update_cursor.execute( - f""" - UPDATE {table_ref} - SET expires_at = %s, updated_at = GETUTCDATE() - WHERE session_id = %s - """, - (new_expires_at, key), - ) - conn.commit() - - return _coerce_bytes(_row_value(row, "data", 0)) + try: + with self._config.provide_connection() as conn: + with SybaseCursor(conn) as cursor: + cursor.execute(sql, (key,)) + row = cursor.fetchone() + + if row is None: + return None + + expires_at = _normalize_utc(_row_value(row, "expires_at", 1)) + if renew_for is not None and expires_at is not None: + new_expires_at = self._calculate_expires_at(renew_for) + if new_expires_at is not None: + with SybaseCursor(conn) as update_cursor: + update_cursor.execute( + f""" + UPDATE {table_ref} + SET expires_at = %s, updated_at = GETUTCDATE() + WHERE session_id = %s + """, + (new_expires_at, key), + ) + conn.commit() + + return _coerce_bytes(_row_value(row, "data", 0)) + except Exception as exc: + raise create_mapped_exception(exc) from exc def _set(self, key: str, value: "str | bytes", expires_in: "int | timedelta | None" = None) -> None: data = self._value_to_bytes(value) @@ -139,22 +148,31 @@ def _set(self, key: str, value: "str | bytes", expires_in: "int | timedelta | No INSERT INTO {table_ref} (session_id, data, expires_at) VALUES (%s, %s, %s) """ - with self._config.provide_connection() as conn, SybaseCursor(conn) as cursor: - cursor.execute(sql, (key, data, expires_at, key, key, data, expires_at)) - conn.commit() + try: + with self._config.provide_connection() as conn, SybaseCursor(conn) as cursor: + cursor.execute(sql, (key, data, expires_at, key, key, data, expires_at)) + conn.commit() + except Exception as exc: + raise create_mapped_exception(exc) from exc def _delete(self, key: str) -> None: table_ref = _table_ref(self._table_name) - with self._config.provide_connection() as conn, SybaseCursor(conn) as cursor: - cursor.execute(f"DELETE FROM {table_ref} WHERE session_id = %s", (key,)) - conn.commit() + try: + with self._config.provide_connection() as conn, SybaseCursor(conn) as cursor: + cursor.execute(f"DELETE FROM {table_ref} WHERE session_id = %s", (key,)) + conn.commit() + except Exception as exc: + raise create_mapped_exception(exc) from exc def _delete_all(self) -> None: table_ref = _table_ref(self._table_name) - with self._config.provide_connection() as conn, SybaseCursor(conn) as cursor: - cursor.execute(f"TRUNCATE TABLE {table_ref}") - conn.commit() - self._log_delete_all() + try: + with self._config.provide_connection() as conn, SybaseCursor(conn) as cursor: + cursor.execute(f"TRUNCATE TABLE {table_ref}") + conn.commit() + self._log_delete_all() + except Exception as exc: + raise create_mapped_exception(exc) from exc def _exists(self, key: str) -> bool: table_ref = _table_ref(self._table_name) @@ -164,23 +182,29 @@ def _exists(self, key: str) -> bool: WHERE session_id = %s AND (expires_at IS NULL OR expires_at > GETUTCDATE()) """ - with self._config.provide_connection() as conn, SybaseCursor(conn) as cursor: - cursor.execute(sql, (key,)) - return cursor.fetchone() is not None + try: + with self._config.provide_connection() as conn, SybaseCursor(conn) as cursor: + cursor.execute(sql, (key,)) + return cursor.fetchone() is not None + except Exception as exc: + raise create_mapped_exception(exc) from exc def _expires_in(self, key: str) -> "int | None": table_ref = _table_ref(self._table_name) - with self._config.provide_connection() as conn, SybaseCursor(conn) as cursor: - cursor.execute(f"SELECT expires_at FROM {table_ref} WHERE session_id = %s", (key,)) - row = cursor.fetchone() - - if row is None: - return None - expires_at = _normalize_utc(_row_value(row, "expires_at", 0)) - if expires_at is None: - return None - remaining = expires_at - datetime.now(UTC) - return max(0, int(remaining.total_seconds())) + try: + with self._config.provide_connection() as conn, SybaseCursor(conn) as cursor: + cursor.execute(f"SELECT expires_at FROM {table_ref} WHERE session_id = %s", (key,)) + row = cursor.fetchone() + + if row is None: + return None + expires_at = _normalize_utc(_row_value(row, "expires_at", 0)) + if expires_at is None: + return None + remaining = expires_at - datetime.now(UTC) + return max(0, int(remaining.total_seconds())) + except Exception as exc: + raise create_mapped_exception(exc) from exc def _delete_expired(self) -> int: table_ref = _table_ref(self._table_name) @@ -189,10 +213,13 @@ def _delete_expired(self) -> int: WHERE expires_at IS NOT NULL AND expires_at < GETUTCDATE() """ - with self._config.provide_connection() as conn, SybaseCursor(conn) as cursor: - cursor.execute(sql) - count = int(getattr(cursor, "rowcount", 0) or 0) - conn.commit() + try: + with self._config.provide_connection() as conn, SybaseCursor(conn) as cursor: + cursor.execute(sql) + count = resolve_rowcount(cursor) + conn.commit() + except Exception as exc: + raise create_mapped_exception(exc) from exc if count > 0: self._log_delete_expired(count) return count diff --git a/sqlspec/adapters/sybase/migrations.py b/sqlspec/adapters/sybase/migrations.py index 1e0a4e613..4a24d1bf7 100644 --- a/sqlspec/adapters/sybase/migrations.py +++ b/sqlspec/adapters/sybase/migrations.py @@ -71,14 +71,16 @@ def _tracking_table_ddl_text(self) -> str: def _existing_columns_query(self) -> str: """Return Sybase ASE catalog query text for migration tracking table columns.""" + has_schema = "." in self.version_table schema_name, table_name = split_schema_and_table(self.version_table) + schema_condition = f"u.name = '{escape_sql_literal(schema_name)}'" if has_schema else "u.name = USER_NAME()" return f""" SELECT c.name AS column_name FROM dbo.syscolumns c INNER JOIN dbo.sysobjects o ON o.id = c.id AND o.type = 'U' INNER JOIN dbo.sysusers u ON u.uid = o.uid WHERE o.name = '{escape_sql_literal(table_name)}' - AND u.name = '{escape_sql_literal(schema_name)}' + AND {schema_condition} """ def _add_column_statement_text(self, column_name: str) -> str | None: diff --git a/sqlspec/adapters/sybase/pool.py b/sqlspec/adapters/sybase/pool.py index 7fe52604a..60abb17fd 100644 --- a/sqlspec/adapters/sybase/pool.py +++ b/sqlspec/adapters/sybase/pool.py @@ -6,12 +6,11 @@ import threading import time from contextlib import contextmanager -from importlib import import_module from typing import TYPE_CHECKING, Any, cast from sqlspec.adapters.sybase._typing import PYMSSQL_MODULE, SybaseConnection from sqlspec.adapters.sybase.core import DEFAULT_SYBASE_CONN_PROPERTIES, DEFAULT_SYBASE_TDS_VERSION -from sqlspec.exceptions import ImproperConfigurationError, MissingDependencyError +from sqlspec.exceptions import ImproperConfigurationError from sqlspec.utils.logging import POOL_LOGGER_NAME, get_logger, log_with_context from sqlspec.utils.uuids import uuid4 @@ -22,7 +21,6 @@ logger = get_logger(POOL_LOGGER_NAME) _ADAPTER_NAME = "sybase" -pymssql = PYMSSQL_MODULE class SybaseConnectionPool: @@ -98,10 +96,7 @@ def new_connection(self) -> SybaseConnection: Returns: SybaseConnection: A newly opened, fully configured connection. """ - module = pymssql - if module is None: - module = _load_pymssql_module() - connection = module.connect(**self._connection_parameters) + connection = PYMSSQL_MODULE.connect(**self._connection_parameters) if self._on_connection_create is not None: self._on_connection_create(connection) @@ -265,11 +260,3 @@ def _close_thread_connection(self) -> None: del self._thread_local.created_at if "last_used" in thread_state: del self._thread_local.last_used - - -def _load_pymssql_module() -> Any: - """Import ``pymssql`` lazily or raise ``MissingDependencyError``.""" - try: - return import_module("pymssql") - except ImportError as exc: - raise MissingDependencyError(package="pymssql", install_package="sybase") from exc diff --git a/sqlspec/builder/_ddl.py b/sqlspec/builder/_ddl.py index 95adab5bc..d8ac921ba 100644 --- a/sqlspec/builder/_ddl.py +++ b/sqlspec/builder/_ddl.py @@ -28,7 +28,14 @@ parse_order_expression, ) from sqlspec.core import SQL -from sqlspec.data_dictionary._types import DDLResult, MetadataFidelity, MetadataSource, MetadataSupport, ObjectIdentity +from sqlspec.data_dictionary import ( + DDLResult, + ForeignKeyMetadata, + MetadataFidelity, + MetadataSource, + MetadataSupport, + ObjectIdentity, +) from sqlspec.exceptions import SQLBuilderError from sqlspec.utils.arrow_helpers import sql_type_for_category, sql_type_from_arrow_type from sqlspec.utils.type_guards import has_sqlglot_expression @@ -3064,6 +3071,19 @@ def _normalize_metadata_row(item: Any) -> dict[str, Any]: result.setdefault("schema_name", nested_ident.schema) return result + if isinstance(item, ForeignKeyMetadata): + return { + "table_name": item.table_name, + "column_name": item.column_name, + "referenced_table": item.referenced_table, + "referenced_column": item.referenced_column, + "constraint_name": item.constraint_name, + "schema": item.schema, + "referenced_schema": item.referenced_schema, + "on_delete": item.on_delete, + "on_update": item.on_update, + } + ident = getattr(item, "identity", None) attrs = getattr(item, "attributes", None) if isinstance(ident, ObjectIdentity) and isinstance(attrs, Mapping): diff --git a/sqlspec/config.py b/sqlspec/config.py index 58e477a5d..4286a2fa3 100644 --- a/sqlspec/config.py +++ b/sqlspec/config.py @@ -349,6 +349,18 @@ class FlaskConfig(TypedDict): Set to ``False`` to explicitly disable this behavior. """ + enable_correlation_middleware: NotRequired[bool] + """Enable request correlation ID middleware. Default: True""" + + correlation_header: NotRequired[str] + """HTTP header to read the request correlation ID from when middleware is enabled. Default: ``X-Request-ID``""" + + correlation_headers: NotRequired[tuple[str, ...] | list[str]] + """Additional HTTP headers to read as correlation ID fallbacks.""" + + auto_trace_headers: NotRequired[bool] + """Read standard trace context headers as correlation ID fallbacks. Default: True.""" + class LitestarConfig(TypedDict): """Configuration options for Litestar SQLSpec plugin. @@ -485,6 +497,18 @@ class StarletteConfig(TypedDict): Set to 'fastapi' when using FastAPI. """ + enable_correlation_middleware: NotRequired[bool] + """Enable request correlation ID middleware. Default: True""" + + correlation_header: NotRequired[str] + """HTTP header to read the request correlation ID from when middleware is enabled. Default: ``X-Request-ID``""" + + correlation_headers: NotRequired[tuple[str, ...] | list[str]] + """Additional HTTP headers to read as correlation ID fallbacks.""" + + auto_trace_headers: NotRequired[bool] + """Read standard trace context headers as correlation ID fallbacks. Default: True.""" + class FastAPIConfig(StarletteConfig): """Configuration options for FastAPI SQLSpec extension. @@ -531,7 +555,7 @@ class SanicConfig(TypedDict): """ enable_correlation_middleware: NotRequired[bool] - """Enable request correlation ID middleware. Default: False.""" + """Enable request correlation ID middleware. Default: True.""" correlation_header: NotRequired[str] """HTTP header to read the request correlation ID from when middleware is enabled. Default: ``X-Request-ID``.""" diff --git a/sqlspec/core/compiler.py b/sqlspec/core/compiler.py index fddba3b28..e5ca5bc37 100644 --- a/sqlspec/core/compiler.py +++ b/sqlspec/core/compiler.py @@ -1057,7 +1057,8 @@ def _operation_profile(expression: "exp.Expr | None", operation_type: "Operation returns_rows = True elif isinstance(expr, (exp.Insert, exp.Update, exp.Delete, exp.Merge)): modifies_rows = True - returns_rows = bool(expr.args.get("returning")) + returning = expr.args.get("returning") + returns_rows = bool(returning) and not (isinstance(returning, exp.Returning) and returning.args.get("into")) elif isinstance(expr, exp.TruncateTable): modifies_rows = True elif isinstance(expr, exp.Copy): diff --git a/sqlspec/core/parameters/_declared.py b/sqlspec/core/parameters/_declared.py index 8b6ae92dc..06ee87d9e 100644 --- a/sqlspec/core/parameters/_declared.py +++ b/sqlspec/core/parameters/_declared.py @@ -58,6 +58,10 @@ def __repr__(self) -> str: f"description={self.description!r}, required={self.required!r})" ) + def __reduce__(self) -> "tuple[Callable[..., ParameterDeclaration], tuple[str, str, str | None, bool]]": + """Preserve declaration metadata through constructor-based reconstruction.""" + return (_rebuild_parameter_declaration, (self.name, self.type_str, self.description, self.required)) + def register_param_type(name: str, py_type: ParamTypeMatcher) -> None: """Register or override a declared-type-string matcher. @@ -101,6 +105,13 @@ def matches_param_type(type_str: str, value: object) -> bool: return resolved(value) +def _rebuild_parameter_declaration( + name: str, type_str: str, description: str | None, required: bool +) -> ParameterDeclaration: + """Restore metadata while retaining the constructor's keyword-only required flag.""" + return ParameterDeclaration(name, type_str, description, required=required) + + def _is_json_value(value: object) -> bool: """Return whether a value can be encoded by SQLSpec's JSON serializer.""" if not isinstance(value, _JSON_VALUE_TYPES): diff --git a/sqlspec/core/statement.py b/sqlspec/core/statement.py index a3c928212..100bd90dc 100644 --- a/sqlspec/core/statement.py +++ b/sqlspec/core/statement.py @@ -576,8 +576,9 @@ def returns_rows(self) -> bool: if state.parsed_expression: expr = state.parsed_expression - if isinstance(expr, (exp.Insert, exp.Update, exp.Delete)) and expr.args.get("returning"): - return True + if isinstance(expr, (exp.Insert, exp.Update, exp.Delete)): + returning = expr.args.get("returning") + return bool(returning) and not (isinstance(returning, exp.Returning) and returning.args.get("into")) return False diff --git a/sqlspec/data_dictionary/_types.py b/sqlspec/data_dictionary/_types.py index 57eabc489..1d8396a5a 100644 --- a/sqlspec/data_dictionary/_types.py +++ b/sqlspec/data_dictionary/_types.py @@ -1260,6 +1260,10 @@ class IndexMetadata(TypedDict, total=False): """Whether the index enforces uniqueness.""" is_primary: bool | int """Whether the index backs the primary key constraint.""" + expression: str | None + """Single expression for expression-based indexes.""" + expressions: str | None + """Comma-separated expressions for expression-based indexes.""" @mypyc_attr(allow_interpreted_subclasses=False) diff --git a/sqlspec/data_dictionary/dialects/cockroachdb/__init__.py b/sqlspec/data_dictionary/dialects/cockroachdb/__init__.py index 6dcc444c1..db2bae08f 100644 --- a/sqlspec/data_dictionary/dialects/cockroachdb/__init__.py +++ b/sqlspec/data_dictionary/dialects/cockroachdb/__init__.py @@ -1,9 +1,17 @@ """Dialect configuration for cockroachdb.""" from sqlspec.data_dictionary.dialects.cockroachdb.config import ( + COCKROACH_METADATA_DOMAINS, + COCKROACH_SUPPORTED_DOMAINS, COCKROACHDB_CONFIG, COCKROACHDB_VERSION_PATTERN, resolve_cockroachdb_json_type, ) -__all__ = ("COCKROACHDB_CONFIG", "COCKROACHDB_VERSION_PATTERN", "resolve_cockroachdb_json_type") +__all__ = ( + "COCKROACHDB_CONFIG", + "COCKROACHDB_VERSION_PATTERN", + "COCKROACH_METADATA_DOMAINS", + "COCKROACH_SUPPORTED_DOMAINS", + "resolve_cockroachdb_json_type", +) diff --git a/sqlspec/data_dictionary/dialects/cockroachdb/config.py b/sqlspec/data_dictionary/dialects/cockroachdb/config.py index 5e278943b..9bc5484a3 100644 --- a/sqlspec/data_dictionary/dialects/cockroachdb/config.py +++ b/sqlspec/data_dictionary/dialects/cockroachdb/config.py @@ -10,9 +10,31 @@ "COCKROACHDB_FEATURE_VERSIONS", "COCKROACHDB_TYPE_MAPPINGS", "COCKROACHDB_VERSION_PATTERN", + "COCKROACH_METADATA_DOMAINS", + "COCKROACH_SUPPORTED_DOMAINS", "resolve_cockroachdb_json_type", ) +COCKROACH_METADATA_DOMAINS: tuple[str, ...] = ( + "schemas", + "objects", + "tables", + "columns", + "constraints", + "indexes", + "views", + "sequences", + "comments", + "privileges", + "dependencies", + "ddl", + "crdb_internal", + "system", +) +COCKROACH_SUPPORTED_DOMAINS: frozenset[str] = frozenset( + d for d in COCKROACH_METADATA_DOMAINS if d not in {"crdb_internal", "system"} +) + COCKROACHDB_VERSION_PATTERN = re.compile(r"CockroachDB (?:CCL )?v(\d+)\.(\d+)\.(\d+)") diff --git a/sqlspec/data_dictionary/dialects/cockroachdb/sql/columns.sql b/sqlspec/data_dictionary/dialects/cockroachdb/sql/columns.sql index 258e5fa90..1e1bcc140 100644 --- a/sqlspec/data_dictionary/dialects/cockroachdb/sql/columns.sql +++ b/sqlspec/data_dictionary/dialects/cockroachdb/sql/columns.sql @@ -1,14 +1,15 @@ -- name: by_schema -- dialect: cockroachdb -WITH target_schema AS ( +WITH resolved_table_schema AS ( + SELECT rn.nspname::text AS schema_name + FROM pg_catalog.pg_class rc + JOIN pg_catalog.pg_namespace rn ON rn.oid = rc.relnamespace + WHERE rc.oid = pg_catalog.to_regclass(pg_catalog.quote_ident(:table_name::text)) +), +target_schema AS ( SELECT COALESCE( :schema_name::text, - ( - SELECT rn.nspname::text - FROM pg_catalog.pg_class rc - JOIN pg_catalog.pg_namespace rn ON rn.oid = rc.relnamespace - WHERE rc.oid = pg_catalog.to_regclass(pg_catalog.quote_ident(:table_name::text)) - ), + (SELECT schema_name FROM resolved_table_schema), current_schema() ) AS schema_name ), @@ -81,15 +82,16 @@ ORDER BY c.table_schema, c.table_name, c.ordinal_position; -- name: by_table -- dialect: cockroachdb -WITH target_schema AS ( +WITH resolved_table_schema AS ( + SELECT rn.nspname::text AS schema_name + FROM pg_catalog.pg_class rc + JOIN pg_catalog.pg_namespace rn ON rn.oid = rc.relnamespace + WHERE rc.oid = pg_catalog.to_regclass(pg_catalog.quote_ident(:table_name::text)) +), +target_schema AS ( SELECT COALESCE( :schema_name::text, - ( - SELECT rn.nspname::text - FROM pg_catalog.pg_class rc - JOIN pg_catalog.pg_namespace rn ON rn.oid = rc.relnamespace - WHERE rc.oid = pg_catalog.to_regclass(pg_catalog.quote_ident(:table_name::text)) - ), + (SELECT schema_name FROM resolved_table_schema), current_schema() ) AS schema_name ), diff --git a/sqlspec/data_dictionary/dialects/cockroachdb/sql/constraints.sql b/sqlspec/data_dictionary/dialects/cockroachdb/sql/constraints.sql index 8622bb1bc..4890a6cdf 100644 --- a/sqlspec/data_dictionary/dialects/cockroachdb/sql/constraints.sql +++ b/sqlspec/data_dictionary/dialects/cockroachdb/sql/constraints.sql @@ -1,5 +1,35 @@ -- name: by_schema -- dialect: cockroachdb +WITH con_cols AS ( + SELECT + kcu.constraint_catalog, + kcu.constraint_schema, + kcu.constraint_name, + kcu.table_schema, + kcu.table_name, + array_agg(kcu.column_name::text ORDER BY kcu.ordinal_position) AS columns + FROM information_schema.key_column_usage kcu + WHERE kcu.table_schema = :schema_name + AND (:table_name::text IS NULL OR kcu.table_name = :table_name) + GROUP BY + kcu.constraint_catalog, + kcu.constraint_schema, + kcu.constraint_name, + kcu.table_schema, + kcu.table_name +), +ref_cols AS ( + SELECT + pk_kcu.constraint_catalog, + pk_kcu.constraint_schema, + pk_kcu.constraint_name, + array_agg(pk_kcu.column_name::text ORDER BY pk_kcu.ordinal_position) AS referenced_columns + FROM information_schema.key_column_usage pk_kcu + GROUP BY + pk_kcu.constraint_catalog, + pk_kcu.constraint_schema, + pk_kcu.constraint_name +) SELECT tc.constraint_catalog::text AS constraint_catalog, tc.constraint_schema::text AS constraint_schema, @@ -9,27 +39,17 @@ SELECT tc.constraint_type::text AS constraint_type, tc.is_deferrable::text AS is_deferrable, tc.initially_deferred::text AS initially_deferred, - ARRAY( - SELECT kcu.column_name::text - FROM information_schema.key_column_usage kcu - WHERE kcu.constraint_catalog = tc.constraint_catalog - AND kcu.constraint_schema = tc.constraint_schema - AND kcu.constraint_name = tc.constraint_name - AND kcu.table_schema = tc.table_schema - AND kcu.table_name = tc.table_name - ORDER BY kcu.ordinal_position - )::text[] AS columns, + COALESCE(cc.columns, ARRAY[]::text[]) AS columns, rc.unique_constraint_schema::text AS referenced_schema, pk_tc.table_name::text AS referenced_table, - ARRAY( - SELECT pk_kcu.column_name::text - FROM information_schema.key_column_usage pk_kcu - WHERE pk_kcu.constraint_catalog = rc.unique_constraint_catalog - AND pk_kcu.constraint_schema = rc.unique_constraint_schema - AND pk_kcu.constraint_name = rc.unique_constraint_name - ORDER BY pk_kcu.ordinal_position - )::text[] AS referenced_columns + COALESCE(rc_cols.referenced_columns, ARRAY[]::text[]) AS referenced_columns FROM information_schema.table_constraints tc +LEFT JOIN con_cols cc + ON cc.constraint_catalog = tc.constraint_catalog + AND cc.constraint_schema = tc.constraint_schema + AND cc.constraint_name = tc.constraint_name + AND cc.table_schema = tc.table_schema + AND cc.table_name = tc.table_name LEFT JOIN information_schema.referential_constraints rc ON rc.constraint_catalog = tc.constraint_catalog AND rc.constraint_schema = tc.constraint_schema @@ -38,6 +58,10 @@ LEFT JOIN information_schema.table_constraints pk_tc ON pk_tc.constraint_catalog = rc.unique_constraint_catalog AND pk_tc.constraint_schema = rc.unique_constraint_schema AND pk_tc.constraint_name = rc.unique_constraint_name +LEFT JOIN ref_cols rc_cols + ON rc_cols.constraint_catalog = rc.unique_constraint_catalog + AND rc_cols.constraint_schema = rc.unique_constraint_schema + AND rc_cols.constraint_name = rc.unique_constraint_name WHERE tc.table_schema = :schema_name AND (:table_name::text IS NULL OR tc.table_name = :table_name) ORDER BY tc.table_schema, tc.table_name, tc.constraint_name; diff --git a/sqlspec/data_dictionary/dialects/db2/sql/schemas.sql b/sqlspec/data_dictionary/dialects/db2/sql/schemas.sql index 2367bf22d..05405f292 100644 --- a/sqlspec/data_dictionary/dialects/db2/sql/schemas.sql +++ b/sqlspec/data_dictionary/dialects/db2/sql/schemas.sql @@ -11,7 +11,3 @@ WHERE s.SCHEMANAME NOT LIKE 'SYS%' AND s.SCHEMANAME NOT LIKE 'SQLJ%' AND (:schema_name IS NULL OR s.SCHEMANAME = :schema_name) ORDER BY s.SCHEMANAME; - --- name: current --- dialect: db2 -VALUES CURRENT SCHEMA; diff --git a/sqlspec/data_dictionary/dialects/db2/sql/tables.sql b/sqlspec/data_dictionary/dialects/db2/sql/tables.sql index c47e6e3e1..3ab694ece 100644 --- a/sqlspec/data_dictionary/dialects/db2/sql/tables.sql +++ b/sqlspec/data_dictionary/dialects/db2/sql/tables.sql @@ -75,15 +75,3 @@ WHERE t.TABSCHEMA NOT LIKE 'SYS%' AND t.TABSCHEMA = COALESCE(CAST(:schema_name AS VARCHAR(128)), CURRENT SCHEMA) AND (CAST(:table_name AS VARCHAR(128)) IS NULL OR t.TABNAME = :table_name) ORDER BY t.TABSCHEMA, t.TABNAME; - --- name: names_by_schema --- dialect: db2 -SELECT - RTRIM(t.TABNAME) AS table_name -FROM SYSCAT.TABLES t -WHERE t.TYPE = 'T' - AND t.TABSCHEMA NOT LIKE 'SYS%' - AND t.TABSCHEMA NOT LIKE 'NULLID%' - AND t.TABSCHEMA NOT LIKE 'SQLJ%' - AND t.TABSCHEMA = COALESCE(CAST(:schema_name AS VARCHAR(128)), CURRENT SCHEMA) -ORDER BY t.TABNAME; diff --git a/sqlspec/data_dictionary/dialects/duckdb/sql/tables.sql b/sqlspec/data_dictionary/dialects/duckdb/sql/tables.sql index 63b774447..ee3c20dfb 100644 --- a/sqlspec/data_dictionary/dialects/duckdb/sql/tables.sql +++ b/sqlspec/data_dictionary/dialects/duckdb/sql/tables.sql @@ -55,4 +55,5 @@ FROM information_schema.tables t LEFT JOIN dependency_levels dl ON dl.table_name = t.table_name WHERE t.table_type = 'BASE TABLE' AND t.table_schema = COALESCE(:schema_name, current_schema()) + AND (:table_name IS NULL OR t.table_name = :table_name) ORDER BY level, t.table_name; diff --git a/sqlspec/data_dictionary/dialects/mysql/dictionary.py b/sqlspec/data_dictionary/dialects/mysql/dictionary.py index 207129d2c..b9571e255 100644 --- a/sqlspec/data_dictionary/dialects/mysql/dictionary.py +++ b/sqlspec/data_dictionary/dialects/mysql/dictionary.py @@ -138,7 +138,9 @@ async def get_privileges( """Get privilege metadata from INFORMATION_SCHEMA.""" dialect = await self._get_metadata_dialect(driver) schema_name = self._resolve_metadata_schema(schema) - return await self._select_domain_result(driver, dialect, "privileges", "by_schema", schema_name=schema_name) + return await self._select_domain_result( + driver, dialect, "privileges", "by_schema", schema_name=schema_name, object_name=object_name + ) async def get_system_metadata( self, driver: "AsyncDriverAdapterBase", request: "SystemMetadataRequest | str | None" = None, **kwargs: Any @@ -250,9 +252,18 @@ async def _get_native_object_ddl( _ = include_dependencies, redact if object_type.strip().upper() not in MYSQL_SHOW_CREATE_OBJECT_TYPES: return None - raw_version = await driver.select_value_or_none(self.get_query_text("version", "current")) - sql_mode = await driver.select_value_or_none("SELECT @@sql_mode") - sql_quote_show_create = await driver.select_value_or_none("SELECT @@sql_quote_show_create") + try: + row_ctx = await driver.select_one( + "SELECT VERSION() AS raw_version, @@sql_mode AS sql_mode, @@sql_quote_show_create AS sql_quote_show_create", + schema_type=dict, + ) + raw_version = row_ctx.get("raw_version") + sql_mode = row_ctx.get("sql_mode") + sql_quote_show_create = row_ctx.get("sql_quote_show_create") + except Exception: + raw_version = await driver.select_value_or_none(self.get_query_text("version", "current")) + sql_mode = await driver.select_value_or_none("SELECT @@sql_mode") + sql_quote_show_create = await driver.select_value_or_none("SELECT @@sql_quote_show_create") statement = build_mysql_show_create_statement(object_name, schema, object_type) try: row = await driver.select_one(statement, schema_type=dict) @@ -317,15 +328,18 @@ async def _select_domain_result( *, schema_name: "str | None", table_name: "str | None" = None, + object_name: "str | None" = None, ) -> "MetadataResult": query = get_data_dictionary_loader().get_domain_query(dialect, domain, query_name) if not query.is_supported or query.sql is None: return MetadataResult.unsupported(domain) capability = build_mysql_metadata_capability_profile(dialect, type(self).__name__).get(domain) - if table_name is None: - rows = await driver.select(query.sql, schema_name=schema_name) - else: - rows = await driver.select(query.sql, schema_name=schema_name, table_name=table_name) + params: dict[str, Any] = {"schema_name": schema_name} + if table_name is not None: + params["table_name"] = table_name + if domain == "privileges": + params["object_name"] = object_name + rows = await driver.select(query.sql, **params) return MetadataResult(domain, capability=capability, items=tuple(rows)) @@ -420,7 +434,9 @@ def get_privileges( """Get privilege metadata from INFORMATION_SCHEMA.""" dialect = self._get_metadata_dialect(driver) schema_name = self._resolve_metadata_schema(schema) - return self._select_domain_result(driver, dialect, "privileges", "by_schema", schema_name=schema_name) + return self._select_domain_result( + driver, dialect, "privileges", "by_schema", schema_name=schema_name, object_name=object_name + ) def get_system_metadata( self, driver: "SyncDriverAdapterBase", request: "SystemMetadataRequest | str | None" = None, **kwargs: Any @@ -532,9 +548,18 @@ def _get_native_object_ddl( _ = include_dependencies, redact if object_type.strip().upper() not in MYSQL_SHOW_CREATE_OBJECT_TYPES: return None - raw_version = driver.select_value_or_none(self.get_query_text("version", "current")) - sql_mode = driver.select_value_or_none("SELECT @@sql_mode") - sql_quote_show_create = driver.select_value_or_none("SELECT @@sql_quote_show_create") + try: + row_ctx = driver.select_one( + "SELECT VERSION() AS raw_version, @@sql_mode AS sql_mode, @@sql_quote_show_create AS sql_quote_show_create", + schema_type=dict, + ) + raw_version = row_ctx.get("raw_version") + sql_mode = row_ctx.get("sql_mode") + sql_quote_show_create = row_ctx.get("sql_quote_show_create") + except Exception: + raw_version = driver.select_value_or_none(self.get_query_text("version", "current")) + sql_mode = driver.select_value_or_none("SELECT @@sql_mode") + sql_quote_show_create = driver.select_value_or_none("SELECT @@sql_quote_show_create") statement = build_mysql_show_create_statement(object_name, schema, object_type) try: row = driver.select_one(statement, schema_type=dict) @@ -599,15 +624,18 @@ def _select_domain_result( *, schema_name: "str | None", table_name: "str | None" = None, + object_name: "str | None" = None, ) -> "MetadataResult": query = get_data_dictionary_loader().get_domain_query(dialect, domain, query_name) if not query.is_supported or query.sql is None: return MetadataResult.unsupported(domain) capability = build_mysql_metadata_capability_profile(dialect, type(self).__name__).get(domain) - if table_name is None: - rows = driver.select(query.sql, schema_name=schema_name) - else: - rows = driver.select(query.sql, schema_name=schema_name, table_name=table_name) + params: dict[str, Any] = {"schema_name": schema_name} + if table_name is not None: + params["table_name"] = table_name + if domain == "privileges": + params["object_name"] = object_name + rows = driver.select(query.sql, **params) return MetadataResult(domain, capability=capability, items=tuple(rows)) diff --git a/sqlspec/data_dictionary/dialects/mysql/sql/privileges.sql b/sqlspec/data_dictionary/dialects/mysql/sql/privileges.sql index c7900fa34..1002dab00 100644 --- a/sqlspec/data_dictionary/dialects/mysql/sql/privileges.sql +++ b/sqlspec/data_dictionary/dialects/mysql/sql/privileges.sql @@ -10,6 +10,7 @@ SELECT NULL AS column_name FROM information_schema.table_privileges WHERE table_schema = COALESCE(:schema_name, DATABASE()) + AND (:object_name IS NULL OR table_name = :object_name) UNION ALL SELECT grantee, @@ -21,4 +22,5 @@ SELECT column_name FROM information_schema.column_privileges WHERE table_schema = COALESCE(:schema_name, DATABASE()) + AND (:object_name IS NULL OR table_name = :object_name) ORDER BY schema_name, table_name, column_name, grantee, privilege_type; diff --git a/sqlspec/data_dictionary/dialects/oracle/config.py b/sqlspec/data_dictionary/dialects/oracle/config.py index f1dc95a36..ec2ba12a9 100644 --- a/sqlspec/data_dictionary/dialects/oracle/config.py +++ b/sqlspec/data_dictionary/dialects/oracle/config.py @@ -168,7 +168,7 @@ def extract_oracle_version_value(row: object) -> "str | None": Version string when one can be found, otherwise None. """ if isinstance(row, Mapping): - for key in ("version", "VERSION", "Version"): + for key in ("version_full", "VERSION_FULL", "Version_Full", "version", "VERSION", "Version"): value = row.get(key) if value: return str(value) diff --git a/sqlspec/data_dictionary/dialects/oracle/sql/indexes.sql b/sqlspec/data_dictionary/dialects/oracle/sql/indexes.sql index e17d28a92..c18eb142f 100644 --- a/sqlspec/data_dictionary/dialects/oracle/sql/indexes.sql +++ b/sqlspec/data_dictionary/dialects/oracle/sql/indexes.sql @@ -20,16 +20,11 @@ SELECT i.avg_leaf_blocks_per_key, i.avg_data_blocks_per_key, i.clustering_factor, - LISTAGG(ic.column_name, ',') WITHIN GROUP (ORDER BY ic.column_position) AS columns, - LISTAGG(ie.column_expression, ',') WITHIN GROUP (ORDER BY ie.column_position) AS expressions + LISTAGG(ic.column_name, ',') WITHIN GROUP (ORDER BY ic.column_position) AS columns FROM all_indexes i LEFT JOIN all_ind_columns ic ON ic.index_owner = i.owner AND ic.index_name = i.index_name -LEFT JOIN all_ind_expressions ie - ON ie.index_owner = i.owner - AND ie.index_name = i.index_name - AND ie.column_position = ic.column_position WHERE i.owner = COALESCE(:schema_name, USER) AND (:table_name IS NULL OR i.table_name = :table_name) GROUP BY @@ -57,6 +52,7 @@ ORDER BY i.owner, i.table_name, i.index_name; -- name: by_table -- dialect: oracle SELECT + i.owner AS schema_name, i.index_name, i.table_name AS table_name, CASE WHEN i.uniqueness = 'UNIQUE' THEN 1 ELSE 0 END AS is_unique, @@ -75,4 +71,17 @@ LEFT JOIN all_constraints c AND c.constraint_type = 'P' WHERE i.table_name = :table_name AND i.owner = COALESCE(:schema_name, USER) -GROUP BY i.index_name, i.table_name, i.uniqueness; +GROUP BY i.owner, i.index_name, i.table_name, i.uniqueness; + +-- name: expressions +-- dialect: oracle +SELECT + ie.index_owner AS schema_name, + ie.table_name, + ie.index_name, + ie.column_position, + ie.column_expression +FROM all_ind_expressions ie +WHERE ie.index_owner = COALESCE(:schema_name, USER) + AND (:table_name IS NULL OR ie.table_name = :table_name) +ORDER BY ie.index_owner, ie.table_name, ie.index_name, ie.column_position; diff --git a/sqlspec/data_dictionary/dialects/oracle/sql/version.sql b/sqlspec/data_dictionary/dialects/oracle/sql/version.sql index bbf3ec91c..ee88f5ca7 100644 --- a/sqlspec/data_dictionary/dialects/oracle/sql/version.sql +++ b/sqlspec/data_dictionary/dialects/oracle/sql/version.sql @@ -1,12 +1,10 @@ -- name: current -- dialect: oracle SELECT - product AS "product", - version AS "version", - status AS "status" -FROM product_component_version -WHERE product LIKE 'Oracle%' -ORDER BY TO_NUMBER(REGEXP_SUBSTR(version, '^[0-9]+')) DESC, version DESC + p.* +FROM product_component_version p +WHERE p.product LIKE 'Oracle%' +ORDER BY TO_NUMBER(REGEXP_SUBSTR(p.version, '^[0-9]+')) DESC, p.version DESC FETCH FIRST 1 ROWS ONLY; -- name: autonomous_service diff --git a/sqlspec/data_dictionary/dialects/postgres/__init__.py b/sqlspec/data_dictionary/dialects/postgres/__init__.py index d15fc9558..ab166164a 100644 --- a/sqlspec/data_dictionary/dialects/postgres/__init__.py +++ b/sqlspec/data_dictionary/dialects/postgres/__init__.py @@ -2,8 +2,18 @@ from sqlspec.data_dictionary.dialects.postgres.config import ( POSTGRES_CONFIG, + POSTGRES_METADATA_DOMAINS, + POSTGRES_SUPPORTED_DOMAINS, + POSTGRES_SYSTEM_METADATA_QUERIES, POSTGRES_VERSION_PATTERN, resolve_postgres_json_type, ) -__all__ = ("POSTGRES_CONFIG", "POSTGRES_VERSION_PATTERN", "resolve_postgres_json_type") +__all__ = ( + "POSTGRES_CONFIG", + "POSTGRES_METADATA_DOMAINS", + "POSTGRES_SUPPORTED_DOMAINS", + "POSTGRES_SYSTEM_METADATA_QUERIES", + "POSTGRES_VERSION_PATTERN", + "resolve_postgres_json_type", +) diff --git a/sqlspec/data_dictionary/dialects/postgres/config.py b/sqlspec/data_dictionary/dialects/postgres/config.py index 405844603..4bb9353b6 100644 --- a/sqlspec/data_dictionary/dialects/postgres/config.py +++ b/sqlspec/data_dictionary/dialects/postgres/config.py @@ -8,11 +8,44 @@ "POSTGRES_CONFIG", "POSTGRES_FEATURE_FLAGS", "POSTGRES_FEATURE_VERSIONS", + "POSTGRES_METADATA_DOMAINS", + "POSTGRES_SUPPORTED_DOMAINS", + "POSTGRES_SYSTEM_METADATA_QUERIES", "POSTGRES_TYPE_MAPPINGS", "POSTGRES_VERSION_PATTERN", "resolve_postgres_json_type", ) +POSTGRES_METADATA_DOMAINS: tuple[str, ...] = ( + "schemas", + "objects", + "tables", + "columns", + "constraints", + "indexes", + "views", + "materialized_views", + "sequences", + "types", + "routines", + "triggers", + "comments", + "privileges", + "dependencies", + "extensions", + "partitions", + "ddl", + "system", +) +POSTGRES_SUPPORTED_DOMAINS: frozenset[str] = frozenset(d for d in POSTGRES_METADATA_DOMAINS if d != "system") +POSTGRES_SYSTEM_METADATA_QUERIES: dict[str, str] = { + "settings": "settings", + "statement_history": "pg_stat_statements", + "pg_stat_statements": "pg_stat_statements", + "table_statistics": "table_stats", + "table_stats": "table_stats", +} + POSTGRES_VERSION_PATTERN = re.compile(r"PostgreSQL (\d+)\.(\d+)(?:\.(\d+))?") diff --git a/sqlspec/data_dictionary/dialects/postgres/sql/constraints.sql b/sqlspec/data_dictionary/dialects/postgres/sql/constraints.sql index 59351a740..6b41bf929 100644 --- a/sqlspec/data_dictionary/dialects/postgres/sql/constraints.sql +++ b/sqlspec/data_dictionary/dialects/postgres/sql/constraints.sql @@ -1,5 +1,23 @@ -- name: by_schema -- dialect: postgres +WITH con_cols AS ( + SELECT + con.oid AS con_oid, + array_agg(att.attname::text ORDER BY keys.ord) AS columns + FROM pg_catalog.pg_constraint con + CROSS JOIN LATERAL unnest(con.conkey) WITH ORDINALITY AS keys(attnum, ord) + JOIN pg_catalog.pg_attribute att ON att.attrelid = con.conrelid AND att.attnum = keys.attnum + GROUP BY con.oid +), +ref_cols AS ( + SELECT + con.oid AS con_oid, + array_agg(att.attname::text ORDER BY keys.ord) AS referenced_columns + FROM pg_catalog.pg_constraint con + CROSS JOIN LATERAL unnest(con.confkey) WITH ORDINALITY AS keys(attnum, ord) + JOIN pg_catalog.pg_attribute att ON att.attrelid = con.confrelid AND att.attnum = keys.attnum + GROUP BY con.oid +) SELECT current_database()::text AS constraint_catalog, n.nspname::text AS constraint_schema, @@ -15,20 +33,10 @@ SELECT ELSE con.contype::text END AS constraint_type, pg_catalog.pg_get_constraintdef(con.oid, true)::text AS definition, - ARRAY( - SELECT att.attname::text - FROM unnest(con.conkey) WITH ORDINALITY AS keys(attnum, ord) - JOIN pg_catalog.pg_attribute att ON att.attrelid = con.conrelid AND att.attnum = keys.attnum - ORDER BY keys.ord - )::text[] AS columns, + COALESCE(cc.columns, ARRAY[]::text[]) AS columns, ref_ns.nspname::text AS referenced_schema, ref.relname::text AS referenced_table, - ARRAY( - SELECT att.attname::text - FROM unnest(con.confkey) WITH ORDINALITY AS keys(attnum, ord) - JOIN pg_catalog.pg_attribute att ON att.attrelid = con.confrelid AND att.attnum = keys.attnum - ORDER BY keys.ord - )::text[] AS referenced_columns, + COALESCE(rc.referenced_columns, ARRAY[]::text[]) AS referenced_columns, con.condeferrable AS is_deferrable, con.condeferred AS initially_deferred, dep.refobjid IS NOT NULL AS is_extension_owned @@ -38,6 +46,8 @@ LEFT JOIN pg_catalog.pg_class c ON c.oid = con.conrelid LEFT JOIN pg_catalog.pg_class ref ON ref.oid = con.confrelid LEFT JOIN pg_catalog.pg_namespace ref_ns ON ref_ns.oid = ref.relnamespace LEFT JOIN pg_catalog.pg_depend dep ON dep.objid = con.oid AND dep.deptype = 'e' +LEFT JOIN con_cols cc ON cc.con_oid = con.oid +LEFT JOIN ref_cols rc ON rc.con_oid = con.oid WHERE n.nspname = :schema_name AND (:table_name::text IS NULL OR c.relname = :table_name) ORDER BY n.nspname, c.relname, con.conname; diff --git a/sqlspec/data_dictionary/dialects/postgres/sql/indexes.sql b/sqlspec/data_dictionary/dialects/postgres/sql/indexes.sql index 41006915e..fbd122100 100644 --- a/sqlspec/data_dictionary/dialects/postgres/sql/indexes.sql +++ b/sqlspec/data_dictionary/dialects/postgres/sql/indexes.sql @@ -1,5 +1,17 @@ -- name: by_schema -- dialect: postgres +WITH target_schema AS ( + SELECT COALESCE( + :schema_name::text, + ( + SELECT rn.nspname::text + FROM pg_catalog.pg_class rc + JOIN pg_catalog.pg_namespace rn ON rn.oid = rc.relnamespace + WHERE rc.oid = pg_catalog.to_regclass(pg_catalog.quote_ident(:table_name::text)) + ), + current_schema() + ) AS schema_name +) SELECT tn.nspname::text AS schema_name, tbl.relname::text AS table_name, @@ -35,11 +47,11 @@ SELECT FROM pg_catalog.pg_index ix JOIN pg_catalog.pg_class tbl ON tbl.oid = ix.indrelid JOIN pg_catalog.pg_namespace tn ON tn.oid = tbl.relnamespace +JOIN target_schema ts ON ts.schema_name = tn.nspname JOIN pg_catalog.pg_class idx ON idx.oid = ix.indexrelid JOIN pg_catalog.pg_am am ON am.oid = idx.relam LEFT JOIN pg_catalog.pg_depend dep ON dep.objid = idx.oid AND dep.deptype = 'e' -WHERE tn.nspname = :schema_name - AND (:table_name::text IS NULL OR tbl.relname = :table_name) +WHERE (:table_name::text IS NULL OR tbl.relname = :table_name) ORDER BY tn.nspname, tbl.relname, idx.relname; -- name: by_table diff --git a/sqlspec/data_dictionary/dialects/postgres/sql/system.sql b/sqlspec/data_dictionary/dialects/postgres/sql/system.sql index b6b6aae58..0035d9ebb 100644 --- a/sqlspec/data_dictionary/dialects/postgres/sql/system.sql +++ b/sqlspec/data_dictionary/dialects/postgres/sql/system.sql @@ -9,7 +9,7 @@ SELECT total_exec_time, mean_exec_time, rows -FROM pg_catalog.pg_stat_statements +FROM pg_stat_statements ORDER BY total_exec_time DESC LIMIT :limit; diff --git a/sqlspec/data_dictionary/dialects/spanner/config.py b/sqlspec/data_dictionary/dialects/spanner/config.py index be14fba4a..c6e42b9fc 100644 --- a/sqlspec/data_dictionary/dialects/spanner/config.py +++ b/sqlspec/data_dictionary/dialects/spanner/config.py @@ -52,6 +52,7 @@ feature_flags=SPANNER_FEATURE_FLAGS, type_mappings=SPANNER_TYPE_MAPPINGS, version_pattern=SPANNER_VERSION_PATTERN, + default_schema="", ) register_dialect(SPANNER_CONFIG) diff --git a/sqlspec/data_dictionary/dialects/sqlanywhere/config.py b/sqlspec/data_dictionary/dialects/sqlanywhere/config.py index 51246d50a..4652ddc32 100644 --- a/sqlspec/data_dictionary/dialects/sqlanywhere/config.py +++ b/sqlspec/data_dictionary/dialects/sqlanywhere/config.py @@ -57,7 +57,7 @@ "supports_merge": True, "supports_for_update": True, "supports_skip_locked": True, - "supports_transactional_ddl": True, + "supports_transactional_ddl": False, "supports_transactions": True, "supports_prepared_statements": True, "supports_in_memory": False, diff --git a/sqlspec/data_dictionary/dialects/sqlite/sql/columns.sql b/sqlspec/data_dictionary/dialects/sqlite/sql/columns.sql index dcae8f703..510bac0d8 100644 --- a/sqlspec/data_dictionary/dialects/sqlite/sql/columns.sql +++ b/sqlspec/data_dictionary/dialects/sqlite/sql/columns.sql @@ -33,7 +33,7 @@ ORDER BY tl.name, ti.cid; -- name: by_table -- dialect: sqlite SELECT - :schema_name AS schema_name, + COALESCE(:schema_name, 'main') AS schema_name, :table_name AS table_name, ti.name AS column_name, ti.type AS data_type, @@ -54,6 +54,6 @@ SELECT END AS extra, ti.hidden IN (2, 3) AS is_generated, CASE WHEN ti.hidden = 1 THEN 1 ELSE 0 END AS is_hidden -FROM pragma_table_xinfo(:table_name, :schema_name) AS ti +FROM pragma_table_xinfo(:table_name, COALESCE(:schema_name, 'main')) AS ti LEFT JOIN {schema_prefix}sqlite_schema AS sm ON sm.type = 'table' AND sm.name = :table_name ORDER BY ti.cid; diff --git a/sqlspec/data_dictionary/dialects/sqlite/sql/indexes.sql b/sqlspec/data_dictionary/dialects/sqlite/sql/indexes.sql index ccc3f9114..f777f4dd8 100644 --- a/sqlspec/data_dictionary/dialects/sqlite/sql/indexes.sql +++ b/sqlspec/data_dictionary/dialects/sqlite/sql/indexes.sql @@ -30,7 +30,7 @@ ORDER BY tl.name, il.seq, ix.seqno; -- name: by_table -- dialect: sqlite SELECT - :schema_name AS schema_name, + COALESCE(:schema_name, 'main') AS schema_name, :table_name AS table_name, il.name AS index_name, il."unique" AS is_unique, @@ -44,8 +44,8 @@ SELECT ix.coll AS collation, ix.key AS is_key_column, sm.sql AS index_sql -FROM pragma_index_list(:table_name, :schema_name) AS il -JOIN pragma_index_xinfo(il.name, :schema_name) AS ix +FROM pragma_index_list(:table_name, COALESCE(:schema_name, 'main')) AS il +JOIN pragma_index_xinfo(il.name, COALESCE(:schema_name, 'main')) AS ix LEFT JOIN {schema_prefix}sqlite_schema AS sm ON sm.type = 'index' AND sm.name = il.name @@ -55,7 +55,7 @@ ORDER BY il.seq, ix.seqno; -- name: columns_by_index -- dialect: sqlite SELECT - :schema_name AS schema_name, + COALESCE(:schema_name, 'main') AS schema_name, :index_name AS index_name, ix.seqno, ix.cid AS column_id, @@ -63,7 +63,7 @@ SELECT ix.desc AS is_descending, ix.coll AS collation, ix.key AS is_key_column -FROM pragma_index_xinfo(:index_name, :schema_name) AS ix +FROM pragma_index_xinfo(:index_name, COALESCE(:schema_name, 'main')) AS ix ORDER BY ix.seqno; -- name: index_columns_by_index diff --git a/sqlspec/driver/_async.py b/sqlspec/driver/_async.py index 19235d8b9..b581a7a5b 100644 --- a/sqlspec/driver/_async.py +++ b/sqlspec/driver/_async.py @@ -340,8 +340,7 @@ async def dispatch_execute_script(self, cursor: Any, statement: "SQL") -> Execut self._release_pooled_statement(single_stmt) if has_observers: stmt_duration = perf_counter() - stmt_started - stmt_profile: dict[str, Any] | None = self._pending_statement_profile - self._pending_statement_profile = None + stmt_profile = self._consume_statement_profile() self._emit_script_statement_event( sql=stmt, parameters=prepared_parameters, @@ -1659,8 +1658,7 @@ async def _execute_with_observability_async( runtime.span_manager.end_span(span, error=exc) await runtime.emit_error_async(exc, **query_context) raise - profile: dict[str, Any] | None = self._pending_statement_profile - self._pending_statement_profile = None + profile = self._consume_statement_profile() runtime.span_manager.end_span(span) duration = perf_counter() - started await runtime.emit_query_complete_async(**{**query_context, "rows_affected": result.rows_affected}) @@ -2265,7 +2263,7 @@ async def get_object_ddl( Returns: DDL result for the object. """ - resolved_name, resolved_schema = self.resolve_ddl_object_name(driver, object_name, schema) + resolved_name, resolved_schema = type(self).resolve_ddl_object_name(self, driver, object_name, schema) return await self._get_exact_object_ddl( driver, resolved_name, @@ -2379,11 +2377,13 @@ async def get_table_profile( Query failures (for example missing privileges) propagate to the caller. """ - resolved_table, resolved_schema = self.resolve_ddl_object_name(driver, table, schema) + resolved_table, resolved_schema = type(self).resolve_ddl_object_name(self, driver, table, schema) columns = await self.get_columns( driver, - table=self.exact_catalog_identifier(driver, resolved_table), - schema=self.exact_catalog_identifier(driver, resolved_schema) if resolved_schema is not None else None, + table=type(self).exact_catalog_identifier(self, driver, resolved_table), + schema=type(self).exact_catalog_identifier(self, driver, resolved_schema) + if resolved_schema is not None + else None, ) if not columns: return [] @@ -2411,7 +2411,7 @@ async def infer_table_ddl_from_data( profiles = await self.get_table_profile( driver, table, schema=schema, sample_size=sample_size, target_dialect=target_dialect ) - resolved_table, resolved_schema = self.resolve_ddl_object_name(driver, table, schema) + resolved_table, resolved_schema = type(self).resolve_ddl_object_name(self, driver, table, schema) if not profiles: return self._unsupported_ddl(target_table or resolved_table, schema=resolved_schema, object_type="table") return self._compile_profiled_table_ddl( @@ -2604,7 +2604,7 @@ async def _replace_unfaithful_native_tables( if not replace_names: return items catalog = await self._load_table_catalog(driver, schema, replace_names, warnings) - source_dialect = self.resolve_ddl_dialect(driver) + source_dialect = type(self).resolve_ddl_dialect(self, driver) resolved: list[DDLResult] = [] replaced_tables: set[tuple[str, str]] = set() for item in items: @@ -2653,7 +2653,7 @@ async def _collect_synthesized_schema_ddl( they were already loaded for every table, so callers can reuse them for ordering. """ resolved_default_schema = default_schema if default_schema is not None else schema - source_dialect = self.resolve_ddl_dialect(driver) + source_dialect = type(self).resolve_ddl_dialect(self, driver) excluded_names = self._normalize_excluded_table_names(exclude_tables) results: list[DDLResult] = [] loaded_foreign_keys: list[Any] | None = None @@ -2837,7 +2837,7 @@ async def _load_table_catalog( async def _load_table_rows(self, driver: Any, domain: str, table_name: str, schema: "str | None") -> "list[Any]": """Load one catalog domain for one table using its exact catalog name.""" - lookup_table = self.exact_catalog_identifier(driver, table_name) + lookup_table = type(self).exact_catalog_identifier(self, driver, table_name) if domain == "columns": return list(await self.get_columns(driver, table=lookup_table, schema=schema)) if domain == "foreign_keys": @@ -2859,8 +2859,8 @@ async def _synthesize_table_ddl( warnings: tuple[str, ...] = (), ) -> "DDLResult": """Synthesize table DDL from exact catalog names using columns, constraints, foreign keys, and indexes.""" - lookup_table = self.exact_catalog_identifier(driver, table_name) - lookup_schema = self.exact_catalog_identifier(driver, schema) if schema is not None else None + lookup_table = type(self).exact_catalog_identifier(self, driver, table_name) + lookup_schema = type(self).exact_catalog_identifier(self, driver, schema) if schema is not None else None columns = await self.get_columns(driver, table=lookup_table, schema=lookup_schema) if not columns: return self._unsupported_ddl(table_name, schema=schema, object_type="table", warnings=warnings) @@ -2879,7 +2879,7 @@ async def _synthesize_table_ddl( foreign_keys=foreign_keys, indexes=indexes, target_dialect=target_dialect, - source_dialect=self.resolve_ddl_dialect(driver), + source_dialect=type(self).resolve_ddl_dialect(self, driver), warnings=warnings, ) diff --git a/sqlspec/driver/_common.py b/sqlspec/driver/_common.py index 647dce1b0..3a6b89717 100644 --- a/sqlspec/driver/_common.py +++ b/sqlspec/driver/_common.py @@ -128,7 +128,7 @@ from sqlspec.core import ArrowResult, FilterTypeT, OperationType, StatementFilter from sqlspec.core.parameters._types import ConvertedParameters from sqlspec.core.stack import StatementStack - from sqlspec.data_dictionary._types import DialectConfig + from sqlspec.data_dictionary import DialectConfig from sqlspec.storage import ( AsyncStoragePipeline, StorageBridgeJob, @@ -768,9 +768,11 @@ def _resolve_schema_ddl_input(self, driver: Any, schema: str | None) -> tuple[st return (None, None) if _ddl_dialect_or_none(type(self).dialect) is not None and isinstance(self, DataDictionaryDialectMixin): with contextlib.suppress(ValueError): - _, exact_schema = self.resolve_ddl_object_name(driver, "_", schema) + _, exact_schema = type(self).resolve_ddl_object_name(self, driver, "_", schema) getter_schema = ( - self.exact_catalog_identifier(driver, exact_schema) if exact_schema is not None else None + type(self).exact_catalog_identifier(self, driver, exact_schema) + if exact_schema is not None + else None ) return (exact_schema, getter_schema) stripped = schema.strip() @@ -1819,6 +1821,12 @@ def attach_observability(self, runtime: "ObservabilityRuntime") -> None: self._observability = runtime self._refresh_statement_cache_state() + def _consume_statement_profile(self) -> "dict[str, Any] | None": + """Take the profile produced by an execution callback and clear pending ownership.""" + profile = self._pending_statement_profile + self._pending_statement_profile = None + return profile + def storage_capabilities(self) -> "StorageCapabilities": """Return cached storage capabilities for the active driver. diff --git a/sqlspec/driver/_sync.py b/sqlspec/driver/_sync.py index 8f49a40d3..7643c73fe 100644 --- a/sqlspec/driver/_sync.py +++ b/sqlspec/driver/_sync.py @@ -304,8 +304,7 @@ def dispatch_execute_script(self, cursor: Any, statement: "SQL") -> ExecutionRes self._release_pooled_statement(single_stmt) if has_observers: stmt_duration = perf_counter() - stmt_started - stmt_profile: dict[str, Any] | None = self._pending_statement_profile - self._pending_statement_profile = None + stmt_profile = self._consume_statement_profile() self._emit_script_statement_event( sql=stmt, parameters=prepared_parameters, @@ -1640,8 +1639,7 @@ def _execute_with_observability_sync( runtime.span_manager.end_span(span, error=exc) runtime.emit_error_sync(exc, **query_context) raise - profile: dict[str, Any] | None = self._pending_statement_profile - self._pending_statement_profile = None + profile = self._consume_statement_profile() runtime.span_manager.end_span(span) duration = perf_counter() - started runtime.emit_query_complete_sync(**{**query_context, "rows_affected": result.rows_affected}) @@ -2169,7 +2167,7 @@ def get_object_ddl( Returns: DDL result for the object. """ - resolved_name, resolved_schema = self.resolve_ddl_object_name(driver, object_name, schema) + resolved_name, resolved_schema = type(self).resolve_ddl_object_name(self, driver, object_name, schema) return self._get_exact_object_ddl( driver, resolved_name, @@ -2283,11 +2281,13 @@ def get_table_profile( Query failures (for example missing privileges) propagate to the caller. """ - resolved_table, resolved_schema = self.resolve_ddl_object_name(driver, table, schema) + resolved_table, resolved_schema = type(self).resolve_ddl_object_name(self, driver, table, schema) columns = self.get_columns( driver, - table=self.exact_catalog_identifier(driver, resolved_table), - schema=self.exact_catalog_identifier(driver, resolved_schema) if resolved_schema is not None else None, + table=type(self).exact_catalog_identifier(self, driver, resolved_table), + schema=type(self).exact_catalog_identifier(self, driver, resolved_schema) + if resolved_schema is not None + else None, ) if not columns: return [] @@ -2313,7 +2313,7 @@ def infer_table_ddl_from_data( profiles = self.get_table_profile( driver, table, schema=schema, sample_size=sample_size, target_dialect=target_dialect ) - resolved_table, resolved_schema = self.resolve_ddl_object_name(driver, table, schema) + resolved_table, resolved_schema = type(self).resolve_ddl_object_name(self, driver, table, schema) if not profiles: return self._unsupported_ddl(target_table or resolved_table, schema=resolved_schema, object_type="table") return self._compile_profiled_table_ddl( @@ -2506,7 +2506,7 @@ def _replace_unfaithful_native_tables( if not replace_names: return items catalog = self._load_table_catalog(driver, schema, replace_names, warnings) - source_dialect = self.resolve_ddl_dialect(driver) + source_dialect = type(self).resolve_ddl_dialect(self, driver) resolved: list[DDLResult] = [] replaced_tables: set[tuple[str, str]] = set() for item in items: @@ -2555,7 +2555,7 @@ def _collect_synthesized_schema_ddl( they were already loaded for every table, so callers can reuse them for ordering. """ resolved_default_schema = default_schema if default_schema is not None else schema - source_dialect = self.resolve_ddl_dialect(driver) + source_dialect = type(self).resolve_ddl_dialect(self, driver) excluded_names = self._normalize_excluded_table_names(exclude_tables) results: list[DDLResult] = [] loaded_foreign_keys: list[Any] | None = None @@ -2739,7 +2739,7 @@ def _load_table_catalog( def _load_table_rows(self, driver: Any, domain: str, table_name: str, schema: "str | None") -> "list[Any]": """Load one catalog domain for one table using its exact catalog name.""" - lookup_table = self.exact_catalog_identifier(driver, table_name) + lookup_table = type(self).exact_catalog_identifier(self, driver, table_name) if domain == "columns": return list(self.get_columns(driver, table=lookup_table, schema=schema)) if domain == "foreign_keys": @@ -2761,8 +2761,8 @@ def _synthesize_table_ddl( warnings: tuple[str, ...] = (), ) -> "DDLResult": """Synthesize table DDL from exact catalog names using columns, constraints, foreign keys, and indexes.""" - lookup_table = self.exact_catalog_identifier(driver, table_name) - lookup_schema = self.exact_catalog_identifier(driver, schema) if schema is not None else None + lookup_table = type(self).exact_catalog_identifier(self, driver, table_name) + lookup_schema = type(self).exact_catalog_identifier(self, driver, schema) if schema is not None else None columns = self.get_columns(driver, table=lookup_table, schema=lookup_schema) if not columns: return self._unsupported_ddl(table_name, schema=schema, object_type="table", warnings=warnings) @@ -2781,7 +2781,7 @@ def _synthesize_table_ddl( foreign_keys=foreign_keys, indexes=indexes, target_dialect=target_dialect, - source_dialect=self.resolve_ddl_dialect(driver), + source_dialect=type(self).resolve_ddl_dialect(self, driver), warnings=warnings, ) diff --git a/sqlspec/extensions/adk/_table_utils.py b/sqlspec/extensions/adk/_table_utils.py index ed818ed67..0a6d3707d 100644 --- a/sqlspec/extensions/adk/_table_utils.py +++ b/sqlspec/extensions/adk/_table_utils.py @@ -40,17 +40,27 @@ def reset_drop_sql( return unique_statements(statements) -def ensure_table_name(table_name: str) -> None: - """Validate table name for SQL safety.""" +def ensure_table_name(table_name: str, *, max_parts: int = 1) -> None: + """Validate table identifier components supported by the store. + + Args: + table_name: Table name to validate. + max_parts: Maximum number of identifier components the store supports. + + Raises: + ValueError: If a name component is empty, invalid, or too long. + """ if not table_name: msg = "Table name cannot be empty" raise ValueError(msg) - if len(table_name) > _MAX_TABLE_NAME_LENGTH: - msg = f"Table name too long: {len(table_name)} chars (max {_MAX_TABLE_NAME_LENGTH})" - raise ValueError(msg) + parts = table_name.split(".") + for part in parts: + if len(part) > _MAX_TABLE_NAME_LENGTH: + msg = f"Table name too long: {len(part)} chars (max {_MAX_TABLE_NAME_LENGTH})" + raise ValueError(msg) - if not _VALID_TABLE_NAME_PATTERN.match(table_name): + if len(parts) > max_parts or any(not _VALID_TABLE_NAME_PATTERN.fullmatch(part) for part in parts): msg = ( f"Invalid table name: {table_name!r}. " "Must start with letter/underscore and contain only alphanumeric characters and underscores" diff --git a/sqlspec/extensions/adk/memory/store.py b/sqlspec/extensions/adk/memory/store.py index 8a437ac91..3e27e4d2a 100644 --- a/sqlspec/extensions/adk/memory/store.py +++ b/sqlspec/extensions/adk/memory/store.py @@ -2,7 +2,7 @@ import logging from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, Any, Final, Generic, Literal, TypeVar, cast +from typing import TYPE_CHECKING, Any, ClassVar, Final, Generic, Literal, TypeVar, cast from sqlspec.extensions.adk._config_utils import _adk_memory_store_config, _ADKMemoryStoreConfig from sqlspec.extensions.adk._table_utils import ensure_table_name, owner_id_column_name, unique_statements @@ -46,6 +46,8 @@ class _ADKMemoryStoreCommon(Generic[ConfigT]): "_vector_index_type", ) + _max_table_name_parts: ClassVar[int] = 1 + def __init__(self, config: ConfigT) -> None: """Initialize the ADK memory store. @@ -56,7 +58,7 @@ def __init__(self, config: ConfigT) -> None: store_config = self._store_config_from_extension() self._enabled: bool = store_config.get("enable_memory", True) self._memory_table: str = str(store_config["memory_table"]) - ensure_table_name(self._memory_table) + ensure_table_name(self._memory_table, max_parts=self._max_table_name_parts) self._use_fts: bool = bool(store_config.get("use_fts", False)) self._max_results: int = store_config.get("max_results", 20) self._vector_index_type: str = store_config.get("vector_index_type", "hnsw") diff --git a/sqlspec/extensions/adk/migrations/0002_create_mssql_python_memory.py b/sqlspec/extensions/adk/migrations/0002_create_mssql_python_memory.py index 2c5727415..a4be50a1d 100644 --- a/sqlspec/extensions/adk/migrations/0002_create_mssql_python_memory.py +++ b/sqlspec/extensions/adk/migrations/0002_create_mssql_python_memory.py @@ -26,20 +26,16 @@ async def up(context: "MigrationContext | None" = None) -> list[str]: if not any(cls.__module__.startswith("sqlspec.adapters.mssql_python.") for cls in type(config).__mro__): return [] - from sqlspec.adapters.mssql_python.adk.store import ( - MssqlPythonADKMemoryStore, - _create_index_sql, - _escape_sql_literal, - _table_ref, - ) + from sqlspec.adapters.mssql_python.adk.store import MssqlPythonADKMemoryStore, _create_index_sql, _table_ref + from sqlspec.adapters.mssql_python.core import escape_sql_literal store = MssqlPythonADKMemoryStore(config) statements = [store._memory_table_ddl()] # pyright: ignore[reportPrivateUsage] for index_name, table, columns in store._memory_index_specs(): # pyright: ignore[reportPrivateUsage] statements.append( "IF NOT EXISTS (SELECT 1 FROM sys.indexes " - f"WHERE name = N'{_escape_sql_literal(index_name)}' " - f"AND object_id = OBJECT_ID(N'{_escape_sql_literal(_table_ref(table))}')) " + f"WHERE name = N'{escape_sql_literal(index_name)}' " + f"AND object_id = OBJECT_ID(N'{escape_sql_literal(_table_ref(table))}')) " f"BEGIN {_create_index_sql(table, index_name, columns)}; END;" ) return statements diff --git a/sqlspec/extensions/adk/store.py b/sqlspec/extensions/adk/store.py index 3b5af519a..b4a3316f1 100644 --- a/sqlspec/extensions/adk/store.py +++ b/sqlspec/extensions/adk/store.py @@ -4,7 +4,7 @@ import logging from abc import ABC, abstractmethod from datetime import UTC, datetime, timedelta -from typing import TYPE_CHECKING, Any, Final, Generic, TypeVar, cast +from typing import TYPE_CHECKING, Any, ClassVar, Final, Generic, TypeVar, cast from sqlspec.extensions.adk._config_utils import _adk_session_store_config from sqlspec.extensions.adk._table_utils import ensure_table_name, owner_id_column_name, unique_statements @@ -47,6 +47,8 @@ class _ADKStoreCommon(Generic[ConfigT]): "_user_state_table", ) + _max_table_name_parts: ClassVar[int] = 1 + def __init__(self, config: ConfigT) -> None: """Initialize the ADK store. @@ -73,11 +75,11 @@ def __init__(self, config: ConfigT) -> None: self._owner_id_column_name: str | None = ( owner_id_column_name(self._owner_id_column_ddl) if self._owner_id_column_ddl else None ) - ensure_table_name(self._session_table) - ensure_table_name(self._events_table) - ensure_table_name(self._app_state_table) - ensure_table_name(self._user_state_table) - ensure_table_name(self._metadata_table) + ensure_table_name(self._session_table, max_parts=self._max_table_name_parts) + ensure_table_name(self._events_table, max_parts=self._max_table_name_parts) + ensure_table_name(self._app_state_table, max_parts=self._max_table_name_parts) + ensure_table_name(self._user_state_table, max_parts=self._max_table_name_parts) + ensure_table_name(self._metadata_table, max_parts=self._max_table_name_parts) @property def config(self) -> ConfigT: diff --git a/sqlspec/extensions/fastapi/extension.py b/sqlspec/extensions/fastapi/extension.py index 3d61e669d..414c714eb 100644 --- a/sqlspec/extensions/fastapi/extension.py +++ b/sqlspec/extensions/fastapi/extension.py @@ -221,7 +221,7 @@ def _extract_extension_settings(self, config: Any) -> "dict[str, Any]": "extra_commit_statuses": framework_config.get("extra_commit_statuses"), "extra_rollback_statuses": framework_config.get("extra_rollback_statuses"), "disable_di": framework_config.get("disable_di", False), - "enable_correlation_middleware": framework_config.get("enable_correlation_middleware", False), + "enable_correlation_middleware": framework_config.get("enable_correlation_middleware", True), "correlation_header": framework_config.get("correlation_header", "x-request-id"), "correlation_headers": tuple(correlation_headers) if correlation_headers is not None else None, "auto_trace_headers": framework_config.get("auto_trace_headers", True), diff --git a/sqlspec/extensions/flask/_state.py b/sqlspec/extensions/flask/_state.py index 84fce2ae8..d214b007f 100644 --- a/sqlspec/extensions/flask/_state.py +++ b/sqlspec/extensions/flask/_state.py @@ -28,7 +28,7 @@ class FlaskConfigState: extra_rollback_statuses: "set[int] | None" is_async: bool disable_di: bool - enable_correlation_middleware: bool = False + enable_correlation_middleware: bool = True correlation_header: str = "x-request-id" correlation_headers: "tuple[str, ...] | None" = None auto_trace_headers: bool = True diff --git a/sqlspec/extensions/flask/extension.py b/sqlspec/extensions/flask/extension.py index b01acfd0c..b469b75ad 100644 --- a/sqlspec/extensions/flask/extension.py +++ b/sqlspec/extensions/flask/extension.py @@ -218,7 +218,7 @@ def _config_state(self, config: Any) -> FlaskConfigState: extra_rollback_statuses = flask_config.get("extra_rollback_statuses") disable_di = flask_config.get("disable_di", False) - enable_correlation = flask_config.get("enable_correlation_middleware", False) + enable_correlation = flask_config.get("enable_correlation_middleware", True) correlation_header = flask_config.get("correlation_header", "x-request-id") correlation_headers = flask_config.get("correlation_headers") if correlation_headers is not None: diff --git a/sqlspec/extensions/sanic/_state.py b/sqlspec/extensions/sanic/_state.py index 81e68932f..ad4ea2569 100644 --- a/sqlspec/extensions/sanic/_state.py +++ b/sqlspec/extensions/sanic/_state.py @@ -30,7 +30,7 @@ class SanicConfigState: extra_commit_statuses: "set[int] | None" extra_rollback_statuses: "set[int] | None" disable_di: bool - enable_correlation_middleware: bool = False + enable_correlation_middleware: bool = True correlation_header: str = "x-request-id" correlation_headers: "tuple[str, ...] | None" = None auto_trace_headers: bool = True diff --git a/sqlspec/extensions/sanic/extension.py b/sqlspec/extensions/sanic/extension.py index 2f20ecd4c..a23ca30dc 100644 --- a/sqlspec/extensions/sanic/extension.py +++ b/sqlspec/extensions/sanic/extension.py @@ -141,7 +141,7 @@ def _extract_extension_settings(self, config: Any) -> "dict[str, Any]": "extra_commit_statuses": framework_config.get("extra_commit_statuses"), "extra_rollback_statuses": framework_config.get("extra_rollback_statuses"), "disable_di": framework_config.get("disable_di", False), - "enable_correlation_middleware": framework_config.get("enable_correlation_middleware", False), + "enable_correlation_middleware": framework_config.get("enable_correlation_middleware", True), "correlation_header": framework_config.get("correlation_header", "x-request-id"), "correlation_headers": tuple(correlation_headers) if correlation_headers is not None else None, "auto_trace_headers": framework_config.get("auto_trace_headers", True), diff --git a/sqlspec/extensions/starlette/_state.py b/sqlspec/extensions/starlette/_state.py index 4f6d56f15..7f1b28ead 100644 --- a/sqlspec/extensions/starlette/_state.py +++ b/sqlspec/extensions/starlette/_state.py @@ -29,7 +29,7 @@ class SQLSpecConfigState: extra_commit_statuses: "set[int] | None" extra_rollback_statuses: "set[int] | None" disable_di: bool - enable_correlation_middleware: bool = False + enable_correlation_middleware: bool = True correlation_header: str = "x-request-id" correlation_headers: "tuple[str, ...] | None" = None auto_trace_headers: bool = True diff --git a/sqlspec/extensions/starlette/extension.py b/sqlspec/extensions/starlette/extension.py index 89220a40d..957c89e82 100644 --- a/sqlspec/extensions/starlette/extension.py +++ b/sqlspec/extensions/starlette/extension.py @@ -212,7 +212,7 @@ def _extract_extension_settings(self, config: Any) -> "dict[str, Any]": "extra_commit_statuses": framework_config.get("extra_commit_statuses"), "extra_rollback_statuses": framework_config.get("extra_rollback_statuses"), "disable_di": framework_config.get("disable_di", False), - "enable_correlation_middleware": framework_config.get("enable_correlation_middleware", False), + "enable_correlation_middleware": framework_config.get("enable_correlation_middleware", True), "correlation_header": framework_config.get("correlation_header", "x-request-id"), "correlation_headers": tuple(correlation_headers) if correlation_headers is not None else None, "auto_trace_headers": framework_config.get("auto_trace_headers", True), diff --git a/sqlspec/migrations/squash.py b/sqlspec/migrations/squash.py index d7f87bca3..71d4b85ac 100644 --- a/sqlspec/migrations/squash.py +++ b/sqlspec/migrations/squash.py @@ -25,7 +25,7 @@ from sqlglot.errors import ErrorLevel, ParseError, TokenError, UnsupportedError from sqlglot.tokenizer_core import Token, TokenType -from sqlspec.data_dictionary._types import DDLResult +from sqlspec.data_dictionary import DDLResult from sqlspec.driver import SyncDriverAdapterBase from sqlspec.exceptions import SquashValidationError from sqlspec.migrations._backup import create_backup, remove_backup, restore_backup diff --git a/sqlspec/typing.py b/sqlspec/typing.py index 938adca52..51cf32a65 100644 --- a/sqlspec/typing.py +++ b/sqlspec/typing.py @@ -52,6 +52,7 @@ if TYPE_CHECKING: from sqlspec._typing import ( + ArrowException, ArrowRecordBatch, ArrowRecordBatchReader, ArrowRecordBatchReaderProtocol, @@ -106,6 +107,7 @@ "PYDANTIC_USE_FAILFAST", "UNSET", "UUID_UTILS_INSTALLED", + "ArrowException", "ArrowRecordBatch", "ArrowRecordBatchReader", "ArrowRecordBatchReaderProtocol", diff --git a/sqlspec/utils/arrow_helpers.py b/sqlspec/utils/arrow_helpers.py index 8627be80e..a27c08f67 100644 --- a/sqlspec/utils/arrow_helpers.py +++ b/sqlspec/utils/arrow_helpers.py @@ -16,13 +16,19 @@ from sqlspec.data_dictionary._registry import get_dialect_config from sqlspec.exceptions import ImproperConfigurationError, SQLBuilderError from sqlspec.utils.dispatch import TypeDispatcher -from sqlspec.utils.module_loader import ensure_pandas, ensure_polars, ensure_pyarrow, import_optional_attr +from sqlspec.utils.module_loader import ( + ensure_pandas, + ensure_polars, + ensure_pyarrow, + import_optional, + import_optional_attr, +) from sqlspec.utils.type_guards import has_arrow_table_stats, has_get_data from sqlspec.utils.uuids import uuid_from_bytes if TYPE_CHECKING: from sqlspec.core.result import ArrowResult - from sqlspec.data_dictionary._types import DialectConfig + from sqlspec.data_dictionary import DialectConfig from sqlspec.storage import StorageTelemetry from sqlspec.typing import ArrowRecordBatch, ArrowRecordBatchReader, ArrowTable, PandasDataFrame, PolarsDataFrame @@ -46,12 +52,16 @@ "cast_arrow_table_schema", "coerce_arrow_batches", "coerce_arrow_table", + "coerce_oracle_df_to_arrow_batches", + "coerce_oracle_df_to_arrow_table", "convert_dict_to_arrow", "convert_dict_to_arrow_with_schema", "encode_arrow_batch_pg_binary", "ensure_arrow_table", "iter_pgcopy_binary_chunks", "iter_pgcopy_binary_chunks_async", + "oracle_df_batches_to_arrow_reader", + "oracle_df_to_arrow_table", "records_to_arrow_table", "sql_type_for_category", "sql_type_from_arrow_type", @@ -1455,3 +1465,68 @@ async def __anext__(self) -> bytes: return next(self._iterator) except StopIteration as exc: raise StopAsyncIteration from exc + + +def oracle_df_to_arrow_table(oracle_df: Any, pa_module: Any = None, arrow_schema: Any = None) -> "ArrowTable": + """Convert an Oracle DataFrame to a PyArrow Table using zero-copy PyCapsule when supported.""" + if pa_module is None: + ensure_pyarrow() + pa_module = import_optional("pyarrow") + + if hasattr(oracle_df, "__arrow_c_stream__"): + reader = ( + pa_module.RecordBatchReader.from_stream(oracle_df, schema=arrow_schema) + if arrow_schema is not None + else pa_module.RecordBatchReader.from_stream(oracle_df) + ) + return cast("ArrowTable", reader.read_all()) + table = pa_module.table(oracle_df) + if arrow_schema is not None: + table = table.cast(arrow_schema) + return cast("ArrowTable", table) + + +def coerce_oracle_df_to_arrow_table(oracle_df: Any, pa_module: Any = None, arrow_schema: Any = None) -> "ArrowTable": + """Alias for :func:`oracle_df_to_arrow_table`.""" + return oracle_df_to_arrow_table(oracle_df, pa_module=pa_module, arrow_schema=arrow_schema) + + +def coerce_oracle_df_to_arrow_batches( + oracle_df: Any, pa_module: Any = None, arrow_schema: Any = None +) -> "list[ArrowRecordBatch]": + """Convert an Oracle DataFrame batch into PyArrow RecordBatches via PyCapsule when supported.""" + if pa_module is None: + ensure_pyarrow() + pa_module = import_optional("pyarrow") + + if hasattr(oracle_df, "__arrow_c_stream__"): + reader = ( + pa_module.RecordBatchReader.from_stream(oracle_df, schema=arrow_schema) + if arrow_schema is not None + else pa_module.RecordBatchReader.from_stream(oracle_df) + ) + return cast("list[ArrowRecordBatch]", list(reader)) + table = pa_module.table(oracle_df) + if arrow_schema is not None: + table = table.cast(arrow_schema) + return cast("list[ArrowRecordBatch]", table.to_batches()) + + +def oracle_df_batches_to_arrow_reader( + batches: "Iterable[Any]", schema: Any = None, pa_module: Any = None +) -> "ArrowRecordBatchReader": + """Convert an iterable of Oracle DataFrame batches into a PyArrow RecordBatchReader.""" + if pa_module is None: + ensure_pyarrow() + pa_module = import_optional("pyarrow") + + collected_batches: list[Any] = [] + for df_batch in batches: + collected_batches.extend(coerce_oracle_df_to_arrow_batches(df_batch, pa_module=pa_module, arrow_schema=schema)) + if collected_batches: + resolved_schema = schema if schema is not None else collected_batches[0].schema + return cast( + "ArrowRecordBatchReader", pa_module.RecordBatchReader.from_batches(resolved_schema, collected_batches) + ) + resolved_schema = schema if schema is not None else pa_module.schema([]) + return cast("ArrowRecordBatchReader", pa_module.RecordBatchReader.from_batches(resolved_schema, [])) diff --git a/tests/conftest.py b/tests/conftest.py index 82a709a25..b5b0ccd34 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -103,6 +103,12 @@ def pytest_addoption(parser: pytest.Parser) -> None: default=False, help="Run Spanner integration tests locally (enabled by default in CI; local runs accept --run-spanner-tests or SQLSPEC_ENABLE_SPANNER_TESTS=1)", ) + parser.addoption( + "--run-compiled-adbc", + action="store_true", + default=False, + help="Enable targeted ADBC validation against an installed compiled wheel; explicit interpreted-only skips remain active.", + ) def pytest_collection_modifyitems(config: pytest.Config, items: list[pytest.Item]) -> None: @@ -130,7 +136,7 @@ def pytest_collection_modifyitems(config: pytest.Config, items: list[pytest.Item if is_spanner and not spanner_enabled: item.add_marker(skip_spanner) - if not is_compiled(): + if not is_compiled() or config.getoption("--run-compiled-adbc", default=False): return skip_adbc = pytest.mark.skip(reason="Skip ADBC tests when running against mypyc-compiled modules.") diff --git a/tests/integration/adapters/_shared/_cases.py b/tests/integration/adapters/_shared/_cases.py index 6b6d28da4..15e61dd10 100644 --- a/tests/integration/adapters/_shared/_cases.py +++ b/tests/integration/adapters/_shared/_cases.py @@ -156,6 +156,7 @@ def _db2_case(mode: Literal["sync", "async"], marks: tuple[Mark | MarkDecorator, supports_merge=True, supports_native_row_streaming=True, supports_pooling=True, + supports_connection_instance=True, supports_connection_hook=True, supports_lowercase_columns=True, config_factory_fixture=f"lifecycle_config_db2_{mode}", @@ -438,7 +439,7 @@ def _db2_case(mode: Literal["sync", "async"], marks: tuple[Mark | MarkDecorator, supports_arrow_streaming=True, supports_native_arrow=True, arrow_reader_honors_batch_size=True, - arrow_reader_reports_rows_affected=True, + arrow_reader_reports_rows_affected=False, fixture_name="contract_oracle_sync_driver", adapter="oracledb", dialect="oracle", diff --git a/tests/integration/adapters/_shared/_driver_type_system.py b/tests/integration/adapters/_shared/_driver_type_system.py index 2f699c07a..8a7ff883b 100644 --- a/tests/integration/adapters/_shared/_driver_type_system.py +++ b/tests/integration/adapters/_shared/_driver_type_system.py @@ -1,26 +1,17 @@ -"""Driver-type-system registry and source-equality helpers for contract tests.""" +"""Driver-type-system registry for adapter contract tests.""" -import ast -import difflib import importlib -import inspect -import textwrap from dataclasses import dataclass -from typing import Any, cast +from typing import Any __all__ = ( "DRIVER_FEATURE_CONSUMED_KEYS", "DRIVER_FEATURE_CONTRACT_ENABLED", "DRIVER_FEATURE_CONTRACT_GROUPS", "DRIVER_FEATURE_TYPED_DICTS", - "SOURCE_EQUIVALENCE_ALLOWLIST", - "SOURCE_EQUIVALENCE_CASES", "DriverFeatureContractGroup", - "SourceEquivalenceCase", - "assert_source_equivalent", "get_driver_feature_consumed_keys", "get_driver_feature_keys", - "normalize_source", ) @@ -33,17 +24,6 @@ class DriverFeatureContractGroup: enabled_adapters: tuple[str, ...] = () -@dataclass(frozen=True) -class SourceEquivalenceCase: - """One source-equality tripwire across adapter helper implementations.""" - - case_id: str - group: str - symbol: str - module_paths: tuple[str, ...] - allowlist_key: str - - DRIVER_FEATURE_TYPED_DICTS: dict[str, tuple[str, str]] = { "adbc": ("sqlspec.adapters.adbc.config", "AdbcDriverFeatures"), "aiomysql": ("sqlspec.adapters.aiomysql.config", "AiomysqlDriverFeatures"), @@ -97,9 +77,7 @@ class SourceEquivalenceCase: "json_serializer", "enable_cast_detection", "enable_strict_type_coercion", - "strict_type_coercion", "enable_arrow_extension_types", - "arrow_extension_types", "enable_pgvector", "enable_paradedb", "enable_pg_textsearch", @@ -139,6 +117,7 @@ class SourceEquivalenceCase: "extensions", ), "arrow_odbc": ( + "events_backend", "chunk_size", "max_bytes_per_batch", "max_text_size", @@ -146,7 +125,7 @@ class SourceEquivalenceCase: "fetch_concurrently", "query_timeout_sec", "payload_text_encoding", - "enable_driver_pooling", + "enable_pooling", "connection_string", "dbms_name", "connection_autocommit", @@ -165,6 +144,7 @@ class SourceEquivalenceCase: "enable_local_infile_bulk_load", ), "asyncpg": ( + "enable_pgcopy_binary", "json_serializer", "json_deserializer", "enable_json_codecs", @@ -184,12 +164,12 @@ class SourceEquivalenceCase: "events_backend", "on_connection_create", "pooler_mode", - "pgbouncer", - "pgbouncer_mode", "direct_dsn", "type_codecs", ), "bigquery": ( + "default_query_job_config", + "json_deserializer", "connection_instance", "on_connection_create", "json_serializer", @@ -198,7 +178,7 @@ class SourceEquivalenceCase: "events_backend", "job_retry_deadline", "job_result_timeout", - "use_query_and_wait", + "enable_query_and_wait", "request_timeout", "query_page_size", "query_max_results", @@ -261,7 +241,7 @@ class SourceEquivalenceCase: "enable_statement_profiling", ), "mssql_python": ( - "use_pool", + "enable_pooling", "json_serializer", "json_deserializer", "on_connection_create", @@ -297,6 +277,7 @@ class SourceEquivalenceCase: "dbms_cloud_credential", ), "psqlpy": ( + "enable_pgcopy_binary", "enable_cast_detection", "enable_pgvector", "enable_paradedb", @@ -308,11 +289,10 @@ class SourceEquivalenceCase: "enable_events", "events_backend", "pooler_mode", - "pgbouncer", - "pgbouncer_mode", "direct_dsn", ), "psycopg": ( + "enable_pgcopy_binary", "enable_pgvector", "enable_paradedb", "enable_pg_textsearch", @@ -327,12 +307,18 @@ class SourceEquivalenceCase: "alloydb_ip_type", "enable_alloydb_columnar", "pooler_mode", - "pgbouncer", - "pgbouncer_mode", "direct_dsn", - "null_pool", + "enable_pooling", + ), + "pymssql": ( + "enable_pooling", + "enable_uuid_conversion", + "json_serializer", + "json_deserializer", + "on_connection_create", + "enable_events", + "events_backend", ), - "pymssql": ("json_serializer", "json_deserializer", "on_connection_create", "enable_events", "events_backend"), "pymysql": ( "json_serializer", "json_deserializer", @@ -346,6 +332,7 @@ class SourceEquivalenceCase: "cloud_sql_ip_type", ), "spanner": ( + "on_connection_create", "enable_uuid_conversion", "json_serializer", "json_deserializer", @@ -419,184 +406,6 @@ class SourceEquivalenceCase: group_name: group.enabled for group_name, group in DRIVER_FEATURE_CONTRACT_GROUPS.items() } -SOURCE_EQUIVALENCE_CASES: tuple[SourceEquivalenceCase, ...] = ( - SourceEquivalenceCase( - case_id="mysql_four_way_format_identifier", - group="mysql_four_way", - symbol="format_identifier", - module_paths=( - "sqlspec.adapters.aiomysql.core", - "sqlspec.adapters.asyncmy.core", - "sqlspec.adapters.mysqlconnector.core", - "sqlspec.adapters.pymysql.core", - ), - allowlist_key="mysql_four_way:format_identifier", - ), - SourceEquivalenceCase( - case_id="mysql_four_way_build_insert_statement", - group="mysql_four_way", - symbol="build_insert_statement", - module_paths=( - "sqlspec.adapters.aiomysql.core", - "sqlspec.adapters.asyncmy.core", - "sqlspec.adapters.mysqlconnector.core", - "sqlspec.adapters.pymysql.core", - ), - allowlist_key="mysql_four_way:build_insert_statement", - ), - SourceEquivalenceCase( - case_id="mysql_four_way_normalize_execute_parameters", - group="mysql_four_way", - symbol="normalize_execute_parameters", - module_paths=( - "sqlspec.adapters.aiomysql.core", - "sqlspec.adapters.asyncmy.core", - "sqlspec.adapters.mysqlconnector.core", - "sqlspec.adapters.pymysql.core", - ), - allowlist_key="mysql_four_way:normalize_execute_parameters", - ), - SourceEquivalenceCase( - case_id="mysql_four_way_normalize_execute_many_parameters", - group="mysql_four_way", - symbol="normalize_execute_many_parameters", - module_paths=( - "sqlspec.adapters.aiomysql.core", - "sqlspec.adapters.asyncmy.core", - "sqlspec.adapters.mysqlconnector.core", - "sqlspec.adapters.pymysql.core", - ), - allowlist_key="mysql_four_way:normalize_execute_many_parameters", - ), - SourceEquivalenceCase( - case_id="mysql_four_way_resolve_column_names", - group="mysql_four_way", - symbol="resolve_column_names", - module_paths=( - "sqlspec.adapters.aiomysql.core", - "sqlspec.adapters.asyncmy.core", - "sqlspec.adapters.mysqlconnector.core", - "sqlspec.adapters.pymysql.core", - ), - allowlist_key="mysql_four_way:resolve_column_names", - ), - SourceEquivalenceCase( - case_id="mysql_four_way_detect_json_columns_from_description", - group="mysql_four_way", - symbol="detect_json_columns_from_description", - module_paths=( - "sqlspec.adapters.aiomysql.core", - "sqlspec.adapters.asyncmy.core", - "sqlspec.adapters.mysqlconnector.core", - "sqlspec.adapters.pymysql.core", - ), - allowlist_key="mysql_four_way:detect_json_columns_from_description", - ), - SourceEquivalenceCase( - case_id="mysql_four_way_detect_json_columns", - group="mysql_four_way", - symbol="detect_json_columns", - module_paths=( - "sqlspec.adapters.aiomysql.core", - "sqlspec.adapters.asyncmy.core", - "sqlspec.adapters.mysqlconnector.core", - "sqlspec.adapters.pymysql.core", - ), - allowlist_key="mysql_four_way:detect_json_columns", - ), - SourceEquivalenceCase( - case_id="mysql_four_way_resolve_row_plan", - group="mysql_four_way", - symbol="resolve_row_plan", - module_paths=( - "sqlspec.adapters.aiomysql.core", - "sqlspec.adapters.asyncmy.core", - "sqlspec.adapters.mysqlconnector.core", - "sqlspec.adapters.pymysql.core", - ), - allowlist_key="mysql_four_way:resolve_row_plan", - ), - SourceEquivalenceCase( - case_id="mysql_four_way_collect_rows", - group="mysql_four_way", - symbol="collect_rows", - module_paths=( - "sqlspec.adapters.aiomysql.core", - "sqlspec.adapters.asyncmy.core", - "sqlspec.adapters.mysqlconnector.core", - "sqlspec.adapters.pymysql.core", - ), - allowlist_key="mysql_four_way:collect_rows", - ), - SourceEquivalenceCase( - case_id="mysql_four_way_collect_stream_rows", - group="mysql_four_way", - symbol="collect_stream_rows", - module_paths=( - "sqlspec.adapters.aiomysql.core", - "sqlspec.adapters.asyncmy.core", - "sqlspec.adapters.mysqlconnector.core", - "sqlspec.adapters.pymysql.core", - ), - allowlist_key="mysql_four_way:collect_stream_rows", - ), - SourceEquivalenceCase( - case_id="mysql_four_way_create_mapped_exception", - group="mysql_four_way", - symbol="create_mapped_exception", - module_paths=( - "sqlspec.adapters.aiomysql.core", - "sqlspec.adapters.asyncmy.core", - "sqlspec.adapters.mysqlconnector.core", - "sqlspec.adapters.pymysql.core", - ), - allowlist_key="mysql_four_way:create_mapped_exception", - ), - SourceEquivalenceCase( - case_id="sqlite_pair_apply_driver_features", - group="sqlite_pair", - symbol="apply_driver_features", - module_paths=("sqlspec.adapters.sqlite.core", "sqlspec.adapters.aiosqlite.core"), - allowlist_key="sqlite_pair:apply_driver_features", - ), - SourceEquivalenceCase( - case_id="sqlite_pair_build_insert_statement", - group="sqlite_pair", - symbol="build_insert_statement", - module_paths=("sqlspec.adapters.sqlite.core", "sqlspec.adapters.aiosqlite.core"), - allowlist_key="sqlite_pair:build_insert_statement", - ), - SourceEquivalenceCase( - case_id="sqlite_pair_normalize_execute_parameters", - group="sqlite_pair", - symbol="normalize_execute_parameters", - module_paths=("sqlspec.adapters.sqlite.core", "sqlspec.adapters.aiosqlite.core"), - allowlist_key="sqlite_pair:normalize_execute_parameters", - ), - SourceEquivalenceCase( - case_id="sqlite_pair_normalize_execute_many_parameters", - group="sqlite_pair", - symbol="normalize_execute_many_parameters", - module_paths=("sqlspec.adapters.sqlite.core", "sqlspec.adapters.aiosqlite.core"), - allowlist_key="sqlite_pair:normalize_execute_many_parameters", - ), - SourceEquivalenceCase( - case_id="mssql_pair_create_mapped_exception", - group="mssql_pair", - symbol="create_mapped_exception", - module_paths=("sqlspec.adapters.mssql_python.core", "sqlspec.adapters.pymssql.core"), - allowlist_key="mssql_pair:create_mapped_exception", - ), -) - -SOURCE_EQUIVALENCE_ALLOWLIST: dict[str, frozenset[str]] = { - case.allowlist_key: frozenset() for case in SOURCE_EQUIVALENCE_CASES -} - -SOURCE_EQUIVALENCE_ALLOWLIST["mysql_four_way:create_mapped_exception"] = frozenset({ - "sqlspec.adapters.mysqlconnector.core" -}) - def get_driver_feature_keys(adapter: str) -> tuple[str, ...]: """Return the canonical DriverFeatures keys declared by an adapter TypedDict.""" @@ -614,75 +423,6 @@ def get_driver_feature_consumed_keys(adapter: str) -> tuple[str, ...]: return DRIVER_FEATURE_CONSUMED_KEYS[adapter] -class _DocstringStripper(ast.NodeTransformer): - def visit_Module(self, node: ast.Module) -> ast.Module: - node.body = self._strip_body(node.body) - return cast("ast.Module", self.generic_visit(node)) - - def visit_FunctionDef(self, node: ast.FunctionDef) -> ast.FunctionDef: - node.body = self._strip_body(node.body) - return cast("ast.FunctionDef", self.generic_visit(node)) - - def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> ast.AsyncFunctionDef: - node.body = self._strip_body(node.body) - return cast("ast.AsyncFunctionDef", self.generic_visit(node)) - - def visit_ClassDef(self, node: ast.ClassDef) -> ast.ClassDef: - node.body = self._strip_body(node.body) - return cast("ast.ClassDef", self.generic_visit(node)) - - def _strip_body(self, body: list[ast.stmt]) -> list[ast.stmt]: - if body and isinstance(body[0], ast.Expr): - value = body[0].value - if isinstance(value, ast.Constant) and isinstance(value.value, str): - return body[1:] - return body - - -def normalize_source(source: str) -> str: - """Normalize Python source so source-equality checks ignore docstrings and formatting.""" - normalized = textwrap.dedent(source).strip() - tree = ast.parse(normalized) - stripped = _DocstringStripper().visit(tree) - ast.fix_missing_locations(stripped) - return ast.unparse(stripped).strip() - - -def assert_source_equivalent(case: SourceEquivalenceCase) -> None: - """Assert that every unallowlisted module exports the same normalized source.""" - allowed_modules = SOURCE_EQUIVALENCE_ALLOWLIST.get(case.allowlist_key, frozenset()) - normalized_sources: dict[str, str] = {} - for module_path in case.module_paths: - if module_path in allowed_modules: - continue - normalized_sources[module_path] = normalize_source(_load_symbol_source(module_path, case.symbol)) - - if len(normalized_sources) < 2: - msg = f"{case.case_id} does not have enough unallowlisted modules to compare" - raise AssertionError(msg) - - baseline_module, baseline_source = next(iter(normalized_sources.items())) - mismatches: list[str] = [] - for module_path, source in normalized_sources.items(): - if source == baseline_source: - continue - diff = "\n".join( - difflib.unified_diff( - baseline_source.splitlines(), - source.splitlines(), - fromfile=baseline_module, - tofile=module_path, - lineterm="", - ) - ) - mismatches.append(diff) - - if mismatches: - details = "\n\n".join(mismatches) - msg = f"{case.case_id} source mismatch for {case.symbol}:\n{details}" - raise AssertionError(msg) - - def _load_symbol(module_path: str, symbol_name: str) -> Any: try: module = importlib.import_module(module_path) @@ -690,12 +430,3 @@ def _load_symbol(module_path: str, symbol_name: str) -> Any: except (AttributeError, ImportError) as exc: msg = f"Unable to load {symbol_name} from {module_path}: {exc}" raise AssertionError(msg) from exc - - -def _load_symbol_source(module_path: str, symbol_name: str) -> str: - symbol = _load_symbol(module_path, symbol_name) - try: - return inspect.getsource(symbol) - except OSError as exc: - msg = f"Unable to read source for {symbol_name} from {module_path}: {exc}" - raise AssertionError(msg) from exc diff --git a/tests/integration/adapters/_shared/behaviors.py b/tests/integration/adapters/_shared/behaviors.py index 7352d0e26..8cfe6cc46 100644 --- a/tests/integration/adapters/_shared/behaviors.py +++ b/tests/integration/adapters/_shared/behaviors.py @@ -620,34 +620,54 @@ async def _assert_aiosqlite_stream_arraysize(driver: object, case: DriverCase) - async def _assert_asyncpg_stream_transaction(driver: object, case: DriverCase) -> None: async_driver = cast("AsyncContractDriver", driver) + connection = cast("Any", driver).connection + was_in_transaction = connection.is_in_transaction() stream = async_driver.select_stream(case.table.select_ordered_sql, chunk_size=10) - await anext(aiter(stream)) - assert stream._source._transaction is not None # pyright: ignore[reportPrivateUsage] - assert stream._source._cursor is not None # pyright: ignore[reportPrivateUsage] - await stream.aclose() - assert stream._source._transaction is None # pyright: ignore[reportPrivateUsage] - assert stream._source._cursor is None # pyright: ignore[reportPrivateUsage] - await async_driver.execute(case.table.select_count_sql) + try: + first = await anext(aiter(stream)) + assert first["value"] == 0 + assert connection.is_in_transaction() is True + finally: + await stream.aclose() + assert connection.is_in_transaction() is was_in_transaction + assert int(cast("int", await async_driver.select_value(case.table.select_count_sql))) == case.streaming_row_count + await async_driver.begin() + try: + await async_driver.execute(case.table.insert_named_sql, name="stream_caller_probe", value=-1, note=None) + async with async_driver.select_stream(case.table.select_ordered_sql, chunk_size=1) as caller_stream: + assert (await anext(aiter(caller_stream)))["value"] == -1 + assert connection.is_in_transaction() is True + finally: + await async_driver.rollback() + assert await async_driver.select(case.table.select_by_name_named_sql, name="stream_caller_probe") == [] -def _assert_psycopg_named_cursor(driver: object, case: DriverCase) -> None: +def _assert_psycopg_stream_transaction(driver: object, case: DriverCase) -> None: sync_driver = cast("SyncContractDriver", driver) + connection = cast("Any", driver).connection + original_status = connection.info.transaction_status stream = sync_driver.select_stream(case.table.select_ordered_sql, chunk_size=10) try: - next(iter(stream)) - assert stream._source._cursor.name # pyright: ignore[reportPrivateUsage] + assert next(iter(stream))["value"] == 0 + assert connection.info.transaction_status == 2 finally: stream.close() + assert connection.info.transaction_status == original_status + assert int(cast("int", sync_driver.select_value(case.table.select_count_sql))) == case.streaming_row_count -async def _assert_psycopg_named_cursor_async(driver: object, case: DriverCase) -> None: +async def _assert_psycopg_stream_transaction_async(driver: object, case: DriverCase) -> None: async_driver = cast("AsyncContractDriver", driver) + connection = cast("Any", driver).connection + original_status = connection.info.transaction_status stream = async_driver.select_stream(case.table.select_ordered_sql, chunk_size=10) try: - await anext(aiter(stream)) - assert stream._source._cursor.name # pyright: ignore[reportPrivateUsage] + assert (await anext(aiter(stream)))["value"] == 0 + assert connection.info.transaction_status == 2 finally: await stream.aclose() + assert connection.info.transaction_status == original_status + assert int(cast("int", await async_driver.select_value(case.table.select_count_sql))) == case.streaming_row_count def _assert_pymysql_sscursor(driver: object, case: DriverCase) -> None: @@ -693,8 +713,8 @@ async def _assert_mysqlconnector_unbuffered_async(driver: object, case: DriverCa register_sync_extra_assertion("streaming_native:sqlite", _STREAMING_SCOPE, _assert_sqlite_stream_arraysize) register_async_extra_assertion("streaming_native:aiosqlite", _STREAMING_SCOPE, _assert_aiosqlite_stream_arraysize) register_async_extra_assertion("streaming_native:asyncpg", _STREAMING_SCOPE, _assert_asyncpg_stream_transaction) -register_sync_extra_assertion("streaming_native:psycopg", _STREAMING_SCOPE, _assert_psycopg_named_cursor) -register_async_extra_assertion("streaming_native:psycopg", _STREAMING_SCOPE, _assert_psycopg_named_cursor_async) +register_sync_extra_assertion("streaming_native:psycopg", _STREAMING_SCOPE, _assert_psycopg_stream_transaction) +register_async_extra_assertion("streaming_native:psycopg", _STREAMING_SCOPE, _assert_psycopg_stream_transaction_async) register_sync_extra_assertion("streaming_native:pymysql", _STREAMING_SCOPE, _assert_pymysql_sscursor) register_async_extra_assertion("streaming_native:aiomysql", _STREAMING_SCOPE, _assert_async_sscursor) register_async_extra_assertion("streaming_native:asyncmy", _STREAMING_SCOPE, _assert_async_sscursor) @@ -758,58 +778,6 @@ class _StreamErrorProbe(Exception): """Sentinel raised to force a mid-stream chunk fetch failure without a live database error.""" -class _SyncTransactionRecorder: - """Delegating transaction wrapper recording whether commit or rollback fired during stream close.""" - - def __init__(self, inner: Any, log: "list[str]") -> None: - object.__setattr__(self, "_inner", inner) - object.__setattr__(self, "_log", log) - - def __getattr__(self, name: str) -> Any: - return getattr(object.__getattribute__(self, "_inner"), name) - - def __setattr__(self, name: str, value: Any) -> None: - setattr(object.__getattribute__(self, "_inner"), name, value) - - def commit(self, *args: Any, **kwargs: Any) -> Any: - object.__getattribute__(self, "_log").append("commit") - return object.__getattribute__(self, "_inner").commit(*args, **kwargs) - - def rollback(self, *args: Any, **kwargs: Any) -> Any: - object.__getattribute__(self, "_log").append("rollback") - return object.__getattribute__(self, "_inner").rollback(*args, **kwargs) - - def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> Any: - object.__getattribute__(self, "_log").append("rollback" if exc_type is not None else "commit") - return object.__getattribute__(self, "_inner").__exit__(exc_type, exc, tb) - - -class _AsyncTransactionRecorder: - """Delegating async transaction wrapper recording whether commit or rollback fired during stream close.""" - - def __init__(self, inner: Any, log: "list[str]") -> None: - object.__setattr__(self, "_inner", inner) - object.__setattr__(self, "_log", log) - - def __getattr__(self, name: str) -> Any: - return getattr(object.__getattribute__(self, "_inner"), name) - - def __setattr__(self, name: str, value: Any) -> None: - setattr(object.__getattribute__(self, "_inner"), name, value) - - async def commit(self, *args: Any, **kwargs: Any) -> Any: - object.__getattribute__(self, "_log").append("commit") - return await object.__getattribute__(self, "_inner").commit(*args, **kwargs) - - async def rollback(self, *args: Any, **kwargs: Any) -> Any: - object.__getattribute__(self, "_log").append("rollback") - return await object.__getattribute__(self, "_inner").rollback(*args, **kwargs) - - async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> Any: - object.__getattribute__(self, "_log").append("rollback" if exc_type is not None else "commit") - return await object.__getattribute__(self, "_inner").__aexit__(exc_type, exc, tb) - - class _SyncStreamFetchFailSource: """Delegating sync stream source that forces the next chunk fetch to fail and delegates close to the origin.""" @@ -846,21 +814,16 @@ def _assert_sync_stream_error_close_rollback(driver: object, case: DriverCase) - """Assert a mid-stream fetch failure rolls back a stream-owned transaction instead of committing.""" sync_driver = cast("SyncContractDriver", driver) table = case.table + sync_driver.rollback() stream = sync_driver.select_stream(table.select_ordered_sql, chunk_size=1) iterator = iter(stream) next(iterator) real_source = stream._source # pyright: ignore[reportPrivateUsage] - assert real_source._transaction is not None # pyright: ignore[reportPrivateUsage] - log: list[str] = [] - real_source._transaction = _SyncTransactionRecorder(real_source._transaction, log) # pyright: ignore[reportPrivateUsage] + sync_driver.execute(table.insert_named_sql, name="stream_rollback_probe", value=-1, note=None) stream._source = _SyncStreamFetchFailSource(real_source) # pyright: ignore[reportPrivateUsage] with pytest.raises(_StreamErrorProbe): next(iterator) - assert log - assert log[0] == "rollback" - assert "commit" not in log - with contextlib.suppress(Exception): - sync_driver.rollback() + assert sync_driver.select(table.select_by_name_named_sql, name="stream_rollback_probe") == [] assert int(cast("int", sync_driver.select_value(table.select_count_sql))) == case.streaming_row_count @@ -868,21 +831,22 @@ async def _assert_async_stream_error_close_rollback(driver: object, case: Driver """Assert a mid-stream fetch failure rolls back a stream-owned transaction instead of committing.""" async_driver = cast("AsyncContractDriver", driver) table = case.table + await async_driver.rollback() stream = async_driver.select_stream(table.select_ordered_sql, chunk_size=1) iterator = aiter(stream) await anext(iterator) real_source = stream._source # pyright: ignore[reportPrivateUsage] - assert real_source._transaction is not None # pyright: ignore[reportPrivateUsage] - log: list[str] = [] - real_source._transaction = _AsyncTransactionRecorder(real_source._transaction, log) # pyright: ignore[reportPrivateUsage] + # CockroachDB rejects another bound statement while the asyncpg cursor's portal is open. + supports_probe_insert = case.adapter != "cockroach_asyncpg" + if supports_probe_insert: + await async_driver.execute(table.insert_named_sql, name="stream_rollback_probe", value=-1, note=None) stream._source = _AsyncStreamFetchFailSource(real_source) # pyright: ignore[reportPrivateUsage] with pytest.raises(_StreamErrorProbe): await anext(iterator) - assert log - assert log[0] == "rollback" - assert "commit" not in log - with contextlib.suppress(Exception): - await async_driver.rollback() + if supports_probe_insert: + assert await async_driver.select(table.select_by_name_named_sql, name="stream_rollback_probe") == [] + else: + assert cast("Any", driver).connection.is_in_transaction() is False assert int(cast("int", await async_driver.select_value(table.select_count_sql))) == case.streaming_row_count @@ -4936,14 +4900,14 @@ def _bigquery_job_controls(driver: object, case: DriverCase) -> None: assert not hasattr(driver, "export_table_to_storage") execute_parameters = inspect.signature(sync_driver.execute).parameters - for parameter_name in ("job_config", "job_retry", "job_result_timeout", "request_timeout", "use_query_and_wait"): + for parameter_name in ("job_config", "job_retry", "job_result_timeout", "request_timeout", "enable_query_and_wait"): assert parameter_name not in execute_parameters assert getattr(driver, "_job_result_timeout") == 30.0 assert getattr(driver, "_job_retry_deadline") == 0.0 assert getattr(driver, "_job_retry") is None assert getattr(driver, "_request_timeout") == 15.0 - assert getattr(driver, "_use_query_and_wait") is True + assert getattr(driver, "_enable_query_and_wait") is True register_sync_extra_assertion("driver_features:bigquery_job_controls", DRIVER_FEATURES_SCOPE, _bigquery_job_controls) @@ -5064,7 +5028,7 @@ def _pool_contract_table(case: DriverCase) -> str: def _connection_probe_sql(case: DriverCase) -> str: - """Trivial query used to force a connection to be drawn so the create-hook fires.""" + """Scalar query that acquires a native connection for lifecycle contracts.""" if case.dialect == "oracle": return "SELECT 1 FROM DUAL" if case.dialect == "db2": @@ -5168,7 +5132,7 @@ def assert_sync_connection_instance_contract(make_config: SyncConfigFactory, cas try: assert fresh.provide_pool() is not None with fresh.provide_session() as session: - assert session.select_value("SELECT 1") == 1 + assert session.select_value(_connection_probe_sql(case)) == 1 finally: fresh.close_pool() @@ -5199,7 +5163,7 @@ async def assert_async_connection_instance_contract(make_config: AsyncConfigFact try: assert await fresh.provide_pool() is not None async with fresh.provide_session() as session: - assert await session.select_value("SELECT 1") == 1 + assert await session.select_value(_connection_probe_sql(case)) == 1 finally: await fresh.close_pool() @@ -6987,7 +6951,8 @@ async def assert_async_load_from_records_contract(driver: object, case: DriverCa assert await async_driver.select_value(table.select_count_sql) == 5 with pytest.raises(ImproperConfigurationError): - await async_driver.load_from_records(table.name, [], overwrite=True) + await async_driver.load_from_records(table.name, []) + assert await async_driver.select_value(table.select_count_sql) == 5 def _storage_bridge_export_sql(table: ContractTable) -> str: diff --git a/tests/integration/adapters/_shared/suite_driver_features_contract.py b/tests/integration/adapters/_shared/suite_driver_features_contract.py index 0058e2753..2931b2c4a 100644 --- a/tests/integration/adapters/_shared/suite_driver_features_contract.py +++ b/tests/integration/adapters/_shared/suite_driver_features_contract.py @@ -1,6 +1,5 @@ -"""Shared driver-feature contracts: registry scaffolds, source tripwires, and folded driver-specific behaviors.""" +"""Shared driver-feature contracts and folded driver-specific behaviors.""" -import importlib from collections.abc import Sequence from typing import cast @@ -9,7 +8,6 @@ from tests.integration.adapters._shared._cases import ( ASYNC_DRIVER_CASES, - DRIVER_CASES, SYNC_DRIVER_CASES, DriverCase, DriverCaseContext, @@ -17,9 +15,6 @@ ) from tests.integration.adapters._shared._driver_type_system import ( DRIVER_FEATURE_CONTRACT_GROUPS, - SOURCE_EQUIVALENCE_CASES, - SourceEquivalenceCase, - assert_source_equivalent, get_driver_feature_consumed_keys, get_driver_feature_keys, ) @@ -37,30 +32,6 @@ validate_extra_assertions, ) -_DRIVER_TYPE_SYSTEM_MODULE = "tests.integration.adapters._shared._driver_type_system" - - -def test_driver_type_system_registry_module_exports_contract_groups() -> None: - """The driver-type-system registry module is importable and exposes the expected contract groups.""" - module = importlib.import_module(_DRIVER_TYPE_SYSTEM_MODULE) - - assert hasattr(module, "DRIVER_FEATURE_CONTRACT_GROUPS") - assert hasattr(module, "DRIVER_FEATURE_CONTRACT_ENABLED") - assert hasattr(module, "DRIVER_FEATURE_CONSUMED_KEYS") - assert hasattr(module, "SOURCE_EQUIVALENCE_CASES") - assert hasattr(module, "SOURCE_EQUIVALENCE_ALLOWLIST") - assert set(module.DRIVER_FEATURE_CONTRACT_GROUPS) == { - "feature_honesty", - "enable_false_semantics", - "parity", - "row_format", - } - assert module.DRIVER_FEATURE_CONTRACT_ENABLED == { - group_name: group.enabled for group_name, group in module.DRIVER_FEATURE_CONTRACT_GROUPS.items() - } - assert set(module.DRIVER_FEATURE_CONSUMED_KEYS) == {case.adapter for case in DRIVER_CASES} - assert {case.group for case in module.SOURCE_EQUIVALENCE_CASES} == {"mysql_four_way", "sqlite_pair", "mssql_pair"} - def _disabled_first_params(group_name: str, cases: Sequence[DriverCase]) -> tuple[ParameterSet, ...]: group = DRIVER_FEATURE_CONTRACT_GROUPS[group_name] @@ -177,12 +148,6 @@ async def test_async_driver_feature_row_format_contract(async_driver_case: Drive await assert_async_driver_feature_row_format_contract(make_config, async_driver_case.case) -@pytest.mark.parametrize("source_equivalence_case", SOURCE_EQUIVALENCE_CASES, ids=lambda case: case.case_id) -def test_driver_feature_source_equivalence(source_equivalence_case: SourceEquivalenceCase) -> None: - """Duplicate helper implementations stay source-equivalent across adapter families.""" - assert_source_equivalent(source_equivalence_case) - - def test_sync_driver_features_contract(sync_driver_case: DriverCaseContext) -> None: """Sync drivers run folded driver-feature proofs opted in via extra_assertions.""" assert_sync_driver_features_contract(sync_driver_case.driver, sync_driver_case.case) diff --git a/tests/integration/adapters/conftest.py b/tests/integration/adapters/conftest.py index 63acf0525..3fa52d47c 100644 --- a/tests/integration/adapters/conftest.py +++ b/tests/integration/adapters/conftest.py @@ -17,6 +17,7 @@ "oracle": frozenset({"oracledb"}), "postgres": frozenset({"asyncpg", "psqlpy", "psycopg"}), "spanner": frozenset({"spanner"}), + "sqlanywhere": frozenset({"sqlanywhere"}), "sqlite": frozenset({"aiosqlite", "sqlite"}), } _SHARED_TEST_FAMILIES = frozenset((*_FAMILY_ADAPTERS, "gizmosql", "snowflake")) diff --git a/tests/integration/adapters/db2/arrow_odbc/test_db2.py b/tests/integration/adapters/db2/arrow_odbc/test_db2.py new file mode 100644 index 000000000..e725c6642 --- /dev/null +++ b/tests/integration/adapters/db2/arrow_odbc/test_db2.py @@ -0,0 +1,153 @@ +"""Native IBM Db2 acceptance through the 64-bit Arrow-ODBC transport.""" + +import json +import platform +import sys +from typing import TYPE_CHECKING + +import pytest +from pytest_databases.docker.db2 import DB2_PORT +from pytest_databases.helpers import exec_output_to_bytes + +if TYPE_CHECKING: + from pytest_databases.docker.db2 import Db2Service + + from sqlspec.adapters.db2 import Db2ConnectionParams + +pytestmark = [ + pytest.mark.db2, + pytest.mark.integration, + pytest.mark.xdist_group("db2"), + pytest.mark.skipif( + sys.platform != "linux" or platform.machine().lower() not in {"x86_64", "amd64"}, + reason="The IBM 64-bit ODBC runtime acceptance uses a Linux x86-64 client", + ), +] + +_ACCEPTANCE_SCRIPT = """ +import json +import os +from decimal import Decimal +from uuid import uuid4 + +import pyarrow as pa +from sqlspec.adapters.arrow_odbc import ArrowOdbcConfig + +payload = json.loads(os.environ.pop("SQLSPEC_DB2_ODBC_CONFIG")) +config = ArrowOdbcConfig( + connection_config=payload, + driver_features={"enable_pooling": False, "query_timeout_sec": 15}, +) +table_name = "SQLSPEC_ARROW_ODBC_" + uuid4().hex[:12].upper() +index_name = table_name + "_UQ" +created = False +try: + with config.provide_session() as session: + assert session.dialect == "db2" + try: + session.execute( + f"CREATE TABLE {table_name} (" + "id INTEGER NOT NULL PRIMARY KEY, big_value BIGINT, label VARCHAR(64), amount DECIMAL(14, 3))" + ) + session.commit() + created = True + session.execute(f"CREATE UNIQUE INDEX {index_name} ON {table_name} (label)") + session.commit() + insert = f"INSERT INTO {table_name} (id, big_value, label, amount) VALUES (?, ?, ?, ?)" + session.execute(insert, 1, 2**40, "Grüße 東京", Decimal("123.450")) + session.execute(insert, 2, None, None, None) + session.commit() + session.begin() + session.execute(insert, 3, 2**42, "rolled back", Decimal("456.789")) + session.rollback() + assert session.select_value(f"SELECT COUNT(*) FROM {table_name} WHERE id = ?", 3) == 0 + rows = session.select(f"SELECT id, big_value, label, amount FROM {table_name} ORDER BY id") + assert rows[0]["id"] == 1 + assert rows[0]["big_value"] == 2**40 + assert rows[0]["label"] == "Grüße 東京" + assert Decimal(str(rows[0]["amount"])) == Decimal("123.450") + assert rows[1] == {"id": 2, "big_value": None, "label": None, "amount": None} + result = session.select_to_arrow( + f"SELECT id, big_value, label, amount FROM {table_name} WHERE id >= ? ORDER BY id", 1 + ) + table = result.get_data() + assert table.num_rows == 2 + assert table.column_names == ["id", "big_value", "label", "amount"] + assert table.schema.field("big_value").type == pa.int64() + assert table.schema.field("amount").type == pa.decimal128(14, 3) + assert table.column("big_value").to_pylist() == [2**40, None] + assert table.column("label").to_pylist() == ["Grüße 東京", None] + assert table.column("amount").to_pylist() == [Decimal("123.450"), None] + schema_name = str(session.select_value("SELECT CURRENT SCHEMA FROM SYSIBM.SYSDUMMY1")).strip() + dictionary = session.data_dictionary + columns = dictionary.get_columns(session, table=table_name, schema=schema_name) + assert [column["column_name"] for column in columns] == ["ID", "BIG_VALUE", "LABEL", "AMOUNT"] + assert [column["ordinal_position"] for column in columns] == [1, 2, 3, 4] + assert all( + column["table_name"] == table_name and column["schema_name"] == schema_name + for column in columns + ) + assert columns[0]["is_primary"] == 1 + assert columns[2]["is_unique"] == 1 + assert columns[3]["numeric_scale"] == 3 + indexes = dictionary.get_indexes(session, table=table_name, schema=schema_name) + assert len(indexes) == 2 + indexes_by_name = {index["index_name"]: index for index in indexes} + assert len(indexes_by_name) == 2 + assert indexes_by_name[index_name]["is_unique"] == 1 + assert sum(bool(index["is_primary"]) for index in indexes) == 1 + tables = [ + item for item in dictionary.get_tables(session, schema=schema_name) + if item["table_name"] == table_name + ] + assert len(tables) == 1 + assert tables[0]["schema_name"] == schema_name + assert tables[0]["table_type"] == "BASE TABLE" + finally: + session.rollback() + with config.provide_session() as reopened: + try: + persisted = reopened.select(f"SELECT id, big_value FROM {table_name} ORDER BY id") + assert persisted == [{"id": 1, "big_value": 2**40}, {"id": 2, "big_value": None}] + finally: + reopened.rollback() +finally: + try: + if created: + with config.provide_session() as cleanup: + cleanup.rollback() + cleanup.execute(f"DROP TABLE {table_name}") + cleanup.commit() + finally: + config.close_pool() +print(json.dumps({"dialect": "db2", "rows": 2, "transport": "arrow_odbc"})) +""" + + +def test_arrow_odbc_db2_parameter_binding_and_arrow_null_indicators( + db2_service: "Db2Service", + db2_connection_config: "Db2ConnectionParams", + db2_arrow_odbc_runtime: tuple[list[str], dict[str, str], str], +) -> None: + """Transactions, catalog metadata, and Arrow buffers preserve typed Db2 values and NULL indicators.""" + command, environment, driver = db2_arrow_odbc_runtime + password = str(db2_connection_config["password"]) + parameters = { + "driver": driver, + "host": "127.0.0.1", + "database": db2_connection_config["database"], + "uid": db2_connection_config["user"], + "pwd": password, + "extra": {"Port": DB2_PORT, "Protocol": "TCPIP"}, + "login_timeout": 15, + "autocommit": False, + } + result = db2_service.container.exec_run( + [*command, "-c", _ACCEPTANCE_SCRIPT], + environment={**environment, "SQLSPEC_DB2_ODBC_CONFIG": json.dumps(parameters)}, + ) + output = exec_output_to_bytes(result.output).decode("utf-8", errors="replace") + if result.exit_code: + diagnostic = output.replace(password, "") + pytest.fail(f"Native Arrow-ODBC Db2 acceptance failed (exit {result.exit_code}):\n{diagnostic}") + assert json.loads(output) == {"dialect": "db2", "rows": 2, "transport": "arrow_odbc"} diff --git a/tests/integration/adapters/db2/test_adapter.py b/tests/integration/adapters/db2/test_adapter.py index 282c10874..ad5ff4c9e 100644 --- a/tests/integration/adapters/db2/test_adapter.py +++ b/tests/integration/adapters/db2/test_adapter.py @@ -1,15 +1,33 @@ """Integration tests for IBM Db2 behavior outside the shared adapter contracts.""" +import asyncio +import json from collections.abc import Generator +from concurrent.futures import ThreadPoolExecutor +from datetime import date, datetime, time from decimal import Decimal from typing import Any +from uuid import UUID import pytest pytest.importorskip("ibm_db") pytest.importorskip("ibm_db_dbi") +from sqlspec import sql +from sqlspec.adapters.db2.config import Db2AsyncConfig, Db2SyncConfig +from sqlspec.adapters.db2.core import extract_sqlstate from sqlspec.adapters.db2.driver import Db2SyncDriver +from sqlspec.adapters.db2.events import Db2SyncEventQueueStore +from sqlspec.exceptions import ( + CheckViolationError, + ConnectionTimeoutError, + NotNullViolationError, + SQLParsingError, + SQLSpecError, + UniqueViolationError, +) +from sqlspec.extensions.events import EventMessage, SyncTableEventQueue pytestmark = [pytest.mark.db2, pytest.mark.xdist_group("db2")] @@ -113,3 +131,246 @@ def test_db2_integration_arrow_conversion(db2_session: Db2SyncDriver, clean_user assert isinstance(table, pa.Table) assert table.to_pydict() == {"id": [42], "name": ["ArrowUser"]} + + +def test_db2_semicolon_schema_survives_native_dsn_roundtrip(db2_connection_config: dict[str, Any]) -> None: + """The CLI driver receives one semicolon-bearing schema value after DSN parsing and rendering.""" + schema = "SQLSPEC;SEMICOLON" + source = Db2SyncConfig(connection_config={**db2_connection_config, "current_schema": schema}) + config = Db2SyncConfig(connection_config={"dsn": source.get_connection_string()}) + try: + with config.provide_session() as session: + assert session.select_value("SELECT CURRENT SCHEMA FROM SYSIBM.SYSDUMMY1").strip() == schema + finally: + config.close_pool() + + +def test_db2_skip_locked_dequeue_excludes_another_sessions_locked_event( + db2_session: Db2SyncDriver, db2_connection_config: dict[str, Any], db2_contract_table: str +) -> None: + """A consumer claims the second event while a different connection holds the first row locked.""" + config = Db2SyncConfig( + connection_config=dict(db2_connection_config), extension_config={"events": {"queue_table": db2_contract_table}} + ) + store = Db2SyncEventQueueStore(config) + for statement in store.create_statements(): + db2_session.execute_script(statement) + db2_session.commit() + queue = SyncTableEventQueue(config, queue_table=db2_contract_table, select_for_update=True, skip_locked=True) + try: + first_id = queue.publish("native_locking", {"position": 1}) + second_id = queue.publish("native_locking", {"position": 2}) + db2_session.begin() + locked = db2_session.select_one( + sql.select("event_id").from_(db2_contract_table).where_eq("event_id", first_id).for_update() + ) + assert locked["event_id"] == first_id + with ThreadPoolExecutor(max_workers=1) as executor: + pending = executor.submit(queue.dequeue, "native_locking", 0.0) + try: + message = pending.result(timeout=10) + assert isinstance(message, EventMessage) + assert message.event_id == second_id + assert message.payload == {"position": 2} + finally: + db2_session.rollback() + first = queue.dequeue("native_locking", poll_interval=0.0) + assert isinstance(first, EventMessage) + assert first.event_id == first_id + assert queue.dequeue("native_locking", poll_interval=0.0) is None + finally: + db2_session.rollback() + config.close_pool() + + +async def test_db2_async_native_pool_bounds_contending_sessions_and_recovers( + db2_connection_config: dict[str, Any], +) -> None: + """Concurrent native borrowers fill the configured capacity and a third times out without leaking it.""" + config = Db2AsyncConfig(connection_config={**db2_connection_config, "max_size": 2, "acquire_timeout": 0.1}) + entered: asyncio.Queue[int] = asyncio.Queue() + release = asyncio.Event() + + async def hold(value: int) -> None: + async with config.provide_session() as session: + observed = await session.select_value("SELECT CAST(:value AS INTEGER) FROM SYSIBM.SYSDUMMY1", value=value) + await entered.put(observed) + await release.wait() + + workers = [asyncio.create_task(hold(value)) for value in (1, 2)] + try: + values = await asyncio.wait_for(asyncio.gather(entered.get(), entered.get()), timeout=10) + assert set(values) == {1, 2} + pool = await config.provide_pool() + assert (pool.size(), pool.checked_out()) == (2, 2) + with pytest.raises(ConnectionTimeoutError): + async with config.provide_session(): + pytest.fail("A third session exceeded max_size while both native borrowers were held.") + release.set() + await asyncio.wait_for(asyncio.gather(*workers), timeout=10) + assert (pool.size(), pool.checked_out()) == (2, 0) + async with config.provide_session() as session: + assert await session.select_value("SELECT 42 FROM SYSIBM.SYSDUMMY1") == 42 + finally: + release.set() + try: + await asyncio.gather(*workers) + finally: + await config.close_pool() + + +@pytest.mark.parametrize( + ("statement", "parameters", "error_type", "sqlstate"), + [ + ( + "INSERT INTO TEST_INTEGRATION_USERS (id, name, balance) VALUES (:id, :name, :balance)", + {"id": 1, "name": "duplicate", "balance": 1}, + UniqueViolationError, + "23505", + ), + ( + "INSERT INTO TEST_INTEGRATION_USERS (id, name, balance) VALUES (:id, :name, :balance)", + {"id": 2, "name": None, "balance": 1}, + NotNullViolationError, + "23502", + ), + ( + "INSERT INTO TEST_INTEGRATION_USERS (id, name, balance) VALUES (:id, :name, :balance)", + {"id": 2, "name": "negative", "balance": -1}, + CheckViolationError, + "23513", + ), + ("SELECT * FROM SQLSPEC_MISSING_NATIVE_TABLE", {}, SQLParsingError, "42704"), + ], + ids=["unique", "not-null", "check", "missing-table"], +) +def test_db2_native_sqlstate_maps_exact_public_exception( + db2_session: Db2SyncDriver, + clean_users_table: None, + statement: str, + parameters: dict[str, Any], + error_type: type[SQLSpecError], + sqlstate: str, +) -> None: + """Real CLI diagnostics preserve representative SQLSTATEs and their specific mapped exception types.""" + db2_session.execute_script( + "ALTER TABLE TEST_INTEGRATION_USERS ADD CONSTRAINT SQLSPEC_BALANCE_CHECK CHECK (balance >= 0)" + ) + db2_session.execute("INSERT INTO TEST_INTEGRATION_USERS (id, name) VALUES (:id, :name)", id=1, name="retained") + db2_session.commit() + with pytest.raises(error_type) as caught: + db2_session.execute(statement, **parameters) + assert extract_sqlstate(caught.value) == sqlstate + db2_session.rollback() + + +@pytest.mark.parametrize("multiple_rows", [False, True], ids=["mapping", "two-rows"]) +def test_db2_merge_extended_typed_source_roundtrip( + db2_session: Db2SyncDriver, db2_contract_table: str, multiple_rows: bool +) -> None: + """The native engine executes the builder's bool, temporal, UUID, CLOB, binary, float, JSON and NULL casts.""" + db2_session.execute_script( + f"CREATE TABLE {db2_contract_table} (" + "record_id INTEGER NOT NULL PRIMARY KEY, label VARCHAR(40), enabled BOOLEAN, amount DECFLOAT(34), " + "captured_at TIMESTAMP, uuid_value VARCHAR(36), body CLOB, day_value DATE, time_value TIME, " + "binary_value VARBINARY(64), ratio DOUBLE, payload CLOB, absent CLOB)" + ) + db2_session.commit() + identifier = UUID("12345678-1234-5678-1234-567812345678") + values = { + "record_id": 1, + "label": "typed", + "enabled": True, + "amount": Decimal("12345678901234567890.123456789"), + "captured_at": datetime(2026, 6, 7, 8, 9, 10, 345678), + "uuid_value": identifier, + "body": "x" * 32673, + "day_value": date(2026, 6, 7), + "time_value": time(12, 34, 56), + "binary_value": b"\x00\xff", + "ratio": 0.5, + "payload": {"nested": "value"}, + "absent": None, + } + query = ( + sql + .merge(dialect="db2") + .into(db2_contract_table, alias="target") + .using([values, {**values, "record_id": 2, "payload": None}] if multiple_rows else values, alias="src") + .on("target.record_id = src.record_id") + .when_not_matched_then_insert({column: f"src.{column}" for column in values}) + ) + db2_session.execute(query) + db2_session.commit() + row = db2_session.select_one(f"SELECT * FROM {db2_contract_table} WHERE record_id = :record_id", record_id=1) + assert (row["record_id"], row["label"], row["enabled"]) == (1, "typed", 1) + assert Decimal(str(row["amount"])) == Decimal("12345678901234567890.123456789") + assert row["captured_at"] == datetime(2026, 6, 7, 8, 9, 10, 345678) + assert str(row["uuid_value"]) == "12345678-1234-5678-1234-567812345678" + assert row["body"] == "x" * 32673 + assert (row["day_value"], row["time_value"]) == (date(2026, 6, 7), time(12, 34, 56)) + assert bytes(row["binary_value"]) == b"\x00\xff" + assert row["ratio"] == 0.5 + payload = row["payload"] if isinstance(row["payload"], dict) else json.loads(row["payload"]) + assert payload == {"nested": "value"} + assert row["absent"] is None + if multiple_rows: + assert db2_session.select_value(f"SELECT COUNT(*) FROM {db2_contract_table}") == 2 + assert ( + db2_session.select_value( + f"SELECT payload FROM {db2_contract_table} WHERE record_id = :record_id", record_id=2 + ) + is None + ) + + +def test_db2_catalog_uniqueness_and_distinct_objects(db2_session: Db2SyncDriver, db2_contract_table: str) -> None: + """Catalog flags distinguish single-column uniqueness and composite keys without duplicated objects/indexes.""" + db2_session.execute_script( + f"CREATE TABLE {db2_contract_table} (" + "id INTEGER NOT NULL PRIMARY KEY, email VARCHAR(64) NOT NULL UNIQUE, " + "left_id INTEGER NOT NULL, right_id INTEGER NOT NULL, UNIQUE (left_id, right_id))" + ) + db2_session.commit() + dictionary = db2_session.data_dictionary + columns = { + column["column_name"]: column for column in dictionary.get_columns(db2_session, table=db2_contract_table) + } + assert columns["ID"]["is_primary"] is True + assert columns["EMAIL"]["is_primary"] is False + assert columns["EMAIL"]["is_unique"] is True + assert columns["LEFT_ID"]["is_unique"] is False + assert columns["RIGHT_ID"]["is_unique"] is False + indexes = dictionary.get_indexes(db2_session, table=db2_contract_table) + identities = [(index["schema_name"], index["table_name"], index["index_name"]) for index in indexes] + assert len(identities) == len(set(identities)) == 3 + assert sorted(tuple(index["columns"] or ()) for index in indexes) == [("EMAIL",), ("ID",), ("LEFT_ID", "RIGHT_ID")] + objects = dictionary.get_objects(db2_session).items + matching = [item for item in objects if isinstance(item, dict) and item.get("object_name") == db2_contract_table] + assert [(item["object_name"], item["object_type"]) for item in matching] == [(db2_contract_table, "TABLE")] + tables = dictionary.get_tables(db2_session) + assert [table["table_name"] for table in tables if table["table_name"] == db2_contract_table] == [ + db2_contract_table + ] + + +@pytest.fixture +def db2_contract_table(db2_session: Db2SyncDriver) -> Generator[str, None, None]: + """Reserve a fixed scratch table and clean it before and after every native residual case.""" + table = "SQLSPEC_DB2_NATIVE_CONTRACT" + _drop_contract_table(db2_session, table) + try: + yield table + finally: + db2_session.rollback() + _drop_contract_table(db2_session, table) + + +def _drop_contract_table(session: Db2SyncDriver, table: str) -> None: + exists = session.select_value( + "SELECT COUNT(*) FROM SYSCAT.TABLES WHERE TABSCHEMA = CURRENT SCHEMA AND TABNAME = :table_name", + table_name=table, + ) + if exists: + session.execute_script(f"DROP TABLE {table}") + session.commit() diff --git a/tests/integration/adapters/duckdb/duckdb/extensions/adk/test_maintenance.py b/tests/integration/adapters/duckdb/duckdb/extensions/adk/test_maintenance.py index fcc484305..8650ffddd 100644 --- a/tests/integration/adapters/duckdb/duckdb/extensions/adk/test_maintenance.py +++ b/tests/integration/adapters/duckdb/duckdb/extensions/adk/test_maintenance.py @@ -6,7 +6,7 @@ import pytest -from sqlspec.adapters.duckdb.adk import DuckdbADKMemoryStore, DuckdbADKStore +from sqlspec.adapters.duckdb.adk import DuckDBADKMemoryStore, DuckDBADKStore from sqlspec.adapters.duckdb.config import DuckDBConfig from sqlspec.extensions.adk import StoredEvent, StoredMemory from sqlspec.extensions.adk.maintenance import ( @@ -57,7 +57,7 @@ def _event(*, app_name: str, session_id: str, timestamp: datetime) -> StoredEven def test_duckdb_prune_sessions_scopes_to_app_name(tmp_path: Path) -> None: """Pruning sessions for one application leaves other applications untouched.""" config = _config(tmp_path, "prune_sessions") - store = DuckdbADKStore(config) + store = DuckDBADKStore(config) store.create_tables() stale = datetime.now(UTC) - timedelta(days=120) @@ -77,7 +77,7 @@ def test_duckdb_prune_sessions_scopes_to_app_name(tmp_path: Path) -> None: def test_duckdb_prune_events_scopes_to_app_name(tmp_path: Path) -> None: """Pruning events for one application leaves other applications' events intact.""" config = _config(tmp_path, "prune_events") - store = DuckdbADKStore(config) + store = DuckDBADKStore(config) store.create_tables() stale = datetime.now(UTC) - timedelta(days=200) @@ -95,7 +95,7 @@ def test_duckdb_prune_events_scopes_to_app_name(tmp_path: Path) -> None: def test_duckdb_prune_user_state_scopes_to_app_name(tmp_path: Path) -> None: """User state pruning honors both the idle threshold and the application filter.""" config = _config(tmp_path, "prune_user_state") - store = DuckdbADKStore(config) + store = DuckDBADKStore(config) store.create_tables() for app in ("app_keep", "app_prune"): @@ -115,7 +115,7 @@ def test_duckdb_prune_user_state_scopes_to_app_name(tmp_path: Path) -> None: def test_duckdb_prune_memory_respects_scope(tmp_path: Path) -> None: """Memory pruning removes user-scoped entries while preserving app-scoped ones.""" config = _config(tmp_path, "prune_memory") - store = DuckdbADKMemoryStore(config) + store = DuckDBADKMemoryStore(config) store.create_tables() stale = datetime.now(UTC) - timedelta(days=200) diff --git a/tests/integration/adapters/duckdb/duckdb/extensions/adk/test_memory_store.py b/tests/integration/adapters/duckdb/duckdb/extensions/adk/test_memory_store.py index 1295b3eae..1b393e8ea 100644 --- a/tests/integration/adapters/duckdb/duckdb/extensions/adk/test_memory_store.py +++ b/tests/integration/adapters/duckdb/duckdb/extensions/adk/test_memory_store.py @@ -7,7 +7,7 @@ import pytest -from sqlspec.adapters.duckdb.adk import DuckdbADKMemoryStore +from sqlspec.adapters.duckdb.adk import DuckDBADKMemoryStore from sqlspec.adapters.duckdb.config import DuckDBConfig from sqlspec.extensions.adk import StoredMemory @@ -35,20 +35,20 @@ def _build_record( ) -def _build_store(tmp_path: Path) -> DuckdbADKMemoryStore: +def _build_store(tmp_path: Path) -> DuckDBADKMemoryStore: db_path = tmp_path / "test_adk_memory.duckdb" config = DuckDBConfig(connection_config={"database": str(db_path)}) - store = DuckdbADKMemoryStore(config) + store = DuckDBADKMemoryStore(config) store.create_tables() return store -def _build_fts_store(tmp_path: Path) -> DuckdbADKMemoryStore: +def _build_fts_store(tmp_path: Path) -> DuckDBADKMemoryStore: db_path = tmp_path / "test_adk_memory_fts.duckdb" config = DuckDBConfig( connection_config={"database": str(db_path)}, extension_config={"adk": {"memory_use_fts": True}} ) - store = DuckdbADKMemoryStore(config) + store = DuckDBADKMemoryStore(config) with config.provide_connection() as conn: if not store._ensure_fts_extension(conn): # pyright: ignore[reportPrivateUsage] pytest.skip("DuckDB FTS extension is unavailable") diff --git a/tests/integration/adapters/duckdb/duckdb/extensions/adk/test_store.py b/tests/integration/adapters/duckdb/duckdb/extensions/adk/test_store.py index 1f8663e98..16e19a3a0 100644 --- a/tests/integration/adapters/duckdb/duckdb/extensions/adk/test_store.py +++ b/tests/integration/adapters/duckdb/duckdb/extensions/adk/test_store.py @@ -13,7 +13,7 @@ import pytest -from sqlspec.adapters.duckdb.adk import DuckdbADKStore +from sqlspec.adapters.duckdb.adk import DuckDBADKStore from sqlspec.adapters.duckdb.config import DuckDBConfig from sqlspec.extensions.adk import StoredEvent @@ -21,7 +21,7 @@ @pytest.fixture -def duckdb_adk_store(tmp_path: Path) -> "Generator[DuckdbADKStore, None, None]": +def duckdb_adk_store(tmp_path: Path) -> "Generator[DuckDBADKStore, None, None]": """Create DuckDB ADK store with temporary file-based database. Args: @@ -39,7 +39,7 @@ def duckdb_adk_store(tmp_path: Path) -> "Generator[DuckdbADKStore, None, None]": connection_config={"database": str(db_path)}, extension_config={"adk": {"session_table": "test_sessions", "events_table": "test_events"}}, ) - store = DuckdbADKStore(config) + store = DuckDBADKStore(config) store.create_tables() yield store finally: @@ -47,7 +47,7 @@ def duckdb_adk_store(tmp_path: Path) -> "Generator[DuckdbADKStore, None, None]": db_path.unlink() -def test_event_with_optional_fields(duckdb_adk_store: DuckdbADKStore) -> None: +def test_event_with_optional_fields(duckdb_adk_store: DuckDBADKStore) -> None: """Test creating events with optional fields stored in event_data.""" session_id = "session-008" duckdb_adk_store.create_session(session_id, "test-app", "user-008", {}) @@ -87,7 +87,7 @@ def test_event_with_optional_fields(duckdb_adk_store: DuckdbADKStore) -> None: assert event_data["turn_complete"] is False -def test_event_ordering_by_timestamp(duckdb_adk_store: DuckdbADKStore) -> None: +def test_event_ordering_by_timestamp(duckdb_adk_store: DuckDBADKStore) -> None: """Test events are ordered by timestamp ascending.""" session_id = "session-009" duckdb_adk_store.create_session(session_id, "test-app", "user-009", {}) @@ -136,7 +136,7 @@ def test_event_ordering_by_timestamp(duckdb_adk_store: DuckdbADKStore) -> None: assert event_ids == ["event-first", "event-middle", "event-last"] -def test_session_state_with_complex_data(duckdb_adk_store: DuckdbADKStore) -> None: +def test_session_state_with_complex_data(duckdb_adk_store: DuckDBADKStore) -> None: """Test session state with nested JSON structures.""" session_id = "session-complex" complex_state = { @@ -158,7 +158,7 @@ def test_session_state_with_complex_data(duckdb_adk_store: DuckdbADKStore) -> No assert session["state"]["conversation"]["turn_count"] == 5 -def test_event_data_round_trip(duckdb_adk_store: DuckdbADKStore) -> None: +def test_event_data_round_trip(duckdb_adk_store: DuckDBADKStore) -> None: """Test storing and retrieving event data via event_data.""" session_id = "session-json-rt" duckdb_adk_store.create_session(session_id, "test-app", "user-012", {}) @@ -185,7 +185,7 @@ def test_event_data_round_trip(duckdb_adk_store: DuckdbADKStore) -> None: assert events[0]["event_data"]["content"] == {"data": "value"} -def test_concurrent_session_updates(duckdb_adk_store: DuckdbADKStore) -> None: +def test_concurrent_session_updates(duckdb_adk_store: DuckDBADKStore) -> None: """Test multiple updates to same session.""" session_id = "session-concurrent" duckdb_adk_store.create_session(session_id, "test-app", "user-013", {"counter": 0}) @@ -222,7 +222,7 @@ def test_owner_id_column_with_integer(tmp_path: Path) -> None: } }, ) - store = DuckdbADKStore(config_with_extension) + store = DuckDBADKStore(config_with_extension) store.create_tables() assert store.owner_id_column_name == "tenant_id" @@ -265,7 +265,7 @@ def test_owner_id_column_with_ubigint(tmp_path: Path) -> None: } }, ) - store = DuckdbADKStore(config_with_extension) + store = DuckDBADKStore(config_with_extension) store.create_tables() assert store.owner_id_column_name == "owner_id" @@ -311,7 +311,7 @@ def test_owner_id_column_foreign_key_constraint(tmp_path: Path) -> None: } }, ) - store = DuckdbADKStore(config_with_extension) + store = DuckDBADKStore(config_with_extension) store.create_tables() store.create_session( @@ -353,7 +353,7 @@ def test_owner_id_column_without_value(tmp_path: Path) -> None: } }, ) - store = DuckdbADKStore(config_with_extension) + store = DuckDBADKStore(config_with_extension) store.create_tables() session = store.create_session( @@ -390,7 +390,7 @@ def test_owner_id_column_with_varchar(tmp_path: Path) -> None: } }, ) - store = DuckdbADKStore(config_with_extension) + store = DuckDBADKStore(config_with_extension) store.create_tables() session = store.create_session( @@ -434,7 +434,7 @@ def test_owner_id_column_multiple_sessions(tmp_path: Path) -> None: } }, ) - store = DuckdbADKStore(config_with_extension) + store = DuckDBADKStore(config_with_extension) store.create_tables() for i in range(5): @@ -477,7 +477,7 @@ def test_owner_id_column_query_by_fk(tmp_path: Path) -> None: } }, ) - store = DuckdbADKStore(config_with_extension) + store = DuckDBADKStore(config_with_extension) store.create_tables() store.create_session("s1", "app", "u1", {"val": 1}, owner_id=1) diff --git a/tests/integration/adapters/duckdb/duckdb/extensions/litestar/test_store.py b/tests/integration/adapters/duckdb/duckdb/extensions/litestar/test_store.py index cb14f3f6a..ee30bdd89 100644 --- a/tests/integration/adapters/duckdb/duckdb/extensions/litestar/test_store.py +++ b/tests/integration/adapters/duckdb/duckdb/extensions/litestar/test_store.py @@ -14,13 +14,13 @@ import pytest from sqlspec.adapters.duckdb.config import DuckDBConfig -from sqlspec.adapters.duckdb.litestar.store import DuckdbStore +from sqlspec.adapters.duckdb.litestar.store import DuckDBStore pytestmark = [pytest.mark.duckdb, pytest.mark.integration] @pytest.fixture -async def duckdb_store(tmp_path: Path) -> AsyncGenerator[DuckdbStore, None]: +async def duckdb_store(tmp_path: Path) -> AsyncGenerator[DuckDBStore, None]: """Create DuckDB store with temporary file-based database. Args: @@ -42,7 +42,7 @@ async def duckdb_store(tmp_path: Path) -> AsyncGenerator[DuckdbStore, None]: connection_config={"database": str(db_path)}, extension_config={"litestar": {"session_table": "test_sessions"}} ) try: - store = DuckdbStore(config) + store = DuckDBStore(config) await store.create_table() yield store await store.delete_all() @@ -52,7 +52,7 @@ async def duckdb_store(tmp_path: Path) -> AsyncGenerator[DuckdbStore, None]: db_path.unlink() -async def test_sync_to_thread_concurrency(duckdb_store: DuckdbStore) -> None: +async def test_sync_to_thread_concurrency(duckdb_store: DuckDBStore) -> None: """Test concurrent access via sync_to_thread wrapper. DuckDB has write serialization, so we test sequential writes diff --git a/tests/integration/adapters/oracle/oracledb/extensions/adk/test_oracle_specific.py b/tests/integration/adapters/oracle/oracledb/extensions/adk/test_oracle_specific.py index 6ff61f13a..283c89087 100644 --- a/tests/integration/adapters/oracle/oracledb/extensions/adk/test_oracle_specific.py +++ b/tests/integration/adapters/oracle/oracledb/extensions/adk/test_oracle_specific.py @@ -1,7 +1,7 @@ """Oracle-specific ADK store tests for LOB handling, JSON types, and FK columns.""" import json -from collections.abc import AsyncGenerator +from collections.abc import AsyncGenerator, Generator from datetime import UTC, datetime from typing import Any, cast from uuid import uuid4 @@ -10,8 +10,13 @@ import pytest from sqlspec.adapters.oracledb import OracleAsyncConfig, OracleSyncConfig -from sqlspec.adapters.oracledb.adk import OracleAsyncADKStore, OracleSyncADKStore -from sqlspec.extensions.adk import StoredEvent +from sqlspec.adapters.oracledb.adk import ( + OracleAsyncADKMemoryStore, + OracleAsyncADKStore, + OracleSyncADKMemoryStore, + OracleSyncADKStore, +) +from sqlspec.extensions.adk import StoredEvent, StoredMemory pytestmark = [pytest.mark.xdist_group("oracle"), pytest.mark.oracledb, pytest.mark.integration] @@ -485,3 +490,123 @@ async def test_create_session_with_owner_id_sync(oracle_store_sync_with_fk: "Ora retrieved = oracle_store_sync_with_fk.get_session(app_name, user_id, session_id) assert retrieved is not None assert retrieved["id"] == session_id + + +async def _cleanup_async_memory_store(store: "OracleAsyncADKMemoryStore", config: "OracleAsyncConfig") -> None: + """Drop ADK memory tables for async stores.""" + async with config.provide_connection() as conn: + cursor = conn.cursor() + for stmt in store._drop_memory_table_sql(): + try: + await cursor.execute(stmt) + except Exception: + pass + await conn.commit() + + +def _cleanup_sync_memory_store(store: "OracleSyncADKMemoryStore", config: "OracleSyncConfig") -> None: + """Drop ADK memory tables for sync stores.""" + with config.provide_connection() as conn: + cursor = conn.cursor() + for stmt in store._drop_memory_table_sql(): + try: + cursor.execute(stmt) + except Exception: + pass + conn.commit() + + +@pytest.fixture +async def oracle_async_memory_store( + oracle_async_config: "OracleAsyncConfig", +) -> "AsyncGenerator[OracleAsyncADKMemoryStore, None]": + """Create an async Oracle ADK memory store with tables created per test.""" + store = OracleAsyncADKMemoryStore(oracle_async_config) + await store.create_tables() + try: + yield store + finally: + await _cleanup_async_memory_store(store, oracle_async_config) + + +@pytest.fixture +def oracle_sync_memory_store( + oracle_sync_config: "OracleSyncConfig", +) -> "Generator[OracleSyncADKMemoryStore, None, None]": + """Create a sync Oracle ADK memory store with tables created per test.""" + store = OracleSyncADKMemoryStore(oracle_sync_config) + store.create_tables() + try: + yield store + finally: + _cleanup_sync_memory_store(store, oracle_sync_config) + + +async def test_oracle_async_adk_memory_store_insert_and_search_embedding( + oracle_async_memory_store: "OracleAsyncADKMemoryStore", +) -> None: + """Async memory store inserts embedding vector and executes vector search.""" + app_name = "test-app" + user_id = "user-vec" + now = datetime.now(UTC) + embedding = [0.1 * i for i in range(768)] + entries: list[StoredMemory] = [ + { + "id": _unique_session_id("mem-vec"), + "session_id": _unique_session_id("sess-vec"), + "app_name": app_name, + "user_id": user_id, + "scope": "user", + "event_id": _unique_session_id("ev-vec"), + "author": "user", + "timestamp": now, + "content_json": {"query": "vector search test"}, + "content_text": "vector search content", + "metadata_json": {"tag": "embedding"}, + "inserted_at": now, + "embedding": embedding, + } + ] + inserted = await oracle_async_memory_store.insert_memory_entries(entries) + assert inserted == 1 + + results = await oracle_async_memory_store.search_entries( + query="vector search", app_name=app_name, user_id=user_id, embedding=embedding + ) + assert len(results) >= 1 + assert results[0]["embedding"] is not None + + +def test_oracle_sync_adk_memory_store_insert_and_search_embedding( + oracle_sync_memory_store: "OracleSyncADKMemoryStore", +) -> None: + """Sync memory store inserts embedding vector and executes vector search.""" + app_name = "test-app" + user_id = "user-vec-sync" + now = datetime.now(UTC) + embedding = [0.1 * i for i in range(768)] + entries: list[StoredMemory] = [ + { + "id": _unique_session_id("mem-vec-sync"), + "session_id": _unique_session_id("sess-vec-sync"), + "app_name": app_name, + "user_id": user_id, + "scope": "user", + "event_id": _unique_session_id("ev-vec-sync"), + "author": "user", + "timestamp": now, + "content_json": {"query": "sync vector search test"}, + "content_text": "sync vector search content", + "metadata_json": {"tag": "embedding"}, + "inserted_at": now, + "embedding": embedding, + } + ] + inserted = oracle_sync_memory_store.insert_memory_entries(entries) + assert inserted == 1 + + results = oracle_sync_memory_store.search_entries( + query="sync vector search", app_name=app_name, user_id=user_id, embedding=embedding + ) + assert len(results) >= 1 + assert results[0]["embedding"] is not None diff --git a/tests/integration/adapters/oracle/oracledb/test_features.py b/tests/integration/adapters/oracle/oracledb/test_features.py index c8d0355a0..fba498148 100644 --- a/tests/integration/adapters/oracle/oracledb/test_features.py +++ b/tests/integration/adapters/oracle/oracledb/test_features.py @@ -5,8 +5,9 @@ import pytest -from sqlspec.adapters.oracledb import OracleAsyncDriver, OracleSyncDriver +from sqlspec.adapters.oracledb import OracleAsyncDriver, OracleSyncDriver, OracleVersionInfo from sqlspec.core import SQL, SQLResult, StatementConfig +from sqlspec.data_dictionary.dialects.oracle import parse_oracle_version_components def _lower_keys(row: dict[str, object]) -> dict[str, object]: @@ -16,6 +17,24 @@ def _lower_keys(row: dict[str, object]) -> dict[str, object]: pytestmark = pytest.mark.xdist_group("oracle") +def test_oracle_release_metadata_matches_native_vector_capable_server(oracle_sync_session: OracleSyncDriver) -> None: + """The 23ai fixture reports its release update and actual vector-compatible level.""" + native_version = str( + oracle_sync_session.select_value( + "SELECT COALESCE(version_full, version) FROM product_component_version " + "WHERE product LIKE 'Oracle%' FETCH FIRST 1 ROWS ONLY" + ) + ) + native_compatible = str(oracle_sync_session.select_value("SELECT value FROM v$parameter WHERE name = 'compatible'")) + version = oracle_sync_session.data_dictionary.get_version(oracle_sync_session) + assert isinstance(version, OracleVersionInfo) + assert (version.major, version.minor, version.patch) == parse_oracle_version_components(native_version) + assert version.compatible == native_compatible + compatible = parse_oracle_version_components(native_compatible) + assert isinstance(compatible, tuple), f"Invalid native COMPATIBLE: {native_compatible!r}" + assert compatible >= (23, 4, 0), f"Vector fixture COMPATIBLE: {native_compatible!r}" + + def test_sync_oracle_data_types(oracle_sync_session: OracleSyncDriver) -> None: """Test Oracle-specific data types (NUMBER, VARCHAR2, CLOB, DATE).""" @@ -226,3 +245,28 @@ async def test_async_oracle_exception_handling(oracle_async_session: OracleAsync await oracle_async_session.execute_script( "BEGIN EXECUTE IMMEDIATE 'DROP TABLE test_exception_table_oracledb_async'; EXCEPTION WHEN OTHERS THEN IF SQLCODE != -942 THEN RAISE; END IF; END;" ) + + +def test_oracle_get_indexes_function_based_index(oracle_sync_session: OracleSyncDriver) -> None: + """Test get_indexes returns index and expression for function-based indexes.""" + oracle_sync_session.execute_script( + "BEGIN EXECUTE IMMEDIATE 'DROP TABLE test_fn_idx_table'; EXCEPTION WHEN OTHERS THEN IF SQLCODE != -942 THEN RAISE; END IF; END;" + ) + oracle_sync_session.execute_script(""" + CREATE TABLE test_fn_idx_table ( + id NUMBER(10) PRIMARY KEY, + email VARCHAR2(255) + ) + """) + oracle_sync_session.execute_script(""" + CREATE INDEX idx_test_fn_email ON test_fn_idx_table (UPPER(email)) + """) + + indexes = oracle_sync_session.data_dictionary.get_indexes(oracle_sync_session, table="test_fn_idx_table") + fn_indexes = [idx for idx in indexes if idx.get("index_name") == "IDX_TEST_FN_EMAIL"] + assert len(fn_indexes) == 1 + assert fn_indexes[0].get("expressions") is not None or fn_indexes[0].get("expression") is not None + + oracle_sync_session.execute_script( + "BEGIN EXECUTE IMMEDIATE 'DROP TABLE test_fn_idx_table'; EXCEPTION WHEN OTHERS THEN IF SQLCODE != -942 THEN RAISE; END IF; END;" + ) diff --git a/tests/integration/adapters/postgres/adbc/extensions/litestar/test_store.py b/tests/integration/adapters/postgres/adbc/extensions/litestar/test_store.py index 1c1eea697..1f61784fd 100644 --- a/tests/integration/adapters/postgres/adbc/extensions/litestar/test_store.py +++ b/tests/integration/adapters/postgres/adbc/extensions/litestar/test_store.py @@ -13,16 +13,16 @@ import pytest from sqlspec.adapters.adbc import AdbcConfig -from sqlspec.adapters.adbc.litestar import ADBCStore +from sqlspec.adapters.adbc.litestar import AdbcStore pytestmark = [pytest.mark.xdist_group("postgres"), pytest.mark.adbc, pytest.mark.integration] @pytest.fixture -async def adbc_store(adbc_postgres_config: AdbcConfig) -> AsyncGenerator[ADBCStore, None]: +async def adbc_store(adbc_postgres_config: AdbcConfig) -> AsyncGenerator[AdbcStore, None]: """Create an ADBC store with PostgreSQL backend.""" adbc_postgres_config.extension_config = {"litestar": {"session_table": "test_adbc_sessions"}} - store = ADBCStore(adbc_postgres_config) + store = AdbcStore(adbc_postgres_config) await store.create_table() try: yield store @@ -31,7 +31,7 @@ async def adbc_store(adbc_postgres_config: AdbcConfig) -> AsyncGenerator[ADBCSto await store.delete_all() -async def test_sync_to_thread_concurrency(adbc_store: ADBCStore) -> None: +async def test_sync_to_thread_concurrency(adbc_store: AdbcStore) -> None: """ADBC PostgreSQL supports concurrent store reads and writes through sync_to_thread.""" async def write_session(session_id: int) -> None: diff --git a/tests/integration/adapters/postgres/adbc/test_driver.py b/tests/integration/adapters/postgres/adbc/test_driver.py index 38e46dfad..8c58c9eac 100644 --- a/tests/integration/adapters/postgres/adbc/test_driver.py +++ b/tests/integration/adapters/postgres/adbc/test_driver.py @@ -1,11 +1,16 @@ """PostgreSQL-backed ADBC driver residuals.""" +from collections.abc import Generator +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from threading import Barrier +from typing import TYPE_CHECKING, Literal, cast from uuid import UUID, uuid4 import pytest -from sqlspec.adapters.adbc import AdbcDriver -from sqlspec.exceptions import SQLSpecError +from sqlspec.adapters.adbc import AdbcConfig, AdbcDriver +from sqlspec.exceptions import ConnectionTimeoutError, DatabaseConnectionError, SQLSpecError from tests.integration.adapters._shared.adbc_backends import postgresql_session, test_postgresql_specific_features from tests.integration.adapters._shared.adbc_connection import ( test_connection, @@ -18,6 +23,9 @@ test_execute_script_edge_cases, ) +if TYPE_CHECKING: + from pyarrow import RecordBatch + __all__ = ( "postgresql_session", "test_adbc_postgresql_statement_stack_continue_on_error", @@ -223,3 +231,235 @@ def test_postgresql_uuid_row_stream_and_arrow_boundary(postgresql_session: AdbcD assert arrow_type.extension_name == "arrow.opaque" assert arrow_type.type_name == "uuid" assert arrow_table.to_pylist() == [{"value": value.bytes}] + + +@pytest.mark.xdist_group("postgres") +@pytest.mark.adbc +def test_postgresql_pool_reuses_pid_and_preserves_transactional_wrapper( + adbc_postgres_connection_config: dict[str, str], +) -> None: + """A reset lease keeps its backend and still supports real rollback and commit.""" + with _pooled_postgresql_config(adbc_postgres_connection_config) as config: + with config.provide_session() as session: + first_pid = session.select_value("SELECT pg_backend_pid()") + session.execute_script("CREATE TEMP TABLE adbc_pool_rollback (value INTEGER)") + session.commit() + session.execute("INSERT INTO adbc_pool_rollback VALUES (?)", 7) + session.rollback() + assert session.select_value("SELECT COUNT(*) FROM adbc_pool_rollback") == 0 + with config.provide_session() as session: + assert session.select_value("SELECT pg_backend_pid()") == first_pid + session.execute_script("CREATE TEMP TABLE adbc_pool_after_reset (value INTEGER)") + session.commit() + session.execute("INSERT INTO adbc_pool_after_reset VALUES (?)", 11) + session.rollback() + assert session.select_value("SELECT COUNT(*) FROM adbc_pool_after_reset") == 0 + session.execute("INSERT INTO adbc_pool_after_reset VALUES (?)", 13) + session.commit() + assert session.select_value("SELECT value FROM adbc_pool_after_reset") == 13 + + +@pytest.mark.xdist_group("postgres") +@pytest.mark.adbc +def test_postgresql_pool_simultaneous_leases_are_exclusive(adbc_postgres_connection_config: dict[str, str]) -> None: + """Overlapping native operations hold distinct backend connections.""" + barrier = Barrier(2, timeout=30) + with _pooled_postgresql_config(adbc_postgres_connection_config, max_size=2) as config: + + def borrower(value: int) -> tuple[int, int]: + with config.provide_session() as session: + pid = int(session.select_value("SELECT pg_backend_pid()")) + barrier.wait() + observed = int(session.select_value("SELECT CAST(? AS INTEGER)", value)) + barrier.wait() + return pid, observed + + with ThreadPoolExecutor(max_workers=2) as executor: + results = list(executor.map(borrower, [3, 5])) + assert len({pid for pid, _value in results}) == 2 + assert [value for _pid, value in results] == [3, 5] + + +@pytest.mark.xdist_group("postgres") +@pytest.mark.adbc +def test_postgresql_pool_nested_checkout_times_out_without_stealing_outer_lease( + adbc_postgres_connection_config: dict[str, str], +) -> None: + """A saturated nested checkout refuses while the outer backend stays usable.""" + with _pooled_postgresql_config(adbc_postgres_connection_config) as config, config.provide_session() as outer: + pid = outer.select_value("SELECT pg_backend_pid()") + with pytest.raises(ConnectionTimeoutError), config.provide_session(): + pytest.fail("nested checkout bypassed the configured capacity") + assert outer.select_value("SELECT pg_backend_pid()") == pid + + +@pytest.mark.xdist_group("postgres") +@pytest.mark.adbc +def test_postgresql_pool_failed_body_evicts_native_connection(adbc_postgres_connection_config: dict[str, str]) -> None: + """A failed transaction is not reused even after driver-level rollback.""" + with _pooled_postgresql_config(adbc_postgres_connection_config) as config: + with config.provide_session() as session: + failed_pid = session.select_value("SELECT pg_backend_pid()") + with pytest.raises(SQLSpecError), config.provide_session() as session: + session.execute("SELECT 1 / CAST(? AS INTEGER)", 0) + with config.provide_session() as session: + assert session.select_value("SELECT pg_backend_pid()") != failed_pid + assert session.select_value("SELECT CAST(? AS INTEGER)", 17) == 17 + + +@pytest.mark.xdist_group("postgres") +@pytest.mark.adbc +def test_postgresql_pool_active_shutdown_drains_before_new_generation( + adbc_postgres_connection_config: dict[str, str], +) -> None: + """Shutdown preserves an active borrower and refuses checkout until it exits.""" + with _pooled_postgresql_config(adbc_postgres_connection_config) as config: + with config.provide_session() as session: + old_pid = session.select_value("SELECT pg_backend_pid()") + config.close_pool() + assert session.select_value("SELECT pg_backend_pid()") == old_pid + with pytest.raises(DatabaseConnectionError), config.provide_session(): + pytest.fail("a draining generation accepted a new borrower") + with config.provide_session() as session: + assert session.select_value("SELECT pg_backend_pid()") != old_pid + assert session.select_value("SELECT CAST(? AS INTEGER)", 19) == 19 + config.close_pool() + config.close_pool() + + +@pytest.mark.xdist_group("postgres") +@pytest.mark.adbc +def test_postgresql_pool_discards_committed_session_state(adbc_postgres_connection_config: dict[str, str]) -> None: + """Reuse clears committed session state that a rollback alone cannot undo.""" + with _pooled_postgresql_config(adbc_postgres_connection_config) as config: + with config.provide_session() as session: + pid = session.select_value("SELECT pg_backend_pid()") + search_path = session.select_value("SELECT current_setting('search_path')") + session.execute_script("CREATE TEMP TABLE adbc_pool_reset (value INTEGER)") + session.select_value("SELECT set_config('search_path', ?, false)", "pg_catalog") + session.execute_script("LISTEN adbc_pool_notifications") + session.select_value("SELECT pg_advisory_lock(CAST(? AS BIGINT))", 728311) + session.execute_script("PREPARE adbc_pool_prepared AS SELECT 1") + session.commit() + with config.provide_session() as session: + assert session.select_value("SELECT pg_backend_pid()") == pid + assert session.select_value("SELECT current_setting('search_path')") == search_path + assert session.select_value("SELECT to_regclass(?)", "pg_temp.adbc_pool_reset") is None + assert session.select_value("SELECT COUNT(*) FROM pg_listening_channels()") == 0 + assert ( + session.select_value( + "SELECT COUNT(*) FROM pg_locks WHERE locktype = ? AND pid = pg_backend_pid()", "advisory" + ) + == 0 + ) + assert ( + session.select_value("SELECT COUNT(*) FROM pg_prepared_statements WHERE name = ?", "adbc_pool_prepared") + == 0 + ) + + +@pytest.mark.xdist_group("postgres") +@pytest.mark.adbc +def test_postgresql_pool_clones_accept_initial_connection_options( + adbc_postgres_connection_config: dict[str, str], +) -> None: + """Fresh clones support explicit native autocommit=False initialization.""" + with ( + _pooled_postgresql_config( + adbc_postgres_connection_config, conn_kwargs={"adbc.connection.autocommit": "false"}, max_size=2 + ) as config, + config.provide_session() as first, + config.provide_session() as second, + ): + first_pid = first.select_value("SELECT pg_backend_pid()") + assert second.select_value("SELECT pg_backend_pid()") != first_pid + assert first.select_value("SELECT CAST(? AS INTEGER)", 29) == 29 + assert second.select_value("SELECT CAST(? AS INTEGER)", 31) == 31 + + +@pytest.mark.xdist_group("postgres") +@pytest.mark.adbc +@pytest.mark.parametrize("kind", ["rows", "arrow"]) +@pytest.mark.parametrize("consumption", ["never", "partial", "exhausted"]) +def test_postgresql_pool_lazy_streams_evict_after_every_consumption_mode( + adbc_postgres_connection_config: dict[str, str], kind: str, consumption: str +) -> None: + """A lazy stream prevents reuse even when never started or fully exhausted.""" + with _pooled_postgresql_config(adbc_postgres_connection_config) as config: + with config.provide_session() as session: + pid = session.select_value("SELECT pg_backend_pid()") + statement = "SELECT generate_series(1, 3) AS value" + stream = ( + session.select_stream(statement, native_only=True, chunk_size=1) + if kind == "rows" + else session.select_to_arrow(statement, return_format="reader").data + ) + iterator = iter(stream) + if consumption == "partial": + next(iterator) + elif consumption == "exhausted": + parts = list(iterator) + if kind == "rows": + assert parts == [{"value": 1}, {"value": 2}, {"value": 3}] + else: + assert [ + row["value"] for batch in cast("list[RecordBatch]", parts) for row in batch.to_pylist() + ] == [1, 2, 3] + with config.provide_session() as session: + assert session.select_value("SELECT pg_backend_pid()") != pid + assert session.select_value("SELECT CAST(? AS INTEGER)", 23) == 23 + + +@pytest.mark.xdist_group("postgres") +@pytest.mark.adbc +@pytest.mark.parametrize("return_format", ["table", "batches"]) +def test_postgresql_pool_eager_arrow_results_survive_session_and_allow_reuse( + adbc_postgres_connection_config: dict[str, str], return_format: Literal["table", "batches"] +) -> None: + """Eager Arrow data remains owned by the caller while its connection is reused.""" + with _pooled_postgresql_config(adbc_postgres_connection_config) as config: + with config.provide_session() as session: + pid = session.select_value("SELECT pg_backend_pid()") + data = session.select_to_arrow("SELECT generate_series(1, 3) AS value", return_format=return_format).data + with config.provide_session() as session: + assert session.select_value("SELECT pg_backend_pid()") == pid + rows = data.to_pylist() if return_format == "table" else [row for batch in data for row in batch.to_pylist()] + assert rows == [{"value": 1}, {"value": 2}, {"value": 3}] + + +@pytest.mark.xdist_group("postgres") +@pytest.mark.adbc +def test_postgresql_autocommit_transaction_boundaries(adbc_postgres_connection_config: dict[str, str]) -> None: + """Explicit transactions commit or roll back native work and restore autocommit.""" + config = AdbcConfig(connection_config={**adbc_postgres_connection_config, "autocommit": True}) + with config.provide_session() as session: + session.execute("CREATE TEMP TABLE sqlspec_adbc_autocommit_transactions (value INTEGER)") + session.begin() + session.execute("INSERT INTO sqlspec_adbc_autocommit_transactions (value) VALUES (?)", 41) + session.rollback() + assert session.select_value("SELECT COUNT(*) FROM sqlspec_adbc_autocommit_transactions") == 0 + assert session.connection.adbc_connection.get_option("adbc.connection.autocommit") == "true" + + session.begin() + session.execute("INSERT INTO sqlspec_adbc_autocommit_transactions (value) VALUES (?)", 42) + session.commit() + assert session.select_value("SELECT value FROM sqlspec_adbc_autocommit_transactions") == 42 + assert session.connection.adbc_connection.get_option("adbc.connection.autocommit") == "true" + + +@contextmanager +def _pooled_postgresql_config( + connection_config: dict[str, str], *, max_size: int = 1, conn_kwargs: dict[str, str] | None = None +) -> Generator[AdbcConfig, None, None]: + config = AdbcConfig( + connection_config={ + "driver_name": "adbc_driver_postgresql", + **connection_config, + "conn_kwargs": conn_kwargs or {}, + }, + pool_config={"max_size": max_size, "timeout": 0.05}, + ) + try: + yield config + finally: + config.close_pool() diff --git a/tests/integration/adapters/postgres/asyncpg/test_driver.py b/tests/integration/adapters/postgres/asyncpg/test_driver.py index 31239e022..e236bac6d 100644 --- a/tests/integration/adapters/postgres/asyncpg/test_driver.py +++ b/tests/integration/adapters/postgres/asyncpg/test_driver.py @@ -483,12 +483,18 @@ async def test_asyncpg_statement_stack_continue_on_error(asyncpg_session: "Async assert verify.get_data()[0]["total"] == 2 -async def test_asyncpg_statement_stack_continue_on_error_inside_transaction(asyncpg_session: "AsyncpgDriver") -> None: +@pytest.mark.parametrize("raw_begin", [False, True]) +async def test_asyncpg_statement_stack_continue_on_error_inside_transaction( + asyncpg_session: "AsyncpgDriver", raw_begin: bool +) -> None: """Continue-on-error inside a user transaction must isolate the failing op via a savepoint.""" await asyncpg_session.execute_script("DELETE FROM test_table_asyncpg") - await asyncpg_session.begin() + if raw_begin: + await asyncpg_session.execute("BEGIN") + else: + await asyncpg_session.begin() stack = ( StatementStack() .push_execute("INSERT INTO test_table_asyncpg (id, name, value) VALUES ($1, $2, $3)", (1, "initial", 5)) @@ -501,6 +507,7 @@ async def test_asyncpg_statement_stack_continue_on_error_inside_transaction(asyn assert results[0].error is None assert results[1].error is not None assert results[2].error is None + assert asyncpg_session.connection.is_in_transaction() verify = await asyncpg_session.execute("SELECT COUNT(*) AS total FROM test_table_asyncpg") assert verify.get_data()[0]["total"] == 2 diff --git a/tests/integration/adapters/postgres/psqlpy/extensions/adk/test_memory_embeddings.py b/tests/integration/adapters/postgres/psqlpy/extensions/adk/test_memory_embeddings.py new file mode 100644 index 000000000..a7fcd6c6b --- /dev/null +++ b/tests/integration/adapters/postgres/psqlpy/extensions/adk/test_memory_embeddings.py @@ -0,0 +1,63 @@ +"""Native vector preservation during batched psqlpy memory insertion.""" + +from datetime import UTC, datetime +from typing import TYPE_CHECKING +from uuid import uuid4 + +import pytest + +from sqlspec.adapters.psqlpy import PsqlpyConfig +from sqlspec.adapters.psqlpy.adk import PsqlpyADKMemoryStore + +if TYPE_CHECKING: + from sqlspec.extensions.adk import StoredMemory + +pytestmark = pytest.mark.xdist_group("pgvector") + + +@pytest.mark.parametrize("with_owner", [False, True]) +async def test_bulk_memory_insertion_preserves_embeddings( + pgvector_config_psqlpy: PsqlpyConfig, with_owner: bool +) -> None: + """Vectors and nulls remain aligned with their rows, including owner-scoped batches.""" + config = PsqlpyConfig( + connection_config=dict(pgvector_config_psqlpy.connection_config), + extension_config={ + "adk": { + "memory_table": f"memory_vectors_{uuid4().hex}", + "vector_dimensions": 2, + "vector_index_type": "none", + "owner_id_column": "tenant_id VARCHAR(128)" if with_owner else None, + } + }, + ) + store = PsqlpyADKMemoryStore(config) + now = datetime.now(UTC) + entries: list[StoredMemory] = [ + { + "id": f"memory-{index}", + "session_id": "session", + "app_name": "app", + "user_id": "user", + "scope": "user", + "event_id": f"event-{index}", + "author": "user", + "timestamp": now, + "content_json": {"text": "memory"}, + "content_text": "memory", + "metadata_json": None, + "inserted_at": now, + "embedding": embedding, + } + for index, embedding in enumerate(((1.0, 0.0), None, (0.0, 1.0))) + ] + try: + await store.create_tables() + assert await store.insert_memory_entries(entries, owner_id="tenant") == 3 + assert await store.insert_memory_entries(entries, owner_id="tenant") == 0 + results = await store.search_entries("", "app", "user", embedding=(1.0, 0.0)) + assert [row["id"] for row in results] == ["memory-0", "memory-2"] + assert [row["embedding"] for row in results] == [[1.0, 0.0], [0.0, 1.0]] + finally: + await store.drop_tables() + await config.close_pool() diff --git a/tests/integration/adapters/postgres/psqlpy/test_driver.py b/tests/integration/adapters/postgres/psqlpy/test_driver.py index 129f538d6..9f3a16287 100644 --- a/tests/integration/adapters/postgres/psqlpy/test_driver.py +++ b/tests/integration/adapters/postgres/psqlpy/test_driver.py @@ -4,6 +4,7 @@ from sqlspec import SQLResult, StatementStack from sqlspec.adapters.psqlpy import PsqlpyConfig, PsqlpyDriver +from sqlspec.exceptions import SQLSpecError pytestmark = pytest.mark.xdist_group("postgres") @@ -244,3 +245,17 @@ async def test_extensions_not_enabled_on_standard_postgres(psqlpy_config: "Psqlp assert psqlpy_config._pgvector_available is False # pyright: ignore[reportPrivateUsage] assert psqlpy_config._paradedb_available is False # pyright: ignore[reportPrivateUsage] assert psqlpy_config.statement_config.dialect == "postgres" + + +async def test_psqlpy_cached_dml_and_errors(psqlpy_session: "PsqlpyDriver") -> None: + """Native psqlpy cache hits preserve counts and SQLSpec exception translation.""" + from psqlpy.exceptions import Error + + for name in ("cached-first", "cached-second"): + result = await psqlpy_session.execute("INSERT INTO test_table_psqlpy (name) VALUES (?)", name) + assert result.rows_affected == 1 + query = "SELECT 10 / ?::integer AS value" + assert (await psqlpy_session.execute(query, 2)).get_data() == [{"value": 5}] + with pytest.raises(SQLSpecError) as error: + await psqlpy_session.execute(query, 0) + assert isinstance(error.value.__cause__, Error) diff --git a/tests/integration/adapters/postgres/psycopg/test_execution_regressions.py b/tests/integration/adapters/postgres/psycopg/test_execution_regressions.py new file mode 100644 index 000000000..7dfc38eb6 --- /dev/null +++ b/tests/integration/adapters/postgres/psycopg/test_execution_regressions.py @@ -0,0 +1,103 @@ +"""PostgreSQL execution contracts across cached and native COPY paths.""" + +import inspect +from typing import TYPE_CHECKING, Any + +import pytest + +from sqlspec.adapters.psycopg import PsycopgAsyncConfig, PsycopgSyncConfig +from sqlspec.exceptions import UniqueViolationError + +if TYPE_CHECKING: + from collections.abc import AsyncIterator + + from sqlspec.adapters.psycopg import PsycopgAsyncDriver, PsycopgSyncDriver + from sqlspec.core import SQLResult + +pytestmark = pytest.mark.xdist_group("postgres") + + +@pytest.fixture(params=["sync", "async"]) +async def psycopg_regression_driver( + request: pytest.FixtureRequest, psycopg_sync_config: PsycopgSyncConfig, psycopg_async_config: PsycopgAsyncConfig +) -> "AsyncIterator[PsycopgSyncDriver | PsycopgAsyncDriver]": + if request.param == "sync": + with psycopg_sync_config.provide_session() as driver: + yield driver + else: + async with psycopg_async_config.provide_session() as async_driver: + yield async_driver + + +async def _execute(driver: "PsycopgSyncDriver | PsycopgAsyncDriver", query: str, *parameters: Any) -> "SQLResult": + result = driver.execute(query, *parameters) + return await result if inspect.isawaitable(result) else result + + +@pytest.mark.parametrize( + ("query", "parameters", "expected"), + [ + ("SELECT '100%' AS value", (), "100%"), + ("SELECT 10 % 3 AS value", (), 1), + ("SELECT ?::integer % 3 AS value", (10,), 1), + ], +) +async def test_psycopg_cached_percent_expressions( + psycopg_regression_driver: "PsycopgSyncDriver | PsycopgAsyncDriver", + query: str, + parameters: tuple[int, ...], + expected: str | int, +) -> None: + """Native parameter parsing gives equal values before and after a cache hit.""" + for _ in range(2): + result = await _execute(psycopg_regression_driver, query, *parameters) + assert result.get_data() == [{"value": expected}] + + +async def test_psycopg_copy_constraint_error_is_mapped( + psycopg_regression_driver: "PsycopgSyncDriver | PsycopgAsyncDriver", +) -> None: + """The constraint error received when COPY finishes is translated consistently.""" + import psycopg + + pa = pytest.importorskip("pyarrow") + driver = psycopg_regression_driver + await _execute(driver, "CREATE TEMP TABLE copy_constraint_probe (id INTEGER PRIMARY KEY) ON COMMIT DROP") + with pytest.raises(UniqueViolationError) as error: + job = driver.load_from_arrow("copy_constraint_probe", pa.table({"id": [1, 1]})) + if inspect.isawaitable(job): + await job + assert isinstance(error.value.__cause__, psycopg.errors.UniqueViolation) + + +@pytest.mark.parametrize("is_async", [False, True]) +async def test_psycopg_sql_connection_hook_leaves_pool_usable( + is_async: bool, psycopg_sync_config: PsycopgSyncConfig +) -> None: + """SQL issued by a connection hook cannot leave the pool's configure callback in INTRANS.""" + connection_config = {**psycopg_sync_config.connection_config, "min_size": 1, "max_size": 1, "timeout": 2} + if is_async: + + async def async_hook(connection: Any) -> None: + await connection.execute("SET timezone TO 'Pacific/Honolulu'") + + config = PsycopgAsyncConfig( + connection_config=connection_config, driver_features={"on_connection_create": async_hook} + ) + try: + async with config.provide_session() as driver: + assert await driver.select_value("SELECT current_setting('TimeZone')") == "Pacific/Honolulu" + finally: + await config.close_pool() + else: + sync_config = PsycopgSyncConfig( + connection_config=connection_config, + driver_features={ + "on_connection_create": lambda connection: connection.execute("SET timezone TO 'Pacific/Honolulu'") + }, + ) + try: + with sync_config.provide_session() as sync_driver: + assert sync_driver.select_value("SELECT current_setting('TimeZone')") == "Pacific/Honolulu" + finally: + sync_config.close_pool() diff --git a/tests/integration/adapters/postgres/test_column_metadata.py b/tests/integration/adapters/postgres/test_column_metadata.py index 1432ef58c..f5c534370 100644 --- a/tests/integration/adapters/postgres/test_column_metadata.py +++ b/tests/integration/adapters/postgres/test_column_metadata.py @@ -70,6 +70,28 @@ async def test_unqualified_columns_resolve_through_search_path(asyncpg_async_dri await driver.execute_script(f"RESET search_path; {TEARDOWN_SQL}") +async def test_unqualified_indexes_resolve_through_search_path(asyncpg_async_driver: "AsyncpgDriver") -> None: + """Index lookup resolves the table on search_path and honors an explicit schema.""" + driver = asyncpg_async_driver + await driver.execute_script(SETUP_SQL) + try: + await driver.execute_script( + 'CREATE INDEX "MixedLookup" ON column_metadata."MixedItems" ("userName"); ' + "SET search_path TO public, column_metadata" + ) + + indexes = await driver.data_dictionary.get_indexes(driver, table='"MixedItems"') + explicit_public = await driver.data_dictionary.get_indexes(driver, table='"MixedItems"', schema="public") + + lookup = next(index for index in indexes if index["index_name"] == "MixedLookup") + assert lookup["schema_name"] == "column_metadata" + assert lookup["table_name"] == "MixedItems" + assert lookup["columns"] == ["userName"] + assert explicit_public == [] + finally: + await driver.execute_script(f"RESET search_path; {TEARDOWN_SQL}") + + def test_columns_report_keys_identity_and_sequences_sync(psycopg_sync_config: "PsycopgSyncConfig") -> None: """The sync dictionary reports the same column metadata.""" with psycopg_sync_config.provide_session() as driver: diff --git a/tests/integration/adapters/sqlanywhere/__init__.py b/tests/integration/adapters/sqlanywhere/__init__.py new file mode 100644 index 000000000..05cb6dbac --- /dev/null +++ b/tests/integration/adapters/sqlanywhere/__init__.py @@ -0,0 +1 @@ +"""SAP SQL Anywhere integration tests.""" diff --git a/tests/integration/adapters/sqlanywhere/test_driver.py b/tests/integration/adapters/sqlanywhere/test_driver.py new file mode 100644 index 000000000..166443c22 --- /dev/null +++ b/tests/integration/adapters/sqlanywhere/test_driver.py @@ -0,0 +1,99 @@ +"""SAP SQL Anywhere driver round trips, transactions, and exception mapping.""" + +from collections.abc import Generator + +import pytest + +pytest.importorskip("sqlanydb") + +from sqlspec.adapters.sqlanywhere import SqlanywhereConfig, SqlanywhereDriver +from sqlspec.exceptions import NotNullViolationError, UniqueViolationError + +pytestmark = [pytest.mark.sqlanywhere, pytest.mark.xdist_group("sqlanywhere")] + +_TABLE = "sqlspec_sa_items" + + +@pytest.fixture +def items_table(sqlanywhere_session: SqlanywhereDriver) -> Generator[str, None, None]: + """Create a fresh items table and drop it after the test.""" + sqlanywhere_session.execute_script(f"DROP TABLE IF EXISTS {_TABLE}") + sqlanywhere_session.execute_script( + f"CREATE TABLE {_TABLE} (id INTEGER NOT NULL PRIMARY KEY, name VARCHAR(50) NOT NULL, note LONG VARCHAR NULL)" + ) + sqlanywhere_session.commit() + try: + yield _TABLE + finally: + sqlanywhere_session.rollback() + sqlanywhere_session.execute_script(f"DROP TABLE IF EXISTS {_TABLE}") + sqlanywhere_session.commit() + + +def test_select_scalar(sqlanywhere_session: SqlanywhereDriver) -> None: + assert sqlanywhere_session.select_value("SELECT 1") == 1 + + +def test_parameterized_round_trip(sqlanywhere_session: SqlanywhereDriver, items_table: str) -> None: + result = sqlanywhere_session.execute( + f"INSERT INTO {items_table} (id, name, note) VALUES (?, ?, ?)", (1, "alpha", "x" * 5000) + ) + assert result.rows_affected == 1 + sqlanywhere_session.execute_many( + f"INSERT INTO {items_table} (id, name, note) VALUES (?, ?, ?)", [(2, "beta", None), (3, "gamma", "short")] + ) + sqlanywhere_session.commit() + + rows = sqlanywhere_session.select(f"SELECT id, name, note FROM {items_table} ORDER BY id") + assert [row["name"] for row in rows] == ["alpha", "beta", "gamma"] + assert rows[0]["note"] == "x" * 5000 + assert rows[1]["note"] is None + assert sqlanywhere_session.select_value(f"SELECT name FROM {items_table} WHERE id = ?", (3,)) == "gamma" + + +def test_rollback_discards_uncommitted_rows(sqlanywhere_session: SqlanywhereDriver, items_table: str) -> None: + sqlanywhere_session.begin() + sqlanywhere_session.execute(f"INSERT INTO {items_table} (id, name) VALUES (?, ?)", (10, "pending")) + assert sqlanywhere_session.select_value(f"SELECT COUNT(*) FROM {items_table}") == 1 + sqlanywhere_session.rollback() + assert sqlanywhere_session.select_value(f"SELECT COUNT(*) FROM {items_table}") == 0 + + +def test_commit_is_visible_to_new_connection( + sqlanywhere_config: SqlanywhereConfig, sqlanywhere_session: SqlanywhereDriver, items_table: str +) -> None: + sqlanywhere_session.begin() + sqlanywhere_session.execute(f"INSERT INTO {items_table} (id, name) VALUES (?, ?)", (20, "kept")) + sqlanywhere_session.commit() + + connection = sqlanywhere_config.create_connection() + try: + cursor = connection.cursor() + cursor.execute(f"SELECT name FROM {items_table} WHERE id = 20") + assert cursor.fetchall() == [("kept",)] + finally: + connection.close() + + +def test_unique_violation_is_mapped(sqlanywhere_session: SqlanywhereDriver, items_table: str) -> None: + sqlanywhere_session.execute(f"INSERT INTO {items_table} (id, name) VALUES (?, ?)", (30, "first")) + with pytest.raises(UniqueViolationError): + sqlanywhere_session.execute(f"INSERT INTO {items_table} (id, name) VALUES (?, ?)", (30, "duplicate")) + + +def test_not_null_violation_is_mapped(sqlanywhere_session: SqlanywhereDriver, items_table: str) -> None: + with pytest.raises(NotNullViolationError): + sqlanywhere_session.execute(f"INSERT INTO {items_table} (id, name) VALUES (?, ?)", (40, None)) + + +def test_data_dictionary_reports_version_and_tables(sqlanywhere_session: SqlanywhereDriver, items_table: str) -> None: + dictionary = sqlanywhere_session.data_dictionary + version = dictionary.get_version(sqlanywhere_session) + assert version is not None + assert version.major >= 16 + table_names = {table["table_name"].lower() for table in dictionary.get_tables(sqlanywhere_session)} + assert items_table in table_names + column_names = [ + column["column_name"].lower() for column in dictionary.get_columns(sqlanywhere_session, items_table) + ] + assert column_names == ["id", "name", "note"] diff --git a/tests/integration/adapters/sqlite/sqlite/extensions/litestar/test_store.py b/tests/integration/adapters/sqlite/sqlite/extensions/litestar/test_store.py index bb100e104..374f044eb 100644 --- a/tests/integration/adapters/sqlite/sqlite/extensions/litestar/test_store.py +++ b/tests/integration/adapters/sqlite/sqlite/extensions/litestar/test_store.py @@ -20,11 +20,11 @@ from litestar.stores.base import Store from litestar.stores.registry import StoreRegistry - from sqlspec.adapters.sqlite.litestar.store import SQLiteStore + from sqlspec.adapters.sqlite.litestar.store import SqliteStore else: Store = import_optional_attr("litestar.stores.base", "Store") StoreRegistry = import_optional_attr("litestar.stores.registry", "StoreRegistry") - SQLiteStore = import_optional_attr("sqlspec.adapters.sqlite.litestar.store", "SQLiteStore") + SqliteStore = import_optional_attr("sqlspec.adapters.sqlite.litestar.store", "SqliteStore") pytest.importorskip("litestar") @@ -32,19 +32,19 @@ @pytest.fixture -async def sqlite_store() -> "AsyncGenerator[SQLiteStore, None]": +async def sqlite_store() -> "AsyncGenerator[SqliteStore, None]": """Create SQLite store with shared in-memory database.""" config = SqliteConfig( connection_config={"database": "file:test_sessions_mem?mode=memory&cache=shared", "uri": True}, extension_config={"litestar": {"session_table": "test_sessions"}}, ) - store = SQLiteStore(config) + store = SqliteStore(config) await store.create_table() yield store await store.delete_all() -async def test_sync_to_thread_concurrency(sqlite_store: "SQLiteStore") -> None: +async def test_sync_to_thread_concurrency(sqlite_store: "SqliteStore") -> None: """Test concurrent access via sync_to_thread wrapper. SQLite has write serialization, so we test sequential writes @@ -62,7 +62,7 @@ async def read_session(session_id: int) -> bytes | None: assert result == f"data_{i}".encode() -async def test_store_implements_litestar_store_contract(sqlite_store: "SQLiteStore") -> None: +async def test_store_implements_litestar_store_contract(sqlite_store: "SqliteStore") -> None: """SQLite stores satisfy Litestar's nominal store and registry contract.""" registry = StoreRegistry() diff --git a/tests/integration/adapters/sqlite/sqlite/test_pooling.py b/tests/integration/adapters/sqlite/sqlite/test_pooling.py index a3c776db6..21e936419 100644 --- a/tests/integration/adapters/sqlite/sqlite/test_pooling.py +++ b/tests/integration/adapters/sqlite/sqlite/test_pooling.py @@ -48,7 +48,7 @@ def test_regular_memory_auto_conversion(sqlite_config_regular_memory: SqliteConf connection_config = build_connection_config(config.connection_config) db_uri = connection_config["database"] - assert db_uri.startswith("file:memory_") and "cache=private" in db_uri + assert db_uri.startswith("file:memory_") and "cache=shared" in db_uri assert connection_config["uri"] is True with config.provide_session() as session1: @@ -248,7 +248,7 @@ def test_config_memory_database_conversion() -> None: try: db_uri = config.connection_config["database"] - assert db_uri.startswith("file:memory_") and "cache=private" in db_uri + assert db_uri.startswith("file:memory_") and "cache=shared" in db_uri assert config.connection_config["uri"] is True with config.provide_session() as session: @@ -267,7 +267,7 @@ def test_config_default_database() -> None: try: db_uri = config.connection_config["database"] - assert db_uri.startswith("file:memory_") and "cache=private" in db_uri + assert db_uri.startswith("file:memory_") and "cache=shared" in db_uri assert config.connection_config["uri"] is True with config.provide_session() as session: diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index 7b67de289..8bcf3d721 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -27,6 +27,7 @@ "pytest_databases.docker.rustfs", "pytest_databases.docker.snowflake", "pytest_databases.docker.spanner", + "pytest_databases.docker.sqlanywhere", ) _NON_PLUGIN_FIXTURE_MODULES: frozenset[str] = frozenset({"cleanup"}) @@ -48,6 +49,7 @@ "spanner_adk": (("google.cloud.spanner",),), "spanner_events": (("google.cloud.spanner",),), "spanner_litestar": (("google.cloud.spanner", "litestar"),), + "sqlanywhere": (("sqlanydb",),), "sqlite": ((),), } diff --git a/tests/integration/fixtures/adapter_cases.py b/tests/integration/fixtures/adapter_cases.py index e2bb886af..75ca1b5d7 100644 --- a/tests/integration/fixtures/adapter_cases.py +++ b/tests/integration/fixtures/adapter_cases.py @@ -76,7 +76,7 @@ from sqlspec.adapters.adbc import AdbcConfig, AdbcDriver from sqlspec.adapters.adbc.adk import AdbcADKStore - from sqlspec.adapters.adbc.litestar import ADBCStore + from sqlspec.adapters.adbc.litestar import AdbcStore from sqlspec.adapters.aiomysql import AiomysqlConfig, AiomysqlDriver, AiomysqlDriverFeatures from sqlspec.adapters.aiomysql.adk import AiomysqlADKStore from sqlspec.adapters.aiomysql.litestar import AiomysqlStore @@ -111,8 +111,8 @@ from sqlspec.adapters.db2.adk import Db2AsyncADKStore, Db2SyncADKStore from sqlspec.adapters.db2.litestar import Db2AsyncStore, Db2SyncStore from sqlspec.adapters.duckdb import DuckDBConfig, DuckDBDriver, DuckDBDriverFeatures - from sqlspec.adapters.duckdb.adk import DuckdbADKStore - from sqlspec.adapters.duckdb.litestar import DuckdbStore + from sqlspec.adapters.duckdb.adk import DuckDBADKStore + from sqlspec.adapters.duckdb.litestar import DuckDBStore from sqlspec.adapters.mssql_python import MssqlPythonConfig, MssqlPythonDriver from sqlspec.adapters.mssql_python.adk import MssqlPythonADKStore from sqlspec.adapters.mysqlconnector import ( @@ -152,7 +152,7 @@ from sqlspec.adapters.pymysql.litestar import PyMysqlStore from sqlspec.adapters.spanner import SpannerAsyncConfig, SpannerSyncConfig from sqlspec.adapters.sqlite.adk import SqliteADKStore - from sqlspec.adapters.sqlite.litestar import SQLiteStore + from sqlspec.adapters.sqlite.litestar import SqliteStore from tests.integration.adapters.bigquery._wedge import describe_wedge, is_emulator_wedge from tests.integration.fixtures.bigquery import _bigquery_connection_config from tests.integration.fixtures.mysql import _mysql_connection_config @@ -161,7 +161,7 @@ AdbcConfig = import_optional_attr("sqlspec.adapters.adbc", "AdbcConfig") AdbcDriver = import_optional_attr("sqlspec.adapters.adbc", "AdbcDriver") AdbcADKStore = import_optional_attr("sqlspec.adapters.adbc.adk", "AdbcADKStore") - ADBCStore = import_optional_attr("sqlspec.adapters.adbc.litestar", "ADBCStore") + AdbcStore = import_optional_attr("sqlspec.adapters.adbc.litestar", "AdbcStore") AiomysqlConfig = import_optional_attr("sqlspec.adapters.aiomysql", "AiomysqlConfig") AiomysqlDriver = import_optional_attr("sqlspec.adapters.aiomysql", "AiomysqlDriver") AiomysqlDriverFeatures = import_optional_attr("sqlspec.adapters.aiomysql", "AiomysqlDriverFeatures") @@ -230,8 +230,8 @@ DuckDBConfig = import_optional_attr("sqlspec.adapters.duckdb", "DuckDBConfig") DuckDBDriver = import_optional_attr("sqlspec.adapters.duckdb", "DuckDBDriver") DuckDBDriverFeatures = import_optional_attr("sqlspec.adapters.duckdb", "DuckDBDriverFeatures") - DuckdbADKStore = import_optional_attr("sqlspec.adapters.duckdb.adk", "DuckdbADKStore") - DuckdbStore = import_optional_attr("sqlspec.adapters.duckdb.litestar", "DuckdbStore") + DuckDBADKStore = import_optional_attr("sqlspec.adapters.duckdb.adk", "DuckDBADKStore") + DuckDBStore = import_optional_attr("sqlspec.adapters.duckdb.litestar", "DuckDBStore") MssqlPythonConfig = import_optional_attr("sqlspec.adapters.mssql_python", "MssqlPythonConfig") MssqlPythonDriver = import_optional_attr("sqlspec.adapters.mssql_python", "MssqlPythonDriver") MssqlPythonADKStore = import_optional_attr("sqlspec.adapters.mssql_python.adk", "MssqlPythonADKStore") @@ -289,7 +289,7 @@ SpannerAsyncConfig = import_optional_attr("sqlspec.adapters.spanner", "SpannerAsyncConfig") SpannerSyncConfig = import_optional_attr("sqlspec.adapters.spanner", "SpannerSyncConfig") SqliteADKStore = import_optional_attr("sqlspec.adapters.sqlite.adk", "SqliteADKStore") - SQLiteStore = import_optional_attr("sqlspec.adapters.sqlite.litestar", "SQLiteStore") + SqliteStore = import_optional_attr("sqlspec.adapters.sqlite.litestar", "SqliteStore") describe_wedge = import_optional_attr("tests.integration.adapters.bigquery._wedge", "describe_wedge") is_emulator_wedge = import_optional_attr("tests.integration.adapters.bigquery._wedge", "is_emulator_wedge") _bigquery_connection_config = import_optional_attr( @@ -1846,14 +1846,14 @@ def async_migration_case(request: pytest.FixtureRequest) -> MigrationCaseContext @pytest.fixture -async def contract_sqlite_store() -> "AsyncGenerator[SQLiteStore, None]": +async def contract_sqlite_store() -> "AsyncGenerator[SqliteStore, None]": """Provide a ready SQLite Litestar store for contract tests.""" pytest.importorskip("litestar") config = SqliteConfig( connection_config={"database": "file:contract_store_sqlite?mode=memory&cache=shared", "uri": True}, extension_config=_STORE_EXTENSION_CONFIG, ) - store = SQLiteStore(config) + store = SqliteStore(config) await store.create_table() yield store with contextlib.suppress(Exception): @@ -1876,7 +1876,7 @@ async def contract_aiosqlite_store() -> "AsyncGenerator[AiosqliteStore, None]": @pytest.fixture -async def contract_duckdb_store(tmp_path: Path) -> "AsyncGenerator[DuckdbStore, None]": +async def contract_duckdb_store(tmp_path: Path) -> "AsyncGenerator[DuckDBStore, None]": """Provide a ready DuckDB Litestar store for contract tests.""" pytest.importorskip("duckdb") pytest.importorskip("litestar") @@ -1884,7 +1884,7 @@ async def contract_duckdb_store(tmp_path: Path) -> "AsyncGenerator[DuckdbStore, connection_config={"database": str(tmp_path / f"store_{uuid4().hex}.duckdb")}, extension_config=_STORE_EXTENSION_CONFIG, ) - store = DuckdbStore(config) + store = DuckDBStore(config) await store.create_table() yield store with contextlib.suppress(Exception): @@ -2077,7 +2077,7 @@ async def contract_arrow_odbc_store(mssql_service: MSSQLService) -> "AsyncGenera @pytest.fixture -async def contract_adbc_store(postgres_service: PostgresService) -> "AsyncGenerator[ADBCStore, None]": +async def contract_adbc_store(postgres_service: PostgresService) -> "AsyncGenerator[AdbcStore, None]": """Provide a ready PostgreSQL-backed ADBC Litestar store for contract tests.""" pytest.importorskip("adbc_driver_manager") pytest.importorskip("pyarrow") @@ -2086,7 +2086,7 @@ async def contract_adbc_store(postgres_service: PostgresService) -> "AsyncGenera connection_config={"uri": _adbc_postgres_uri(postgres_service), "driver_name": "adbc_driver_postgresql"}, extension_config=_STORE_EXTENSION_CONFIG, ) - store = ADBCStore(config) + store = AdbcStore(config) await store.create_table() yield store with contextlib.suppress(Exception): @@ -2195,7 +2195,7 @@ def make() -> "tuple[Any, Any]": connection_config={"database": str(tmp_path / f"adk_{suffix}.duckdb")}, extension_config=_adk_extension_config(suffix), ) - return config, DuckdbADKStore(config) + return config, DuckDBADKStore(config) return make diff --git a/tests/integration/fixtures/db2.py b/tests/integration/fixtures/db2.py index 01f62eae8..af4431aeb 100644 --- a/tests/integration/fixtures/db2.py +++ b/tests/integration/fixtures/db2.py @@ -1,8 +1,15 @@ """IBM Db2 integration test fixtures.""" +import json import os +import platform import shlex +import sys +import sysconfig from collections.abc import AsyncGenerator, Generator +from importlib.metadata import version +from importlib.util import find_spec +from pathlib import Path from typing import TYPE_CHECKING, Any import pytest @@ -15,6 +22,7 @@ ) from pytest_databases.helpers import exec_output_to_bytes, get_xdist_worker_num +import sqlspec from sqlspec.adapters.db2.config import Db2AsyncConfig, Db2ConnectionParams, Db2SyncConfig if TYPE_CHECKING: @@ -24,6 +32,7 @@ from sqlspec.adapters.db2.driver import Db2AsyncDriver, Db2SyncDriver __all__ = ( + "db2_arrow_odbc_runtime", "db2_async_config", "db2_async_session", "db2_connection_config", @@ -42,6 +51,7 @@ def db2_image() -> str: @pytest.fixture(scope="session") def db2_service( + request: pytest.FixtureRequest, container_service: "ContainerService", db2_image: str, db2_dbname: str, @@ -51,6 +61,21 @@ def db2_service( xdist_db2_isolation_level: "XdistIsolationLevel", ) -> Generator[Db2Service, None, None]: """Start and yield a ready IBM Db2 service container with automatic 4Kn Direct-IO recovery.""" + client_needed = ( + sys.platform == "linux" + and platform.machine().lower() in {"x86_64", "amd64"} + and any( + isinstance(item, pytest.Function) and "db2_arrow_odbc_runtime" in item.fixturenames + for item in request.session.items + ) + ) + client_mounts = _db2_odbc_client_mounts() if client_needed else {} + if client_needed and not client_mounts: + msg = ( + "Arrow-ODBC Db2 acceptance requires an isolated or managed Linux x86-64 Python installation " + "and an available glibc dynamic loader; broad system Python prefixes cannot be mounted safely." + ) + raise RuntimeError(msg) worker_num = get_xdist_worker_num() name = "db2" initial_db = ( @@ -115,6 +140,7 @@ def check(service: "ServiceContainer") -> bool: transient=xdist_db2_isolation_level == "server", privileged=True, ipc_mode="host", + volumes=client_mounts or None, ) as service: if db2_dbname != initial_db: create_worker_database(service.container, db2_dbname, instance_user=db2_username) @@ -141,6 +167,90 @@ def db2_connection_config(db2_service: Db2Service) -> Db2ConnectionParams: } +@pytest.fixture(scope="session") +def db2_arrow_odbc_runtime(db2_service: Db2Service) -> tuple[list[str], dict[str, str], str]: + """Run the host client against the container's installed 64-bit IBM ODBC runtime. + + Linux x86-64 clients use their own Python and system loader from read-only + mounts. IBM's installation registry and instance profile remain inside the + existing Db2 service container; no host driver registration is changed. + """ + python_base = Path(sys.base_prefix).resolve() + executable = Path(sys.executable).resolve() + package_paths = tuple( + dict.fromkeys(str(Path(sysconfig.get_path(kind)).resolve()) for kind in ("purelib", "platlib")) + ) + library_path = ":".join(( + "/opt/sqlspec-host-libraries", + str(python_base / "lib"), + "/opt/ibm/db2/V11.5/lib64", + "/opt/ibm/db2/V11.5/lib64/gskit", + )) + driver = "/opt/ibm/db2/V11.5/lib64/libdb2o.so.1" + client_mounts = _db2_odbc_client_mounts() + if not client_mounts: + msg = "Arrow-ODBC Db2 acceptance requires a Linux x86-64 glibc host with its dynamic loader available." + raise RuntimeError(msg) + db2_service.container.reload() + actual_mounts = {item["Destination"]: item for item in db2_service.container.attrs.get("Mounts", ())} + missing_mounts = sorted( + item["bind"] + for item in client_mounts.values() + if item["bind"] not in actual_mounts or actual_mounts[item["bind"]].get("RW", True) + ) + if missing_mounts: + msg = ( + "The existing Db2 service lacks required read-only client mounts; cached containers cannot " + f"acquire new mounts in place. Missing or writable destinations: {missing_mounts!r}" + ) + raise RuntimeError(msg) + result = db2_service.container.exec_run(["test", "-r", str(executable)]) + if result.exit_code: + msg = f"The mounted canonical Python executable is not readable inside Db2: {str(executable)!r}" + raise RuntimeError(msg) + result = db2_service.container.exec_run(["test", "-r", driver]) + if result.exit_code: + msg = f"The pinned Db2 service image does not provide the required IBM 64-bit ODBC driver: {driver}" + raise RuntimeError(msg) + command = ["/opt/sqlspec-host-libraries/ld-linux-x86-64.so.2", "--library-path", library_path, str(executable)] + environment = { + "LD_LIBRARY_PATH": "/opt/ibm/db2/V11.5/lib64:/opt/ibm/db2/V11.5/lib64/gskit", + "DB2INSTANCE": db2_service.username, + "PYTHONHOME": str(python_base), + "PYTHONPATH": ":".join(dict.fromkeys((str(Path(sqlspec.__file__).resolve().parent.parent), *package_paths))), + } + probe = ( + "import json,sys,sqlspec,pyarrow,arrow_odbc; from importlib.metadata import version; " + "from sqlspec.adapters.arrow_odbc import ArrowOdbcConfig; " + "from pathlib import Path; " + "print(json.dumps({'python':list(sys.version_info[:3]), 'sqlspec':str(Path(sqlspec.__file__).resolve()), " + "'pyarrow':str(Path(pyarrow.__file__).resolve()), 'arrow_odbc':str(Path(arrow_odbc.__file__).resolve()), " + "'versions':{name:version(name) for name in ('sqlspec','pyarrow','arrow-odbc')}}))" + ) + result = db2_service.container.exec_run([*command, "-c", probe], environment=environment) + output = exec_output_to_bytes(result.output).decode("utf-8", errors="replace") + if result.exit_code: + msg = f"The mounted Db2 client Python environment failed initialization (exit {result.exit_code}):\n{output}" + raise RuntimeError(msg) + expected_origins: dict[str, str] = {} + for module_name in ("pyarrow", "arrow_odbc"): + spec = find_spec(module_name) + if spec is None or spec.origin is None: + msg = f"The host environment lacks the required {module_name!r} package origin." + raise RuntimeError(msg) + expected_origins[module_name] = str(Path(spec.origin).resolve()) + expected = { + "python": list(sys.version_info[:3]), + "sqlspec": str(Path(sqlspec.__file__).resolve()), + **expected_origins, + "versions": {name: version(name) for name in ("sqlspec", "pyarrow", "arrow-odbc")}, + } + if json.loads(output) != expected: + msg = f"The mounted Db2 client imported a different Python/package environment: {output}" + raise RuntimeError(msg) + return command, environment, driver + + @pytest.fixture(scope="session") def db2_sync_config(db2_connection_config: Db2ConnectionParams) -> Generator[Db2SyncConfig, None, None]: """Yield a session-scoped Db2SyncConfig and close its connection pool on teardown.""" @@ -175,6 +285,34 @@ async def db2_async_session(db2_async_config: Db2AsyncConfig) -> "AsyncGenerator yield driver +def _db2_odbc_client_mounts() -> dict[str, dict[str, str]]: + """Expose only the Python client, SQLSpec package, and host loader read-only.""" + if sys.platform != "linux" or platform.machine().lower() not in {"x86_64", "amd64"}: + return {} + python_base = Path(sys.base_prefix).resolve() + executable_base = Path(sys.executable).resolve().parent.parent + if any( + prefix in {Path("/"), Path("/usr"), Path("/usr/local")} + for prefix in (python_base, executable_base, Path(sys.prefix).resolve()) + ): + return {} + system_libraries = next( + (path for path in (Path("/lib/x86_64-linux-gnu"), Path("/lib64")) if (path / "ld-linux-x86-64.so.2").is_file()), + None, + ) + if system_libraries is None: + return {} + paths = { + python_base, + executable_base, + *(Path(sysconfig.get_path(kind)).resolve() for kind in ("purelib", "platlib")), + Path(sqlspec.__file__).resolve().parent, + } + mounts = {str(path): {"bind": str(path), "mode": "ro"} for path in paths} + mounts[str(system_libraries.resolve())] = {"bind": "/opt/sqlspec-host-libraries", "mode": "ro"} + return mounts + + def _ensure_db2_4k_database(container: Any, db_name: str, *, instance_user: str) -> None: """Configure 4Kn Direct-IO registry settings and create the initial Db2 database if missing.""" validate_db2_database_name(db_name) diff --git a/tests/integration/fixtures/sqlanywhere.py b/tests/integration/fixtures/sqlanywhere.py new file mode 100644 index 000000000..3ef595d9a --- /dev/null +++ b/tests/integration/fixtures/sqlanywhere.py @@ -0,0 +1,57 @@ +"""SAP SQL Anywhere integration test fixtures.""" + +from collections.abc import Generator +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import pytest +from pytest_databases.docker.sqlanywhere import SQLAnywhereService, sqlanywhere_client_libs + +from sqlspec.adapters.sqlanywhere import SqlanywhereConfig + +if TYPE_CHECKING: + from sqlspec.adapters.sqlanywhere import SqlanywhereDriver + +__all__ = ( + "sqlanywhere_client_libs", + "sqlanywhere_config", + "sqlanywhere_connection_config", + "sqlanywhere_session", + "xdist_sqlanywhere_isolation_level", +) + + +@pytest.fixture(scope="session") +def xdist_sqlanywhere_isolation_level() -> str: + """Run one SQL Anywhere server per worker because the demo image cannot create databases on demand.""" + return "server" + + +@pytest.fixture(scope="session") +def sqlanywhere_connection_config( + sqlanywhere_service: "SQLAnywhereService", sqlanywhere_client_libs: Path +) -> dict[str, Any]: + """Return ``sqlanydb`` connection parameters for the SQL Anywhere service.""" + return { + "uid": sqlanywhere_service.user, + "pwd": sqlanywhere_service.password, + "dbn": sqlanywhere_service.database, + "host": f"{sqlanywhere_service.host}:{sqlanywhere_service.port}", + } + + +@pytest.fixture +def sqlanywhere_config(sqlanywhere_connection_config: dict[str, Any]) -> Generator[SqlanywhereConfig, None, None]: + """Yield a SQL Anywhere config whose pool is closed after the test.""" + config = SqlanywhereConfig(connection_config=dict(sqlanywhere_connection_config)) + try: + yield config + finally: + config.close_pool() + + +@pytest.fixture +def sqlanywhere_session(sqlanywhere_config: SqlanywhereConfig) -> Generator["SqlanywhereDriver", None, None]: + """Yield a SQL Anywhere driver session.""" + with sqlanywhere_config.provide_session() as session: + yield session diff --git a/tests/unit/adapters/test_adapter_implementations.py b/tests/unit/adapters/test_adapter_implementations.py index ed02e73a5..d0f2e68e7 100644 --- a/tests/unit/adapters/test_adapter_implementations.py +++ b/tests/unit/adapters/test_adapter_implementations.py @@ -1,6 +1,5 @@ """Parameterized tests for real adapter implementations.""" -import ast import importlib import inspect import operator @@ -11,6 +10,7 @@ import pytest +import sqlspec.adapters from sqlspec.adapters.sqlite.driver import SqliteDriver from sqlspec.config import DatabaseConfigProtocol from sqlspec.core import SQL, ParameterStyle, ParameterStyleConfig, SQLResult, StatementConfig @@ -83,44 +83,10 @@ ) CREATE_MAPPED_EXCEPTION_MODULES = ADAPTER_CORE_MODULES APPLY_DRIVER_FEATURES_MODULES = ADAPTER_CORE_MODULES -ADAPTER_DRIVER_MODULES_WITH_FEATURE_HELPERS = ( - "sqlspec/adapters/cockroach_psycopg/driver.py", - "sqlspec/adapters/duckdb/driver.py", -) -MYSQL_FAMILY_UNCOMPILED_MODULES = ( - "sqlspec/adapters/aiomysql/config.py", - "sqlspec/adapters/aiomysql/driver.py", - "sqlspec/adapters/asyncmy/config.py", - "sqlspec/adapters/asyncmy/driver.py", - "sqlspec/adapters/mysqlconnector/config.py", - "sqlspec/adapters/mysqlconnector/driver.py", - "sqlspec/adapters/pymysql/driver.py", - "sqlspec/adapters/pymysql/pool.py", -) -MYSQL_VENDOR_IMPORT_ROOTS = ("aiomysql", "asyncmy", "mysql", "pymysql") -SPANNER_MODULE_PROXY_PATHS = ("sqlspec/adapters/spanner/config.py", "sqlspec/adapters/spanner/driver.py") SQLITE_EXCEPTION_HANDLER_CLASSES = ( ("sqlspec.adapters.aiosqlite.driver", "AiosqliteExceptionHandler"), ("sqlspec.adapters.sqlite.driver", "SqliteExceptionHandler"), ) -BRIDGE_CURSOR_CLOSE_METHODS = ( - ("sqlspec/adapters/adbc/_typing.py", "AdbcCursor", "__exit__"), - ("sqlspec/adapters/aiomysql/_typing.py", "AiomysqlCursor", "__aexit__"), - ("sqlspec/adapters/aiosqlite/_typing.py", "AiosqliteCursor", "__aexit__"), - ("sqlspec/adapters/asyncmy/_typing.py", "AsyncmyCursor", "__aexit__"), - ("sqlspec/adapters/db2/_typing.py", "Db2SyncCursor", "__exit__"), - ("sqlspec/adapters/db2/_typing.py", "Db2AsyncCursor", "__aexit__"), - ("sqlspec/adapters/mssql_python/_typing.py", "MssqlPythonCursor", "__exit__"), - ("sqlspec/adapters/mysqlconnector/_typing.py", "MysqlConnectorSyncCursor", "__exit__"), - ("sqlspec/adapters/mysqlconnector/_typing.py", "MysqlConnectorAsyncCursor", "__aexit__"), - ("sqlspec/adapters/oracledb/_typing.py", "OracleSyncCursor", "__exit__"), - ("sqlspec/adapters/oracledb/_typing.py", "OracleAsyncCursor", "__aexit__"), - ("sqlspec/adapters/psycopg/_typing.py", "PsycopgSyncCursor", "__exit__"), - ("sqlspec/adapters/psycopg/_typing.py", "PsycopgAsyncCursor", "__aexit__"), - ("sqlspec/adapters/pymssql/_typing.py", "PymssqlCursor", "__exit__"), - ("sqlspec/adapters/pymysql/_typing.py", "PyMysqlCursor", "__exit__"), - ("sqlspec/adapters/sqlite/_typing.py", "SqliteCursor", "__exit__"), -) @pytest.fixture(params=ADAPTER_CONFIGS, ids=operator.itemgetter("name")) @@ -205,95 +171,12 @@ def test_adapter_apply_driver_features_signature_and_return_shape(module_name: s assert features["custom"] == "value" -@pytest.mark.parametrize("module_path", ADAPTER_DRIVER_MODULES_WITH_FEATURE_HELPERS) -def test_adapter_drivers_do_not_apply_driver_features(module_path: str) -> None: - tree = ast.parse(Path(module_path).read_text()) - - imports_helper = any( - isinstance(node, ast.ImportFrom) - and node.module is not None - and node.module.endswith(".core") - and any(alias.name == "apply_driver_features" for alias in node.names) - for node in ast.walk(tree) - ) - calls_helper = any( - isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "apply_driver_features" - for node in ast.walk(tree) - ) - - assert imports_helper is False - assert calls_helper is False - - -@pytest.mark.parametrize("module_path", MYSQL_FAMILY_UNCOMPILED_MODULES) -def test_mysql_family_uncompiled_modules_use_adapter_typing_boundary(module_path: str) -> None: - tree = ast.parse(Path(module_path).read_text()) - direct_imports: list[str] = [] - - for node in ast.walk(tree): - if isinstance(node, ast.Import): - direct_imports.extend(alias.name for alias in node.names) - elif isinstance(node, ast.ImportFrom) and node.module is not None: - direct_imports.append(node.module) - - direct_vendor_imports = [ - import_name - for import_name in direct_imports - if import_name in MYSQL_VENDOR_IMPORT_ROOTS - or import_name.startswith(tuple(f"{root}." for root in MYSQL_VENDOR_IMPORT_ROOTS)) - ] - - assert direct_vendor_imports == [] - - -@pytest.mark.parametrize("module_path", SPANNER_MODULE_PROXY_PATHS) -def test_spanner_modules_do_not_define_module_proxy_getattr(module_path: str) -> None: - tree = ast.parse(Path(module_path).read_text()) - - module_getattrs = [node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name == "__getattr__"] - - assert module_getattrs == [] - - def test_mssql_python_pool_is_implemented_in_pool_module() -> None: from sqlspec.adapters.mssql_python import MssqlPythonConnectionPool assert MssqlPythonConnectionPool.__module__ == "sqlspec.adapters.mssql_python.pool" -@pytest.mark.parametrize(("module_path", "class_name", "method_name"), BRIDGE_CURSOR_CLOSE_METHODS) -def test_bridge_cursor_close_cleanup_suppresses_close_errors( - module_path: str, class_name: str, method_name: str -) -> None: - tree = ast.parse(Path(module_path).read_text()) - class_node = next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == class_name) - method_node = next( - node - for node in class_node.body - if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == method_name - ) - - suppress_calls = [ - item.context_expr - for node in ast.walk(method_node) - if isinstance(node, ast.With) - for item in node.items - if isinstance(item.context_expr, ast.Call) - ] - suppresses_close_errors = any( - isinstance(call.func, ast.Attribute) - and isinstance(call.func.value, ast.Name) - and call.func.value.id == "contextlib" - and call.func.attr == "suppress" - and len(call.args) == 1 - and isinstance(call.args[0], ast.Name) - and call.args[0].id == "Exception" - for call in suppress_calls - ) - - assert suppresses_close_errors - - def test_adapter_parameter_style_handling( adapter_config: "dict[str, Any]", statement_config_for_adapter: StatementConfig ) -> None: @@ -703,3 +586,32 @@ def test_every_config_declares_type_coercion_capabilities() -> None: assert caps is not DatabaseConfigProtocol.type_coercion_capabilities, ( f"{config_cls.__name__}.type_coercion_capabilities is the base default object" ) + + +def _litestar_store_targets() -> "list[tuple[str, str]]": + """Return ``(adapter, config class)`` pairs for adapters that ship a Litestar store.""" + targets: list[tuple[str, str]] = [] + adapters_root = Path(sqlspec.adapters.__file__).parent + for config_path in sorted(adapters_root.glob("*/config.py")): + adapter = config_path.parent.name + if not (config_path.parent / "litestar" / "store.py").exists(): + continue + module = importlib.import_module(f"sqlspec.adapters.{adapter}.config") + for name, value in vars(module).items(): + if ( + isinstance(value, type) + and value.__module__ == module.__name__ + and name.endswith("Config") + and issubclass(value, DatabaseConfigProtocol) + ): + targets.append((adapter, name)) + return targets + + +@pytest.mark.parametrize(("adapter", "config_name"), _litestar_store_targets()) +def test_litestar_session_migration_resolves_store_class(adapter: str, config_name: str) -> None: + """The Litestar session migration derives the store class name from the config class name.""" + store_name = config_name.replace("Config", "Store") + store_module = importlib.import_module(f"sqlspec.adapters.{adapter}.litestar.store") + + assert hasattr(store_module, store_name), f"{adapter}: expected {store_name} in litestar/store.py" diff --git a/tests/unit/adapters/test_adbc/test_adbc_config_normalization.py b/tests/unit/adapters/test_adbc/test_adbc_config_normalization.py index 7e8bcc50f..e543bb13f 100644 --- a/tests/unit/adapters/test_adbc/test_adbc_config_normalization.py +++ b/tests/unit/adapters/test_adbc/test_adbc_config_normalization.py @@ -1,21 +1,58 @@ """Unit tests for ADBC config normalization helpers.""" -from typing import Any, get_type_hints +import gc +import importlib.metadata +import json +import os +import sys +import threading +import weakref +from collections.abc import Callable, Iterator +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass, field +from importlib import import_module +from importlib.machinery import ModuleSpec +from types import ModuleType +from typing import Any, cast, get_type_hints from unittest.mock import MagicMock import adbc_driver_manager.dbapi +import adbc_driver_sqlite.dbapi import pytest from adbc_driver_bigquery import DatabaseOptions -from sqlspec.adapters.adbc import AdbcConfig, AdbcDriver +from sqlspec.adapters.adbc import AdbcConfig, AdbcDriver, AdbcDriverFeatures +from sqlspec.adapters.adbc._typing import AdbcSessionContext from sqlspec.adapters.adbc.config import AdbcConnectionParams from sqlspec.adapters.adbc.core import ( build_connection_config, + get_statement_config, resolve_driver_connect_func, resolve_driver_name_from_config, ) +from sqlspec.adapters.adbc.pool import AdbcConnectionPool +from sqlspec.adapters.adbc.type_converter import get_adbc_type_converter from sqlspec.core import SQL, ParameterStyle -from sqlspec.exceptions import SQLSpecError +from sqlspec.exceptions import ImproperConfigurationError, SQLSpecError +from sqlspec.utils.module_loader import reset_dependency_cache + + +@pytest.fixture +def optional_driver_imports(monkeypatch: pytest.MonkeyPatch) -> Iterator[Callable[[str, bool], None]]: + """Control optional SDK package imports through Python's module registry.""" + configured: set[str] = set() + + def configure(package: str, installed: bool) -> None: + module = ModuleType(package) if installed else None + if module is not None: + module.__spec__ = ModuleSpec(package, loader=None, is_package=True) + monkeypatch.setitem(sys.modules, package, module) + reset_dependency_cache(package) + configured.add(package) + + yield configure + for package in configured: + reset_dependency_cache(package) def test_adbc_config_runs_connection_create_callback(monkeypatch: Any) -> None: @@ -37,6 +74,506 @@ def fake_connect(**_kwargs: Any) -> Any: assert "on_connection_create" not in config.driver_features +def test_adbc_session_json_bind_accepts_unhashable_callable_serializer() -> None: + """A stateful callable serializer remains valid for public named parameter binding.""" + + @dataclass + class Serializer: + marker: str + + def __call__(self, value: Any) -> str: + return json.dumps({self.marker: value}, separators=(",", ":")) + + config = AdbcConfig( + connection_config={"driver_name": "sqlite"}, + statement_config=get_statement_config("postgres"), + driver_features={ + "json_serializer": Serializer("custom"), + "enable_pgvector": False, + "enable_paradedb": False, + "enable_pg_textsearch": False, + }, + ) + with config.provide_session() as session: + assert session.select_value("SELECT :payload", payload={"answer": 42}) == '{"custom":{"answer":42}}' + + +def test_adbc_converter_cache_preserves_distinct_equal_callable_identity() -> None: + """Equal hashable serializer objects retain their distinct conversion behavior.""" + + @dataclass(frozen=True) + class Serializer: + marker: str = field(compare=False) + + def __call__(self, value: Any) -> str: + return json.dumps({self.marker: value}, separators=(",", ":")) + + first_serializer = Serializer("first") + second_serializer = Serializer("second") + first = get_adbc_type_converter("postgres", json_serializer=first_serializer) + second = get_adbc_type_converter("postgres", json_serializer=second_serializer) + assert first.convert_dict({"answer": 42}) == '{"first":{"answer":42}}' + assert second.convert_dict({"answer": 42}) == '{"second":{"answer":42}}' + + +def test_adbc_converter_cache_releases_unused_serializer_after_turnover() -> None: + """Live cache entries reuse converters without retaining discarded serializers indefinitely.""" + + @dataclass + class Serializer: + marker: str + + def __call__(self, value: Any) -> str: + return json.dumps({self.marker: value}, separators=(",", ":")) + + serializer = Serializer("original") + reference = weakref.ref(serializer) + retained = get_adbc_type_converter("postgres", json_serializer=serializer) + assert get_adbc_type_converter("postgres", json_serializer=serializer) is retained + del retained, serializer + gc.collect() + assert reference() is not None + for index in range(512): + get_adbc_type_converter("postgres", json_serializer=Serializer(str(index))) + gc.collect() + assert reference() is None + + +def test_adbc_pool_option_enables_exclusive_reuse_without_changing_standalone_connections( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Only explicit pool settings enable reuse, while direct connections stay caller-owned.""" + seed, lease, standalone = MagicMock(), MagicMock(), MagicMock() + seed.adbc_clone.return_value = lease + create = MagicMock(side_effect=[seed, standalone]) + monkeypatch.setattr(import_module("adbc_driver_postgresql.dbapi"), "connect", create) + config = AdbcConfig( + connection_config={"driver_name": "postgres"}, + pool_config={}, + driver_features={"enable_pgvector": False, "enable_paradedb": False, "enable_pg_textsearch": False}, + ) + pool = config.provide_pool() + create.assert_not_called() + with config.provide_session() as first: + assert first.connection is lease + with config.provide_session() as second: + assert second.connection is lease + assert config.create_pool() is pool + assert config.create_connection() is standalone + standalone.close.assert_not_called() + config.close_pool() + seed.close.assert_called_once_with() + lease.close.assert_called_once_with() + + +def test_adbc_disabled_pool_methods_preserve_borrowed_connection() -> None: + """All disabled pool hooks are inert even when a legacy raw connection is supplied.""" + borrowed = MagicMock() + events: list[Any] = [] + config = AdbcConfig( + connection_config={"driver_name": "sqlite"}, + connection_instance=borrowed, + pool_config=None, + driver_features={"on_pool_create": events.append, "on_pool_destroy": events.append}, + ) + assert config.create_pool() is None + assert config.provide_pool() is None + assert config._create_pool() is None + config._close_pool() + config.close_pool() + assert config.create_connection() is borrowed + borrowed.close.assert_not_called() + assert events == [] + + +@pytest.mark.parametrize( + "invalid", + [ + "driver", + "autocommit", + "native_autocommit", + "injected", + "hook", + "size_zero", + "size_bool", + "size_float", + "timeout", + ], +) +def test_adbc_pool_rejects_unsupported_settings_before_native_io(monkeypatch: pytest.MonkeyPatch, invalid: str) -> None: + """Unsupported pooling settings fail before opening a native resource.""" + create = MagicMock() + monkeypatch.setattr(import_module("adbc_driver_postgresql.dbapi"), "connect", create) + options: dict[str, Any] = {"connection_config": {"driver_name": "postgres"}, "pool_config": {}} + if invalid == "driver": + options["connection_config"]["driver_name"] = "sqlite" + elif invalid == "autocommit": + options["connection_config"]["autocommit"] = True + elif invalid == "native_autocommit": + options["connection_config"]["conn_kwargs"] = {"adbc.connection.autocommit": "true"} + elif invalid == "injected": + options["connection_instance"] = MagicMock() + elif invalid == "hook": + options["driver_features"] = {"on_connection_create": lambda _connection: None} + elif invalid.startswith("size"): + options["pool_config"]["max_size"] = {"size_zero": 0, "size_bool": True, "size_float": 1.5}[invalid] + else: + options["pool_config"]["timeout"] = float("nan") + with pytest.raises(ImproperConfigurationError): + AdbcConfig(**options) + create.assert_not_called() + + +@pytest.mark.parametrize("distribution", ["adbc-driver-manager", "adbc-driver-postgresql"]) +@pytest.mark.parametrize("installed", ["1.11.0", "1.12.0rc1", "1.12.0.dev0", None]) +def test_adbc_pool_requires_supported_native_versions_only_when_enabled( + monkeypatch: pytest.MonkeyPatch, distribution: str, installed: str | None +) -> None: + """Unsupported or absent optional packages do not affect the unpooled default.""" + original_version = importlib.metadata.version + + def version(name: str) -> str: + if name != distribution: + return original_version(name) + if installed is None: + raise importlib.metadata.PackageNotFoundError(name) + return installed + + monkeypatch.setattr(importlib.metadata, "version", version) + assert AdbcConfig(connection_config={"driver_name": "postgres"}).create_pool() is None + with pytest.raises(ImproperConfigurationError, match=r"1\.12\.0"): + _postgres_pool_config() + + +def test_adbc_pool_requires_public_native_clone_capability(monkeypatch: pytest.MonkeyPatch) -> None: + """A sufficient package version without public cloning fails before opening a connection.""" + monkeypatch.setattr(adbc_driver_manager.dbapi.Connection, "adbc_clone", None) + with pytest.raises(ImproperConfigurationError, match="adbc_clone"): + _postgres_pool_config() + + +def test_adbc_pool_atomic_first_publication_and_callbacks_outside_config_lock() -> None: + """Concurrent first requests observe one generation and one outside-lock create callback.""" + events: list[Any] = [] + start = threading.Barrier(5, timeout=3) + + def created(pool: Any) -> None: + with ThreadPoolExecutor(max_workers=1) as executor: + assert executor.submit(config.provide_pool).result(timeout=2) is pool + events.append(pool) + + config = _postgres_pool_config(driver_features={"on_pool_create": created}) + + def acquire_pool() -> Any: + start.wait() + return config.provide_pool() + + with ThreadPoolExecutor(max_workers=4) as executor: + futures = [executor.submit(acquire_pool) for _ in range(4)] + start.wait() + pools = [future.result(timeout=3) for future in futures] + assert all(pool is pools[0] for pool in pools) + assert events == [pools[0]] + config.close_pool() + + +def test_adbc_pool_drain_retains_origin_and_restarts_after_final_release(monkeypatch: pytest.MonkeyPatch) -> None: + """A pending drain keeps ownership and retires exactly once before a new generation.""" + seed, lease = MagicMock(), MagicMock() + seed.adbc_clone.return_value = lease + monkeypatch.setattr(import_module("adbc_driver_postgresql.dbapi"), "connect", lambda **_kwargs: seed) + events: list[tuple[str, Any]] = [] + config = _postgres_pool_config( + driver_features={ + "on_pool_create": lambda pool: events.append(("create", pool)), + "on_pool_destroying": lambda pool: events.append(("destroying", pool)), + "on_pool_destroy": lambda pool: events.append(("destroy", pool)), + } + ) + context = config.provide_session() + context.__enter__() + origin = config.provide_pool() + config.close_pool() + assert config.connection_instance is origin + with pytest.raises(SQLSpecError, match="drain"): + config.provide_session() + seed.close.assert_not_called() + context.__exit__(None, None, None) + replacement = config.provide_pool() + assert replacement is not origin + assert events == [("create", origin), ("destroying", origin), ("destroy", origin), ("create", replacement)] + lease.close.assert_called_once_with() + seed.close.assert_called_once_with() + config.close_pool() + config.close_pool() + assert events[-2:] == [("destroying", replacement), ("destroy", replacement)] + + +def test_adbc_pool_cleanup_failure_retains_generation_for_retry(monkeypatch: pytest.MonkeyPatch) -> None: + """A native-close failure never permits a new generation before cleanup succeeds.""" + seed, lease = MagicMock(), MagicMock() + seed.adbc_clone.return_value = lease + lease.close.side_effect = [RuntimeError("native close failed"), None] + monkeypatch.setattr(import_module("adbc_driver_postgresql.dbapi"), "connect", lambda **_kwargs: seed) + config = _postgres_pool_config() + with config.provide_session(): + pass + origin = config.provide_pool() + with pytest.raises(RuntimeError, match="native close failed"): + config.close_pool() + assert config.connection_instance is origin + with pytest.raises(SQLSpecError, match="drain"): + config.create_pool() + seed.close.assert_not_called() + config.close_pool() + assert config.connection_instance is None + assert lease.close.call_count == 2 + seed.close.assert_called_once_with() + + +@pytest.mark.parametrize("close_fails", [False, True]) +def test_adbc_pool_body_error_evicts_and_preserves_primary(monkeypatch: pytest.MonkeyPatch, close_fails: bool) -> None: + """Failed session bodies discard their exact lease without masking the application error.""" + seed, lease = MagicMock(), MagicMock() + primary = ValueError("application failed") + seed.adbc_clone.return_value = lease + if close_fails: + lease.close.side_effect = [KeyboardInterrupt("native close interrupted"), None] + monkeypatch.setattr(import_module("adbc_driver_postgresql.dbapi"), "connect", lambda **_kwargs: seed) + config = _postgres_pool_config(pool_config={"max_size": 1, "timeout": 0}) + with pytest.raises(ValueError) as caught, config.provide_session() as session: + assert session.connection is lease + raise primary + assert caught.value is primary + lease.rollback.assert_called_once_with() + lease.adbc_connection.set_autocommit.assert_not_called() + if close_fails: + with pytest.raises(SQLSpecError, match="Timed out"): + with config.provide_session(): + pass + config.close_pool() + assert lease.close.call_count == (2 if close_fails else 1) + seed.close.assert_called_once_with() + + +def test_adbc_pool_destroying_reentry_cannot_close_replacement_generation() -> None: + """Reentrant callbacks cannot redirect a shutdown onto a newly published generation.""" + replacements: list[Any] = [] + events: list[tuple[str, Any]] = [] + + def destroying(pool: Any) -> None: + events.append(("destroying", pool)) + if not replacements: + pool.close() + replacements.append(config.provide_pool()) + + config = _postgres_pool_config( + driver_features={ + "on_pool_create": lambda pool: events.append(("create", pool)), + "on_pool_destroying": destroying, + "on_pool_destroy": lambda pool: events.append(("destroy", pool)), + } + ) + origin = config.provide_pool() + config.close_pool() + replacement = replacements[0] + assert config.provide_pool() is replacement + assert events == [("create", origin), ("destroying", origin), ("destroy", origin), ("create", replacement)] + config.close_pool() + + +def test_adbc_pool_session_releases_exact_acquired_generation(monkeypatch: pytest.MonkeyPatch) -> None: + """A handler returns its lease to the captured generation even if the config reference changes.""" + seed, lease = MagicMock(), MagicMock() + seed.adbc_clone.return_value = lease + monkeypatch.setattr(import_module("adbc_driver_postgresql.dbapi"), "connect", lambda **_kwargs: seed) + config = _postgres_pool_config() + context = config.provide_session() + context.__enter__() + origin = cast("AdbcConnectionPool", config.provide_pool()) + replacement = AdbcConnectionPool(MagicMock()) + config.connection_instance = replacement + context.__exit__(None, None, None) + lease.rollback.assert_called_once_with() + lease.cursor.return_value.execute.assert_called_once_with("DISCARD ALL") + assert origin.close() is True + seed.close.assert_called_once_with() + config.close_pool() + + +@pytest.mark.parametrize("event", ["create", "destroying"]) +def test_adbc_pool_callback_failure_retains_owned_generation(monkeypatch: pytest.MonkeyPatch, event: str) -> None: + """Lifecycle failures retain ownership and emit each lifecycle callback at most once.""" + seed, lease = MagicMock(), MagicMock() + seed.adbc_clone.return_value = lease + create = MagicMock(return_value=seed) + monkeypatch.setattr(import_module("adbc_driver_postgresql.dbapi"), "connect", create) + seen: list[Any] = [] + + def failure(pool: Any) -> None: + seen.append(pool) + raise KeyboardInterrupt("lifecycle callback failed") + + config = _postgres_pool_config(driver_features={f"on_pool_{event}": failure}) + if event == "create": + with pytest.raises(KeyboardInterrupt, match="lifecycle callback failed"): + config.create_pool() + assert config.provide_pool() is seen[0] + create.assert_not_called() + with config.provide_session(): + pass + if event == "destroying": + with pytest.raises(KeyboardInterrupt, match="lifecycle callback failed"): + config.close_pool() + assert config.connection_instance is seen[0] + seed.close.assert_not_called() + with pytest.raises(SQLSpecError, match="drain"): + config.create_pool() + config.close_pool() + assert len(seen) == 1 + assert config.connection_instance is None + seed.close.assert_called_once_with() + + +def test_adbc_pool_rejects_inherited_generation_before_config_lock_and_callbacks( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """All pool entrypoints reject inherited ownership before a poisoned config lock.""" + + class _LockProbeConfig(AdbcConfig): + def replace_lock(self, lock: Any) -> None: + self._pool_lock = lock + + def close_protected(self) -> None: + self._close_pool() + + events: list[Any] = [] + config = _LockProbeConfig( + connection_config={"driver_name": "postgres"}, + pool_config={}, + driver_features={"on_pool_create": events.append, "on_pool_destroying": events.append}, + ) + config.provide_pool() + events.clear() + poison = MagicMock() + poison.__enter__.side_effect = AssertionError("inherited config lock acquired") + config.replace_lock(poison) + child_pid = os.getpid() + 1 + monkeypatch.setattr(os, "getpid", lambda: child_pid) + for operation in ( + config.create_pool, + config.provide_pool, + config.provide_session, + config.close_pool, + config.close_protected, + ): + with pytest.raises(SQLSpecError, match="fork"): + operation() + poison.__enter__.assert_not_called() + assert events == [] + + +def test_adbc_pool_derive_preserves_settings_and_validates_overrides() -> None: + """Derivation retains opt-in bounds and starts independent native ownership.""" + config = _postgres_pool_config(pool_config={"max_size": 1, "timeout": 0}) + derived = config.derive_pool() + assert derived.provide_pool() is not config.provide_pool() + assert derived.derive_pool(pool_config=None).provide_pool() is None + with pytest.raises(ImproperConfigurationError, match="PostgreSQL"): + config.derive_pool(connection_config_overrides={"driver_name": "sqlite"}) + derived.close_pool() + config.close_pool() + + +@pytest.mark.parametrize("primary", [ValueError("hook failed"), KeyboardInterrupt("hook interrupted")]) +@pytest.mark.parametrize("cleanup_error", [None, RuntimeError("close failed"), SystemExit("close interrupted")]) +def test_adbc_config_hook_failure_closes_owned_connection( + monkeypatch: pytest.MonkeyPatch, primary: BaseException, cleanup_error: BaseException | None +) -> None: + """A failing initialization hook preserves its error after closing its owned connection.""" + connection = MagicMock() + connection.close.side_effect = cleanup_error + monkeypatch.setattr(adbc_driver_sqlite.dbapi, "connect", lambda **_kwargs: connection) + + def fail_hook(_connection: Any) -> None: + raise primary + + config = AdbcConfig( + connection_config={"driver_name": "sqlite"}, driver_features={"on_connection_create": fail_hook} + ) + expected = ImproperConfigurationError if isinstance(primary, Exception) else type(primary) + with pytest.raises(expected) as caught: + config.create_connection() + if isinstance(primary, Exception): + assert caught.value.__cause__ is primary + else: + assert caught.value is primary + connection.close.assert_called_once_with() + + +@pytest.mark.parametrize("stage", ["constructor", "preparation"]) +@pytest.mark.parametrize("primary", [RuntimeError("setup failed"), KeyboardInterrupt("setup interrupted")]) +@pytest.mark.parametrize("cleanup_error", [None, SystemExit("release interrupted")]) +def test_adbc_session_setup_failure_releases_once( + stage: str, primary: BaseException, cleanup_error: BaseException | None +) -> None: + """Failed session entry releases the acquisition with its original exception details.""" + connection = MagicMock() + released: list[dict[str, Any]] = [] + + class FailingDriver: + def __init__(self, **kwargs: Any) -> None: + raise primary + + def prepare(driver: AdbcDriver) -> AdbcDriver: + raise primary + + def release(conn: Any, **kwargs: Any) -> None: + assert conn is connection + assert kwargs["exc_tb"] is primary.__traceback__ + released.append(kwargs) + connection.close() + if cleanup_error is not None: + raise cleanup_error + + context = AdbcSessionContext( + acquire_connection=lambda: connection, + release_connection=release, + statement_config=AdbcConfig(connection_config={"driver_name": "sqlite"}).statement_config, + driver_features={}, + prepare_driver=prepare, + driver_class=cast("type[AdbcDriver]", FailingDriver) if stage == "constructor" else AdbcDriver, + ) + with pytest.raises(type(primary)) as caught: + context.__enter__() + assert caught.value is primary + assert len(released) == 1 + assert released[0]["exc_type"] is type(primary) + assert released[0]["exc_val"] is primary + context.__exit__(None, None, None) + connection.close.assert_called_once_with() + + +def test_adbc_session_release_once_after_cleanup_failure() -> None: + """An exit failure cannot release an acquisition again on a later exit.""" + connection = MagicMock() + cleanup_error = RuntimeError("close failed") + connection.close.side_effect = cleanup_error + context = AdbcSessionContext( + acquire_connection=lambda: connection, + release_connection=lambda conn, **_kwargs: conn.close(), + statement_config=AdbcConfig(connection_config={"driver_name": "sqlite"}).statement_config, + driver_features={}, + prepare_driver=lambda driver: driver, + ) + context.__enter__() + with pytest.raises(RuntimeError) as caught: + context.__exit__(None, None, None) + assert caught.value is cleanup_error + context.__exit__(None, None, None) + connection.close.assert_called_once_with() + + def test_resolve_driver_name_alias_to_connect_path() -> None: """Resolve short driver aliases to concrete connect paths.""" config = AdbcConfig(connection_config={"driver_name": "sqlite"}) @@ -73,16 +610,16 @@ def test_resolve_driver_name_from_uri() -> None: assert _resolve_driver_name(config) == "adbc_driver_postgresql.dbapi.connect" -def test_resolve_driver_name_gizmosql_alias(monkeypatch: Any) -> None: +def test_resolve_driver_name_gizmosql_alias(optional_driver_imports: Callable[[str, bool], None]) -> None: """Resolve GizmoSQL aliases to the FlightSQL driver when native gizmosql driver is absent.""" - monkeypatch.setattr("sqlspec.adapters.adbc.core.import_optional", lambda _name: None) + optional_driver_imports("adbc_driver_gizmosql", False) config = AdbcConfig(connection_config={"driver_name": "gizmosql"}) assert _resolve_driver_name(config) == "adbc_driver_flightsql.dbapi.connect" -def test_resolve_driver_name_from_gizmosql_uri(monkeypatch: Any) -> None: +def test_resolve_driver_name_from_gizmosql_uri(optional_driver_imports: Callable[[str, bool], None]) -> None: """Detect FlightSQL driver from GizmoSQL URI schemes when native gizmosql driver is absent.""" - monkeypatch.setattr("sqlspec.adapters.adbc.core.import_optional", lambda _name: None) + optional_driver_imports("adbc_driver_gizmosql", False) config = AdbcConfig(connection_config={"uri": "gizmosql://localhost:31337"}) assert _resolve_driver_name(config) == "adbc_driver_flightsql.dbapi.connect" config = AdbcConfig(connection_config={"uri": "gizmo://localhost:31337"}) @@ -153,9 +690,9 @@ def test_connection_config_dict_preserves_db_kwargs_for_non_bigquery() -> None: assert "foo" not in resolved -def test_flightsql_lifts_username_password_into_db_kwargs(monkeypatch: Any) -> None: +def test_flightsql_lifts_username_password_into_db_kwargs(optional_driver_imports: Callable[[str, bool], None]) -> None: """Top-level username/password move into db_kwargs on the FlightSQL fallback route.""" - monkeypatch.setattr("sqlspec.adapters.adbc.core.import_optional", lambda _name: None) + optional_driver_imports("adbc_driver_gizmosql", False) config = AdbcConfig(connection_config={"driver_name": "gizmosql", "username": "alice", "password": "secret"}) resolved = _get_connection_config_dict(config) @@ -164,36 +701,38 @@ def test_flightsql_lifts_username_password_into_db_kwargs(monkeypatch: Any) -> N assert "password" not in resolved -def test_flightsql_lifts_tls_skip_verify_true(monkeypatch: Any) -> None: +def test_flightsql_lifts_tls_skip_verify_true(optional_driver_imports: Callable[[str, bool], None]) -> None: """tls_skip_verify=True maps to the FlightSQL ADBC option key on the FlightSQL fallback route.""" - monkeypatch.setattr("sqlspec.adapters.adbc.core.import_optional", lambda _name: None) + optional_driver_imports("adbc_driver_gizmosql", False) config = AdbcConfig(connection_config={"driver_name": "gizmosql", "tls_skip_verify": True}) resolved = _get_connection_config_dict(config) assert resolved["db_kwargs"]["adbc.flight.sql.client_option.tls_skip_verify"] == "true" -def test_flightsql_lifts_tls_skip_verify_false(monkeypatch: Any) -> None: +def test_flightsql_lifts_tls_skip_verify_false(optional_driver_imports: Callable[[str, bool], None]) -> None: """tls_skip_verify=False maps to the FlightSQL ADBC option key on the FlightSQL fallback route.""" - monkeypatch.setattr("sqlspec.adapters.adbc.core.import_optional", lambda _name: None) + optional_driver_imports("adbc_driver_gizmosql", False) config = AdbcConfig(connection_config={"driver_name": "gizmosql", "tls_skip_verify": False}) resolved = _get_connection_config_dict(config) assert resolved["db_kwargs"]["adbc.flight.sql.client_option.tls_skip_verify"] == "false" -def test_flightsql_lifts_authorization_header(monkeypatch: Any) -> None: +def test_flightsql_lifts_authorization_header(optional_driver_imports: Callable[[str, bool], None]) -> None: """authorization_header maps to the FlightSQL ADBC option key.""" - monkeypatch.setattr("sqlspec.adapters.adbc.core.import_optional", lambda _name: None) + optional_driver_imports("adbc_driver_gizmosql", False) config = AdbcConfig(connection_config={"driver_name": "gizmosql", "authorization_header": "Bearer token"}) resolved = _get_connection_config_dict(config) assert resolved["db_kwargs"]["adbc.flight.sql.authorization_header"] == "Bearer token" -def test_flightsql_explicit_db_kwargs_wins_over_top_level_shortcuts(monkeypatch: Any) -> None: +def test_flightsql_explicit_db_kwargs_wins_over_top_level_shortcuts( + optional_driver_imports: Callable[[str, bool], None], +) -> None: """User-supplied db_kwargs entries take precedence over top-level shortcuts.""" - monkeypatch.setattr("sqlspec.adapters.adbc.core.import_optional", lambda _name: None) + optional_driver_imports("adbc_driver_gizmosql", False) config = AdbcConfig( connection_config={ "driver_name": "gizmosql", @@ -286,9 +825,9 @@ def test_gizmosql_parameter_style_is_qmark() -> None: assert config.statement_config.parameter_config.default_parameter_style == ParameterStyle.QMARK -def test_gizmo_alias_resolves_to_flightsql(monkeypatch: Any) -> None: +def test_gizmo_alias_resolves_to_flightsql(optional_driver_imports: Callable[[str, bool], None]) -> None: """The short 'gizmo' alias resolves to FlightSQL driver when native gizmosql driver is absent.""" - monkeypatch.setattr("sqlspec.adapters.adbc.core.import_optional", lambda _name: None) + optional_driver_imports("adbc_driver_gizmosql", False) config = AdbcConfig(connection_config={"driver_name": "gizmo"}) assert _resolve_driver_name(config) == "adbc_driver_flightsql.dbapi.connect" assert config.statement_config.dialect == "duckdb" @@ -511,11 +1050,11 @@ def fail_connect(_self: AdbcConfig) -> None: def test_adbc_driver_standalone_applies_features_and_enforces_strict_coercion() -> None: - """Standalone AdbcDriver populates default driver_features, aligns dialect, and honors strict_type_coercion.""" + """Standalone AdbcDriver populates default driver_features, aligns dialect, and honors enable_strict_type_coercion.""" conn = MagicMock() conn.adbc_get_info.return_value = {"vendor_name": "duckdb", "driver_name": "duckdb"} - driver = AdbcDriver(conn, driver_features={"strict_type_coercion": True}) + driver = AdbcDriver(conn, driver_features={"enable_strict_type_coercion": True}) assert driver.dialect == "duckdb" assert str(driver.statement_config.dialect) == "duckdb" assert "json_serializer" in driver.driver_features @@ -574,11 +1113,11 @@ def test_db2_driver_aliases_uri_and_parameter_profile() -> None: } -def test_mssql_driver_manager_fallback_and_tsql_statement_config(monkeypatch: Any) -> None: +def test_mssql_driver_manager_fallback_and_tsql_statement_config( + optional_driver_imports: Callable[[str, bool], None], +) -> None: """When adbc_driver_mssql is absent, mssql routes through adbc_driver_manager with driver='mssql'.""" - monkeypatch.setattr( - "sqlspec.adapters.adbc.core.import_optional", lambda name: None if name == "adbc_driver_mssql" else MagicMock() - ) + optional_driver_imports("adbc_driver_mssql", False) for conn_cfg in ( {"driver_name": "mssql", "uri": "sqlserver://sa:Pass@localhost:1433?database=master"}, @@ -599,22 +1138,17 @@ def test_mssql_driver_manager_fallback_and_tsql_statement_config(monkeypatch: An assert resolved["db_kwargs"]["uri"].startswith(("sqlserver://", "mssql://")) -def test_mssql_native_package_preferred_when_installed(monkeypatch: Any) -> None: +def test_mssql_native_package_preferred_when_installed(optional_driver_imports: Callable[[str, bool], None]) -> None: """When adbc_driver_mssql is installed, mssql resolves to adbc_driver_mssql.dbapi.connect.""" - monkeypatch.setattr( - "sqlspec.adapters.adbc.core.import_optional", lambda name: MagicMock() if name == "adbc_driver_mssql" else None - ) + optional_driver_imports("adbc_driver_mssql", True) config = AdbcConfig(connection_config={"driver_name": "mssql", "uri": "sqlserver://localhost:1433"}) assert _resolve_driver_name(config) == "adbc_driver_mssql.dbapi.connect" assert config.statement_config.dialect == "tsql" -def test_gizmosql_prefers_native_driver_when_installed(monkeypatch: Any) -> None: +def test_gizmosql_prefers_native_driver_when_installed(optional_driver_imports: Callable[[str, bool], None]) -> None: """When adbc_driver_gizmosql is installed, gizmosql resolves to adbc_driver_gizmosql.dbapi.connect.""" - monkeypatch.setattr( - "sqlspec.adapters.adbc.core.import_optional", - lambda name: MagicMock() if name == "adbc_driver_gizmosql" else None, - ) + optional_driver_imports("adbc_driver_gizmosql", True) config = AdbcConfig( connection_config={ "driver_name": "gizmosql", @@ -635,9 +1169,11 @@ def test_gizmosql_prefers_native_driver_when_installed(monkeypatch: Any) -> None assert _get_connection_config_dict(explicit_autocommit)["autocommit"] is True -def test_gizmosql_flightsql_fallback_lifts_auth_type_into_db_kwargs(monkeypatch: Any) -> None: +def test_gizmosql_flightsql_fallback_lifts_auth_type_into_db_kwargs( + optional_driver_imports: Callable[[str, bool], None], +) -> None: """When adbc_driver_gizmosql is absent, auth_type is lifted into db_kwargs['adbc.gizmosql.auth_type'].""" - monkeypatch.setattr("sqlspec.adapters.adbc.core.import_optional", lambda _name: None) + optional_driver_imports("adbc_driver_gizmosql", False) config = AdbcConfig( connection_config={"driver_name": "gizmosql", "uri": "grpc+tls://localhost:31337", "auth_type": "external"} ) @@ -790,6 +1326,18 @@ def __arrow_c_stream__(self, requested_schema: object = None) -> object: assert cursor.executemany.call_args_list[5].args == ("INSERT INTO items (id, name) VALUES (?, ?)", arrow_stream) +def _postgres_pool_config(**overrides: Any) -> AdbcConfig: + """Construct a pooled config without unrelated native extension probes.""" + features: dict[str, Any] = {"enable_pgvector": False, "enable_paradedb": False, "enable_pg_textsearch": False} + features.update(overrides.pop("driver_features", {})) + return AdbcConfig( + connection_config=overrides.pop("connection_config", {"driver_name": "postgres"}), + pool_config=overrides.pop("pool_config", {}), + driver_features=features, + **overrides, + ) + + def _resolve_driver_name(config: AdbcConfig) -> str: """Resolve driver name from configuration.""" return resolve_driver_name_from_config(config.connection_config) @@ -798,3 +1346,21 @@ def _resolve_driver_name(config: AdbcConfig) -> str: def _get_connection_config_dict(config: AdbcConfig) -> dict[str, Any]: """Build the normalized connection configuration.""" return build_connection_config(config.connection_config) + + +def test_snowflake_lowercase_default() -> None: + """Snowflake dialect defaults enable_lowercase_column_names to True.""" + config = AdbcConfig(connection_config={"driver_name": "snowflake"}) + assert config.driver_features["enable_lowercase_column_names"] is True + + +def test_driver_features_have_no_legacy_aliases() -> None: + """AdbcDriverFeatures declares only the enable_* keys, and legacy alias keys are not translated.""" + assert "strict_type_coercion" not in AdbcDriverFeatures.__annotations__ + assert "arrow_extension_types" not in AdbcDriverFeatures.__annotations__ + config = AdbcConfig( + connection_config={"driver_name": "sqlite"}, + driver_features={"enable_strict_type_coercion": True, "enable_arrow_extension_types": False}, + ) + assert config.driver_features["enable_strict_type_coercion"] is True + assert config.driver_features["enable_arrow_extension_types"] is False diff --git a/tests/unit/adapters/test_adbc/test_adbc_serialization.py b/tests/unit/adapters/test_adbc/test_adbc_serialization.py index d45887f92..30d947cee 100644 --- a/tests/unit/adapters/test_adbc/test_adbc_serialization.py +++ b/tests/unit/adapters/test_adbc/test_adbc_serialization.py @@ -89,8 +89,8 @@ def test_driver_features_defaults() -> None: assert config.driver_features["json_serializer"] is to_json assert config.driver_features["enable_cast_detection"] is True - assert config.driver_features["strict_type_coercion"] is False - assert config.driver_features["arrow_extension_types"] is True + assert config.driver_features["enable_strict_type_coercion"] is False + assert config.driver_features["enable_arrow_extension_types"] is True def test_driver_features_override() -> None: @@ -103,14 +103,14 @@ def custom_serializer(data: Any) -> str: driver_features={ "json_serializer": custom_serializer, "enable_cast_detection": False, - "strict_type_coercion": True, + "enable_strict_type_coercion": True, } ) assert config.driver_features["json_serializer"] is custom_serializer assert config.driver_features["enable_cast_detection"] is False - assert config.driver_features["strict_type_coercion"] is True - assert config.driver_features["arrow_extension_types"] is True + assert config.driver_features["enable_strict_type_coercion"] is True + assert config.driver_features["enable_arrow_extension_types"] is True def test_backward_compatibility_no_serializer() -> None: diff --git a/tests/unit/adapters/test_adbc/test_adk_dialect_support.py b/tests/unit/adapters/test_adbc/test_adk_dialect_support.py index ee5b6a079..bc1283783 100644 --- a/tests/unit/adapters/test_adbc/test_adk_dialect_support.py +++ b/tests/unit/adapters/test_adbc/test_adk_dialect_support.py @@ -6,6 +6,7 @@ from unittest.mock import MagicMock import pytest +from adbc_driver_manager import AdbcStatusCode, IntegrityError from sqlspec.adapters.adbc import AdbcConfig from sqlspec.adapters.adbc.adk import AdbcADKConfig, AdbcADKMemoryStore, AdbcADKStore @@ -16,9 +17,10 @@ format_pgvector_literal, ) from sqlspec.adapters.adbc.events.store import AdbcEventQueueStore -from sqlspec.adapters.adbc.litestar.store import ADBCStore +from sqlspec.adapters.adbc.litestar.store import AdbcStore from sqlspec.core import StatementConfig from sqlspec.exceptions import ImproperConfigurationError +from sqlspec.extensions.adk import StoredMemory pytestmark = [pytest.mark.xdist_group("sqlite"), pytest.mark.adbc, pytest.mark.integration] @@ -336,9 +338,9 @@ def test_sqlite_adk_store_encodes_and_decodes_timestamps() -> None: def test_adbc_litestar_store_handles_pg_extensions_and_bigquery_snowflake(monkeypatch: Any) -> None: - """ADBCStore should recognize Postgres extension dialects, omit indexes on BigQuery/Snowflake, and use select_one_or_none.""" + """AdbcStore should recognize Postgres extension dialects, omit indexes on BigQuery/Snowflake, and use select_one_or_none.""" config = AdbcConfig(connection_config={"driver_name": "bigquery"}) - store = ADBCStore(config) + store = AdbcStore(config) store._dialect = "bigquery" # pyright: ignore[reportPrivateUsage] ddl = store._table_ddl() # pyright: ignore[reportPrivateUsage] @@ -399,7 +401,7 @@ def test_detect_enterprise_dialects_in_adk_stores(driver_name: str, expected_dia config = AdbcConfig(connection_config={"driver_name": driver_name}) store = AdbcADKStore(config) mem_store = AdbcADKMemoryStore(config) - assert store.dialect() == expected_dialect + assert store.dialect == expected_dialect assert mem_store.dialect == expected_dialect @@ -408,7 +410,7 @@ def test_detect_dialect_fallback_to_statement_config() -> None: config = AdbcConfig(connection_config={}, statement_config=StatementConfig(dialect="oracle")) store = AdbcADKStore(config) mem_store = AdbcADKMemoryStore(config) - assert store.dialect() == "oracle" + assert store.dialect == "oracle" assert mem_store.dialect == "oracle" @@ -760,3 +762,163 @@ def _store_for(driver_name: str) -> AdbcADKStore: def _normalized(sql: str) -> str: return " ".join(sql.split()) + + +def test_oracle_adk_memory_store_pre_queries_existing_event_ids(monkeypatch: pytest.MonkeyPatch) -> None: + """AdbcADKMemoryStore pre-queries existing event_ids on non-returning dialects.""" + executed_statements: list[str] = [] + mock_cursor = MagicMock() + mock_cursor.fetchall.return_value = [("evt-existing",)] + + def record_execute(sql: str, _params: Any = None) -> None: + executed_statements.append(sql) + + mock_cursor.execute.side_effect = record_execute + mock_conn = MagicMock() + mock_conn.cursor.return_value = mock_cursor + mock_ctx = MagicMock() + mock_ctx.__enter__.return_value = mock_conn + mock_ctx.__exit__.return_value = None + monkeypatch.setattr(AdbcConfig, "provide_connection", lambda *_a, **_k: mock_ctx) + + config = AdbcConfig( + connection_config={"driver_name": "oracle", "uri": ":memory:"}, + extension_config={"adk": {"enable_memory": True}}, + ) + mem_store = AdbcADKMemoryStore(config) + now = datetime(2026, 6, 1, 12, 0, 0, tzinfo=UTC) + + count = mem_store.insert_memory_entries([ + { + "id": "mem-1", + "session_id": "sess-1", + "app_name": "app1", + "user_id": "user1", + "scope": "user", + "event_id": "evt-existing", + "author": "user", + "timestamp": now, + "content_json": {"parts": [{"text": "hello"}]}, + "content_text": "hello", + "metadata_json": None, + "inserted_at": now, + "embedding": None, + }, + { + "id": "mem-2", + "session_id": "sess-1", + "app_name": "app1", + "user_id": "user1", + "scope": "user", + "event_id": "evt-new", + "author": "user", + "timestamp": now, + "content_json": {"parts": [{"text": "world"}]}, + "content_text": "world", + "metadata_json": None, + "inserted_at": now, + "embedding": None, + }, + ]) + + assert any("SELECT event_id FROM" in s for s in executed_statements) + assert count == 1 + + +@pytest.mark.parametrize("driver_name", ["oracle", "db2"]) +def test_adbc_memory_batch_deduplicates_fresh_event_ids(driver_name: str, monkeypatch: pytest.MonkeyPatch) -> None: + """Repeated fresh event IDs keep the first entry and preserve the rest of the batch.""" + cursor = MagicMock() + cursor.fetchall.return_value = [] + conn = MagicMock() + conn.cursor.return_value = cursor + context = MagicMock() + context.__enter__.return_value = conn + monkeypatch.setattr(AdbcConfig, "provide_connection", lambda *_a, **_k: context) + store = AdbcADKMemoryStore(AdbcConfig(connection_config={"driver_name": driver_name})) + now = datetime.now(UTC) + entry: StoredMemory = { + "id": "mem-1", + "session_id": "sess-1", + "app_name": "app", + "user_id": "u1", + "scope": "user", + "event_id": "evt-1", + "author": "user", + "timestamp": now, + "content_json": {"text": "hello"}, + "content_text": "hello", + "metadata_json": None, + "inserted_at": now, + "embedding": None, + } + + count = store.insert_memory_entries([ + entry, + {**entry, "id": "mem-duplicate"}, + {**entry, "id": "mem-2", "event_id": "evt-2"}, + ]) + + inserts = [call.args[1] for call in cursor.execute.call_args_list if "INSERT INTO" in call.args[0]] + assert count == 2 + assert [params[0] for params in inserts] == ["mem-1", "mem-2"] + conn.commit.assert_called_once() + + +@pytest.mark.parametrize("driver_name", ["oracle", "db2"]) +@pytest.mark.parametrize( + ("error_message", "event_exists"), + [ + ("ORA-00001: unique constraint violated", True), + ("SQL0803N duplicate key", True), + ("ORA-00001: unique constraint violated", False), + ("foreign key constraint violated", True), + ], +) +def test_adbc_memory_insert_only_ignores_confirmed_event_conflicts( + driver_name: str, error_message: str, event_exists: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + """A duplicate race is idempotent while other integrity errors still roll back.""" + cursor = MagicMock() + cursor.fetchall.return_value = [] + cursor.fetchone.return_value = ("evt-1",) if event_exists else None + error = IntegrityError(error_message, status_code=AdbcStatusCode.INTEGRITY) + + def execute(sql: str, _params: object = None) -> None: + if "INSERT INTO" in sql: + raise error + + cursor.execute.side_effect = execute + conn = MagicMock() + conn.cursor.return_value = cursor + context = MagicMock() + context.__enter__.return_value = conn + monkeypatch.setattr(AdbcConfig, "provide_connection", lambda *_a, **_k: context) + store = AdbcADKMemoryStore(AdbcConfig(connection_config={"driver_name": driver_name})) + now = datetime.now(UTC) + entry: StoredMemory = { + "id": "mem-1", + "session_id": "sess-1", + "app_name": "app", + "user_id": "u1", + "scope": "user", + "event_id": "evt-1", + "author": "user", + "timestamp": now, + "content_json": {"text": "hello"}, + "content_text": "hello", + "metadata_json": None, + "inserted_at": now, + "embedding": None, + } + + if event_exists and "foreign key" not in error_message: + assert store.insert_memory_entries([entry]) == 0 + conn.commit.assert_called_once() + conflict_checks = [call.args[1] for call in cursor.execute.call_args_list if "WHERE event_id =" in call.args[0]] + assert conflict_checks == [("evt-1",)] + else: + with pytest.raises(IntegrityError) as exc_info: + store.insert_memory_entries([entry]) + assert exc_info.value is error + conn.rollback.assert_called_once() diff --git a/tests/unit/adapters/test_adbc/test_arrow_streaming.py b/tests/unit/adapters/test_adbc/test_arrow_streaming.py index 7d97784b4..9124603b5 100644 --- a/tests/unit/adapters/test_adbc/test_arrow_streaming.py +++ b/tests/unit/adapters/test_adbc/test_arrow_streaming.py @@ -1,16 +1,22 @@ """Unit tests for ADBC Arrow streaming paths.""" +from importlib import import_module from typing import TYPE_CHECKING, Any, cast +from unittest.mock import MagicMock import pyarrow as pa import pytest from adbc_driver_manager import AdbcStatusCode, DatabaseError +from sqlspec.adapters.adbc import AdbcConfig +from sqlspec.adapters.adbc._typing import AdbcSessionContext +from sqlspec.adapters.adbc.core import get_statement_config from sqlspec.adapters.adbc.driver import AdbcDriver, AdbcSelectStreamSource from sqlspec.exceptions import SQLSpecError if TYPE_CHECKING: from sqlspec.adapters.adbc._typing import AdbcConnection + from sqlspec.driver import SyncRowStream _STORAGE_CAPABILITIES = { @@ -22,18 +28,44 @@ } +class _AdbcStreamingReader: + def __init__(self) -> None: + table = pa.table({"x": [1, 2, 3]}) + self.schema = table.schema + self._reader = pa.RecordBatchReader.from_batches(table.schema, table.to_batches(max_chunksize=2)) + self.close_count = 0 + self.close_error: BaseException | None = None + + def __iter__(self) -> "_AdbcStreamingReader": + return self + + def __next__(self) -> pa.RecordBatch: + return self._reader.read_next_batch() + + def read_next_batch(self) -> pa.RecordBatch: + return next(self) + + def close(self) -> None: + self.close_count += 1 + self._reader.close() + if self.close_error is not None: + raise self.close_error + + class _AdbcStreamingCursor: def __init__(self) -> None: self.closed = False + self.close_count = 0 + self.close_error: BaseException | None = None + self.reader_obj = _AdbcStreamingReader() self.executed: list[tuple[str, object]] = [] self.ingest_calls: list[tuple[str, object, str, dict[str, Any]]] = [] def execute(self, sql: str, parameters: object = None) -> None: self.executed.append((sql, parameters)) - def fetch_record_batch(self) -> pa.RecordBatchReader: - table = pa.table({"x": [1, 2, 3]}) - return pa.RecordBatchReader.from_batches(table.schema, table.to_batches(max_chunksize=2)) + def fetch_record_batch(self) -> _AdbcStreamingReader: + return self.reader_obj def fetch_arrow_table(self) -> pa.Table: msg = "reader formats should use fetch_record_batch(), not fetch_arrow_table()" @@ -45,6 +77,9 @@ def adbc_ingest(self, table: str, source: object, *, mode: str, **kwargs: Any) - def close(self) -> None: self.closed = True + self.close_count += 1 + if self.close_error is not None: + raise self.close_error class _AdbcStreamingConnection: @@ -58,6 +93,204 @@ def cursor(self) -> _AdbcStreamingCursor: return self.cursor_obj +@pytest.mark.parametrize("kind", ["row", "reader"]) +@pytest.mark.parametrize("progress", ["never", "partial", "exhausted"]) +def test_session_closes_lazy_stream_owners_before_release(kind: str, progress: str) -> None: + """All lazy streams taint their session and close actual native owners before release.""" + connection = _AdbcStreamingConnection() + released: list[AdbcDriver] = [] + + def release(_connection: Any, **kwargs: Any) -> None: + expected_closes = 0 if kind == "row" and progress == "never" else 1 + assert connection.cursor_obj.close_count == expected_closes + assert connection.cursor_obj.reader_obj.close_count == expected_closes + assert kwargs["driver"].has_streams is True + released.append(kwargs["driver"]) + + context = AdbcSessionContext( + acquire_connection=lambda: connection, + release_connection=release, + statement_config=get_statement_config("sqlite"), + driver_features={}, + prepare_driver=lambda driver: driver, + ) + stream: SyncRowStream[dict[str, Any]] | None = None + with context as driver: + if kind == "row": + stream = driver.select_stream("SELECT 1 AS x", native_only=True, chunk_size=1) + if progress == "partial": + assert next(stream) == {"x": 1} + elif progress == "exhausted": + assert list(stream) == [{"x": 1}, {"x": 2}, {"x": 3}] + else: + reader = driver.select_to_arrow("SELECT 1 AS x", return_format="reader").data + if progress == "partial": + assert reader.read_next_batch().to_pydict() == {"x": [1, 2]} + elif progress == "exhausted": + assert reader.read_all().to_pydict() == {"x": [1, 2, 3]} + assert released == [driver] + driver.close_streams() + if stream is not None: + assert list(stream) == [] + + +def test_lazy_fallback_stream_taints_session_before_execution() -> None: + """The generic lazy fallback also prevents reuse before it opens native resources.""" + connection = _AdbcStreamingConnection() + released: list[AdbcDriver] = [] + + def release(_connection: Any, **kwargs: Any) -> None: + assert kwargs["driver"].has_streams is True + assert connection.cursor_obj.executed == [] + assert connection.cursor_obj.close_count == 0 + released.append(kwargs["driver"]) + + context = AdbcSessionContext( + acquire_connection=lambda: connection, + release_connection=release, + statement_config=get_statement_config("sqlite"), + driver_features={}, + prepare_driver=lambda driver: driver, + ) + with context as driver: + stream = driver.select_stream("UPDATE items SET x = 1") + assert released == [driver] + stream.close() + + +@pytest.mark.parametrize("kind,consume", [("row", False), ("reader", False), ("reader", True)]) +def test_pooled_lazy_stream_closes_native_owners_then_evicts_origin( + monkeypatch: pytest.MonkeyPatch, kind: str, consume: bool +) -> None: + """Stream cleanup precedes native connection disposal and the next session gets another clone.""" + + class _PooledStreamingConnection(_AdbcStreamingConnection): + def __init__(self) -> None: + super().__init__() + self.close_count = 0 + self.rollback_count = 0 + + def rollback(self) -> None: + self.rollback_count += 1 + + def close(self) -> None: + expected = 0 if kind == "row" else 1 + assert self.cursor_obj.close_count == expected + assert self.cursor_obj.reader_obj.close_count == expected + self.close_count += 1 + + seed, replacement = MagicMock(), MagicMock() + lease = _PooledStreamingConnection() + seed.adbc_clone.side_effect = [lease, replacement] + monkeypatch.setattr(import_module("adbc_driver_postgresql.dbapi"), "connect", lambda **_kwargs: seed) + config = AdbcConfig( + connection_config={"driver_name": "postgres"}, + pool_config={"max_size": 1, "timeout": 0}, + driver_features={"enable_pgvector": False, "enable_paradedb": False, "enable_pg_textsearch": False}, + ) + with config.provide_session() as session: + assert session.connection is lease + if kind == "row": + session.select_stream("SELECT 1 AS x", native_only=True) + else: + reader = session.select_to_arrow("SELECT 1 AS x", return_format="reader").data + if consume: + assert reader.read_all().to_pydict() == {"x": [1, 2, 3]} + assert lease.close_count == 1 + assert lease.rollback_count == 1 + with config.provide_session() as session: + assert session.connection is replacement + config.close_pool() + seed.close.assert_called_once_with() + + +@pytest.mark.parametrize("stage", ["exit", "body", "setup"]) +def test_session_stream_cleanup_attempts_all_owners_and_preserves_primary(stage: str) -> None: + """Cleanup attempts every native reader/cursor and release, preserving setup/body errors.""" + cursors = [_AdbcStreamingCursor(), _AdbcStreamingCursor()] + available = iter(cursors) + + class _MultipleStreamingConnection(_AdbcStreamingConnection): + def cursor(self) -> _AdbcStreamingCursor: + return next(available) + + connection = _MultipleStreamingConnection() + primary = KeyboardInterrupt("primary interruption") + first_cleanup = RuntimeError("first reader cleanup failed") + cursors[0].reader_obj.close_error = first_cleanup + cursors[0].close_error = RuntimeError("cursor cleanup failed") + cursors[1].reader_obj.close_error = SystemExit("second reader cleanup failed") + released: list[AdbcDriver] = [] + drivers: list[AdbcDriver] = [] + readers: list[Any] = [] + + def open_readers(driver: AdbcDriver) -> None: + drivers.append(driver) + readers.extend(driver.select_to_arrow("SELECT 1 AS x", return_format="reader").data for _ in range(2)) + + def prepare(driver: AdbcDriver) -> AdbcDriver: + if stage == "setup": + open_readers(driver) + raise primary + return driver + + def release(_connection: Any, **kwargs: Any) -> None: + for cursor in cursors: + assert cursor.close_count == 1 + assert cursor.reader_obj.close_count == 1 + released.append(kwargs["driver"]) + raise RuntimeError("release cleanup failed") + + context = AdbcSessionContext( + acquire_connection=lambda: connection, + release_connection=release, + statement_config=get_statement_config("sqlite"), + driver_features={}, + prepare_driver=prepare, + ) + expected = RuntimeError if stage == "exit" else KeyboardInterrupt + with pytest.raises(expected) as caught, context as driver: + open_readers(driver) + if stage == "body": + raise primary + assert caught.value is (first_cleanup if stage == "exit" else primary) + assert released == drivers + context.__exit__(None, None, None) + drivers[0].close_streams() + assert [cursor.close_count for cursor in cursors] == [1, 1] + + +def test_eager_arrow_batches_close_owners_without_tainting_session() -> None: + """Eager batches release their native resources immediately and remain reusable.""" + connection = _AdbcStreamingConnection() + driver = AdbcDriver(cast("AdbcConnection", connection), dialect="sqlite") + result = driver.select_to_arrow("SELECT 1 AS x", return_format="batches") + assert [batch.to_pydict() for batch in result.data] == [{"x": [1, 2]}, {"x": [3]}] + assert driver.has_streams is False + assert connection.cursor_obj.close_count == 1 + assert connection.cursor_obj.reader_obj.close_count == 1 + driver.close_streams() + assert connection.cursor_obj.close_count == 1 + + +def test_materialized_arrow_table_does_not_taint_session() -> None: + """Materialized tables close their cursor and do not require stream eviction.""" + + class _MaterializedCursor(_AdbcStreamingCursor): + def fetch_arrow_table(self) -> pa.Table: + return pa.table({"x": [1, 2, 3]}) + + connection = _AdbcStreamingConnection() + connection.cursor_obj = _MaterializedCursor() + driver = AdbcDriver(cast("AdbcConnection", connection), dialect="sqlite") + result = driver.select_to_arrow("SELECT 1 AS x", return_format="table") + assert result.data.to_pydict() == {"x": [1, 2, 3]} + assert driver.has_streams is False + assert connection.cursor_obj.close_count == 1 + driver.close_streams() + assert connection.cursor_obj.close_count == 1 + + def test_select_to_arrow_reader_uses_fetch_record_batch_and_defers_cursor_close() -> None: connection = _AdbcStreamingConnection() driver = AdbcDriver(cast("AdbcConnection", connection), dialect="sqlite") @@ -104,15 +337,16 @@ def test_load_from_arrow_passes_record_batch_reader_to_adbc_ingest() -> None: assert job.telemetry["destination"] == "target_table" -def test_select_to_arrow_postgres_rolls_back_on_error() -> None: +@pytest.mark.parametrize("return_format", ["table", "reader"]) +@pytest.mark.parametrize("primary", [RuntimeError("postgres failure"), KeyboardInterrupt("query interrupted")]) +def test_select_to_arrow_postgres_rolls_back_on_error(return_format: Any, primary: BaseException) -> None: class _PostgresErrorCursor: def __init__(self, conn: "_PostgresErrorConnection") -> None: self.connection = conn self.closed = False def execute(self, sql: str, parameters: object = None) -> None: - msg = "postgres failure" - raise RuntimeError(msg) + raise primary def close(self) -> None: self.closed = True @@ -120,6 +354,7 @@ def close(self) -> None: class _PostgresErrorConnection: def __init__(self) -> None: self.rolled_back = False + self.cursor_obj = _PostgresErrorCursor(self) def adbc_get_info(self) -> dict[str, str]: return {"vendor_name": "PostgreSQL", "driver_name": "adbc_driver_postgresql"} @@ -128,14 +363,16 @@ def rollback(self) -> None: self.rolled_back = True def cursor(self) -> _PostgresErrorCursor: - return _PostgresErrorCursor(self) + return self.cursor_obj connection = _PostgresErrorConnection() driver = AdbcDriver(cast("AdbcConnection", connection), dialect="postgres") - with pytest.raises((RuntimeError, SQLSpecError)): - driver.select_to_arrow("SELECT 1", return_format="table") + with pytest.raises(type(primary)) as caught: + driver.select_to_arrow("SELECT 1", return_format=return_format) + assert caught.value is primary assert connection.rolled_back is True + assert connection.cursor_obj.closed is True def test_select_stream_start_closes_cursor_and_rolls_back_postgres_on_failure() -> None: diff --git a/tests/unit/adapters/test_adbc/test_core.py b/tests/unit/adapters/test_adbc/test_core.py index 80f035874..0fcd4d518 100644 --- a/tests/unit/adapters/test_adbc/test_core.py +++ b/tests/unit/adapters/test_adbc/test_core.py @@ -1,17 +1,22 @@ """Unit tests for ADBC core execute-many helpers.""" +import json from collections.abc import Sequence from decimal import Decimal -from types import SimpleNamespace +from pathlib import Path +from types import NoneType, SimpleNamespace from typing import Any, cast +from unittest.mock import Mock, call from uuid import UUID import pyarrow as pa import pytest -from adbc_driver_manager import AdbcStatusCode, DatabaseError +from adbc_driver_manager import AdbcStatusCode, DatabaseError, NotSupportedError -from sqlspec.adapters.adbc import AdbcDriver +from sqlspec.adapters.adbc import AdbcConfig, AdbcDriver from sqlspec.adapters.adbc import core as adbc_core +from sqlspec.adapters.adbc import driver as adbc_driver_mod +from sqlspec.adapters.adbc._typing import AdbcSessionContext from sqlspec.adapters.adbc.core import ( build_connection_config, collect_rows, @@ -26,7 +31,9 @@ resolve_dialect_from_config, resolve_many_rowcount, ) +from sqlspec.adapters.adbc.type_converter import get_adbc_type_converter from sqlspec.core import SQL, ParameterStyle +from sqlspec.core.compiler import SQLProcessor from sqlspec.dialects import PGTextSearch from sqlspec.exceptions import ( CheckViolationError, @@ -35,6 +42,7 @@ DeadlockError, ForeignKeyViolationError, IntegrityError, + NotFoundError, NotNullViolationError, OperationalError, OperationCancelledError, @@ -156,56 +164,40 @@ def test_collect_rows_uses_precomputed_column_names() -> None: assert column_names == ["id", "name"] -def test_prepare_parameters_with_casts_uses_type_converter_factory(monkeypatch: Any) -> None: +def test_prepare_parameters_with_casts_uses_custom_json_serializer() -> None: statement_config = get_statement_config("postgres") - factory_calls: list[tuple[str, int]] = [] - - class FakeConverter: - def convert_dict(self, value: dict[str, Any]) -> str: - return f"factory:{sorted(value.items())!r}" - - def fake_factory(dialect: str, cache_size: int = 5000) -> FakeConverter: - factory_calls.append((dialect, cache_size)) - return FakeConverter() - - monkeypatch.setattr(adbc_core, "get_adbc_type_converter", fake_factory) - - prepared = prepare_parameters_with_casts( - [{"id": 1}], {}, statement_config, dialect="postgres", json_serializer=lambda value: str(value) + statement_config = statement_config.replace( + parameter_config=statement_config.parameter_config.replace( + json_serializer=lambda value: json.dumps({"custom": value}, separators=(",", ":")) + ) ) - assert prepared == ["factory:[('id', 1)]"] - assert factory_calls == [("postgres", 5000)] - + prepared = prepare_parameters_with_casts([{"id": 1}], {}, statement_config, dialect="postgres", json_serializer=str) -def test_prepare_batch_with_casts_creates_converter_once(monkeypatch: Any) -> None: - statement_config = get_statement_config("postgres") - factory_calls: list[tuple[str, int]] = [] - - class FakeConverter: - def convert_dict(self, value: dict[str, Any]) -> str: - return f"factory:{sorted(value.items())!r}" + assert prepared == ['{"custom":{"id":1}}'] - def fake_factory(dialect: str, cache_size: int = 5000) -> FakeConverter: - factory_calls.append((dialect, cache_size)) - return FakeConverter() - monkeypatch.setattr(adbc_core, "get_adbc_type_converter", fake_factory) +def test_prepare_batch_with_casts_uses_custom_json_serializer_for_each_row() -> None: + statement_config = get_statement_config("postgres") + statement_config = statement_config.replace( + parameter_config=statement_config.parameter_config.replace( + json_serializer=lambda value: json.dumps({"custom": value}, separators=(",", ":")) + ) + ) prepared = prepare_batch_with_casts( [(1, {"name": "alice"}), (2, {"name": "bob"}), (3, {"name": "carol"})], {1: "INT"}, statement_config, dialect="postgres", - json_serializer=lambda value: str(value), + json_serializer=str, ) assert prepared == [ - (1, "factory:[('name', 'alice')]"), - (2, "factory:[('name', 'bob')]"), - (3, "factory:[('name', 'carol')]"), + (1, '{"custom":{"name":"alice"}}'), + (2, '{"custom":{"name":"bob"}}'), + (3, '{"custom":{"name":"carol"}}'), ] - assert factory_calls == [("postgres", 5000)] def test_prepare_parameters_with_casts_supports_subclass_type_dispatch() -> None: @@ -642,3 +634,375 @@ def test_adbc_driver_dispatch_execute_lowercases_oracle_columns() -> None: assert result.column_names == ["user_id", "MixedCase"] assert result.selected_data == [{"user_id": 1, "MixedCase": "a"}] + + +@pytest.mark.parametrize("autocommit", [False, True]) +def test_adbc_begin_with_native_autocommit(autocommit: bool) -> None: + """Real explicit transactions commit/roll back and restore the native autocommit mode.""" + config = AdbcConfig(connection_config={"driver_name": "sqlite", "autocommit": autocommit}) + expected_mode = "true" if autocommit else "false" + with config.provide_session() as session: + native = session.connection.adbc_connection + session.execute("CREATE TABLE transaction_values (value INTEGER)") + session.commit() + session.begin() + begun_mode = native.get_option("adbc.connection.autocommit") + session.execute("INSERT INTO transaction_values (value) VALUES (:value)", value=1) + session.rollback() + assert session.select_value("SELECT COUNT(*) FROM transaction_values") == 0 + assert begun_mode == "false" + assert native.get_option("adbc.connection.autocommit") == expected_mode + session.begin() + session.execute("INSERT INTO transaction_values (value) VALUES (:value)", value=2) + session.commit() + assert session.select_value("SELECT value FROM transaction_values") == 2 + assert native.get_option("adbc.connection.autocommit") == expected_mode + + +@pytest.mark.parametrize("capability", ["missing", "unsupported"]) +def test_adbc_begin_preserves_wrapper_transactions_without_native_getter(capability: str) -> None: + """Older/unsupported getters retain ordinary DB-API transaction capability behavior.""" + native = SimpleNamespace() + if capability == "unsupported": + native.get_option = Mock(side_effect=NotSupportedError("getter unsupported")) + committed, rolled_back = Mock(), Mock() + connection = SimpleNamespace( + adbc_get_info=lambda: {"driver_name": "sqlite"}, adbc_connection=native, commit=committed, rollback=rolled_back + ) + driver = AdbcDriver(cast("Any", connection), dialect="sqlite") + driver.begin() + driver.commit() + driver.begin() + driver.rollback() + committed.assert_called_once_with() + rolled_back.assert_called_once_with() + + +@pytest.mark.parametrize("operation", ["commit", "rollback"]) +@pytest.mark.parametrize("interrupted", [False, True]) +def test_adbc_native_transaction_preserves_primary_on_restoration_failure(operation: str, interrupted: bool) -> None: + """Explicitly switched native transactions preserve primary errors while attempting restoration.""" + connection = Mock() + connection.adbc_get_info.return_value = {"driver_name": "sqlite"} + native = connection.adbc_connection + native.get_option.return_value = "true" + native.set_autocommit.side_effect = [None, RuntimeError("restore failed")] + primary = ( + KeyboardInterrupt("transaction interrupted") + if interrupted + else DatabaseError("transaction failed", status_code=AdbcStatusCode.UNKNOWN) + ) + getattr(native, operation).side_effect = [primary, None] if operation == "rollback" else primary + driver = AdbcDriver(connection, dialect="sqlite") + driver.begin() + expected = KeyboardInterrupt if interrupted else SQLSpecError + with pytest.raises(expected) as caught: + getattr(driver, operation)() + if interrupted: + assert caught.value is primary + else: + assert caught.value.__cause__ is primary + assert native.set_autocommit.call_args_list == [call(False), call(True)] + connection.commit.assert_not_called() + connection.rollback.assert_not_called() + + +@pytest.mark.parametrize("interrupted", [False, True]) +@pytest.mark.parametrize("operation", ["commit", "rollback"]) +def test_adbc_failed_native_rollback_never_commits_pending_rows( + tmp_path: Path, interrupted: bool, operation: str +) -> None: + """A failed rollback and cleanup leave real work uncommitted until an explicit retry.""" + path = str(tmp_path / "rollback.sqlite") + config = AdbcConfig(connection_config={"driver_name": "sqlite", "uri": path, "autocommit": True}) + primary = ( + KeyboardInterrupt(f"{operation} interrupted") + if interrupted + else DatabaseError(f"{operation} failed", status_code=AdbcStatusCode.UNKNOWN) + ) + + class NativeBoundary: + def __init__(self, native: Any) -> None: + self.native = native + self.remaining = 2 + + def get_option(self, key: str) -> str: + return cast("str", self.native.get_option(key)) + + def set_autocommit(self, value: bool) -> None: + self.native.set_autocommit(value) + + def commit(self) -> None: + if operation == "commit" and self.remaining: + self.remaining -= 1 + raise primary + self.native.commit() + + def rollback(self) -> None: + if self.remaining: + self.remaining -= 1 + if self.remaining == 1: + raise primary + raise RuntimeError("cleanup rollback failed") + self.native.rollback() + + class ConnectionBoundary: + def __init__(self, connection: Any) -> None: + self.connection = connection + self.adbc_connection = NativeBoundary(connection.adbc_connection) + + def adbc_get_info(self) -> Any: + return self.connection.adbc_get_info() + + def cursor(self) -> Any: + return self.connection.cursor() + + def commit(self) -> None: + self.connection.commit() + + def rollback(self) -> None: + self.connection.rollback() + + with config.provide_connection() as connection: + driver = AdbcDriver(cast("Any", ConnectionBoundary(connection)), dialect="sqlite") + driver.execute("CREATE TABLE pending_values (value INTEGER)") + driver.begin() + driver.execute("INSERT INTO pending_values (value) VALUES (:value)", value=1) + expected = KeyboardInterrupt if interrupted else SQLSpecError + with pytest.raises(expected) as caught: + getattr(driver, operation)() + if interrupted: + assert caught.value is primary + else: + assert caught.value.__cause__ is primary + observer = AdbcConfig(connection_config={"driver_name": "sqlite", "uri": path, "autocommit": True}) + with observer.provide_session() as session: + assert session.select_value("SELECT COUNT(*) FROM pending_values") == 0 + assert connection.adbc_connection.get_option("adbc.connection.autocommit") == "false" + driver.begin() + driver.rollback() + assert connection.adbc_connection.get_option("adbc.connection.autocommit") == "true" + assert driver.select_value("SELECT COUNT(*) FROM pending_values") == 0 + + +def test_adbc_native_begin_failure_restores_mode_and_preserves_native_cause() -> None: + """A failed native mode switch does not start a logical transaction or hide its error.""" + connection = Mock() + connection.adbc_get_info.return_value = {"driver_name": "sqlite"} + native = connection.adbc_connection + primary = DatabaseError("begin failed", status_code=AdbcStatusCode.UNKNOWN) + native.get_option.return_value = "true" + native.set_autocommit.side_effect = [primary, RuntimeError("restore failed")] + driver = AdbcDriver(connection, dialect="sqlite") + with pytest.raises(SQLSpecError) as caught: + driver.begin() + assert caught.value.__cause__ is primary + assert native.set_autocommit.call_args_list == [call(False), call(True)] + driver.rollback() + connection.rollback.assert_called_once_with() + + +def test_pending_exception_preserves_cause() -> None: + """_check_pending_exception preserves original exception cause.""" + + class MockConn: + def adbc_get_info(self) -> dict[str, str]: + return {"driver_name": "sqlite"} + + driver = AdbcDriver(cast("Any", MockConn())) + orig_exc = DatabaseError("query failed", status_code=AdbcStatusCode.INTERNAL) + mapped_exc = adbc_core.create_mapped_exception(orig_exc) + handler = driver.handle_database_exceptions() + handler.pending_exception = mapped_exc + + with pytest.raises(OperationalError) as exc_info: + driver._check_pending_exception(handler) + + assert exc_info.value is mapped_exc + assert exc_info.value.__cause__ is orig_exc + + +@pytest.mark.parametrize( + ("status_code", "expected_cls"), + [ + (AdbcStatusCode.TIMEOUT, QueryTimeoutError), + (AdbcStatusCode.CANCELLED, OperationCancelledError), + (AdbcStatusCode.UNAUTHENTICATED, PermissionDeniedError), + (AdbcStatusCode.UNAUTHORIZED, PermissionDeniedError), + (AdbcStatusCode.NOT_FOUND, NotFoundError), + (AdbcStatusCode.ALREADY_EXISTS, IntegrityError), + (AdbcStatusCode.INVALID_ARGUMENT, DataError), + (AdbcStatusCode.IO, DatabaseConnectionError), + (AdbcStatusCode.INTEGRITY, IntegrityError), + (AdbcStatusCode.INVALID_STATE, OperationalError), + (AdbcStatusCode.INTERNAL, OperationalError), + (AdbcStatusCode.NOT_IMPLEMENTED, SQLSpecError), + ], +) +def test_adbc_status_code_mappings(status_code: AdbcStatusCode, expected_cls: type[SQLSpecError]) -> None: + """_map_adbc_status_exception maps ADBC status codes to SQLSpecError classes.""" + error = DatabaseError("adbc error", status_code=status_code) + mapped = adbc_core.create_mapped_exception(error) + assert isinstance(mapped, expected_cls) + + +@pytest.mark.parametrize( + ("status_code", "sqlstate", "message", "expected_cls"), + [ + (AdbcStatusCode.INVALID_ARGUMENT, "23505", "duplicate key", UniqueViolationError), + (AdbcStatusCode.INVALID_ARGUMENT, "23502", "null value", NotNullViolationError), + (AdbcStatusCode.INVALID_ARGUMENT, "23514", "check constraint", CheckViolationError), + (AdbcStatusCode.INTEGRITY, "23503", "foreign key constraint", ForeignKeyViolationError), + (AdbcStatusCode.INVALID_ARGUMENT, "42601", "invalid SQL", SQLParsingError), + (AdbcStatusCode.INVALID_ARGUMENT, None, "syntax error near SELEC", SQLParsingError), + (AdbcStatusCode.INTERNAL, None, "ORA-00001: unique constraint violated", UniqueViolationError), + ], +) +def test_adbc_specific_diagnostics_override_generic_status( + status_code: AdbcStatusCode, sqlstate: str | None, message: str, expected_cls: type[SQLSpecError] +) -> None: + error = DatabaseError(message, status_code=status_code, sqlstate=sqlstate) + + mapped = adbc_core.create_mapped_exception(error) + + assert type(mapped) is expected_cls + assert mapped.__cause__ is error + + +def test_duckdb_copy_ast_transformer_caching() -> None: + """DuckDB statement config omits ast_transformer so compiled SQL can be cached.""" + config = get_statement_config("duckdb").replace(enable_caching=True) + assert config.parameter_config.ast_transformer is None + + stmt = SQL("SELECT 1 WHERE x = :val", {"val": 42}, statement_config=config) + compiled_sql, _ = stmt.compile() + assert compiled_sql + + processor = SQLProcessor(config) + processor.compile("SELECT 1 WHERE x = :val", {"val": 42}) + processor.compile("SELECT 1 WHERE x = :val", {"val": 99}) + assert processor._cache_hits == 1 + + +def test_rewrite_postgres_uuid_placeholders() -> None: + """_rewrite_postgres_uuid_placeholders rewrites positional parameters to explicit UUID casts.""" + sql = "SELECT $1, $2" + rewritten, scalars, arrays = adbc_core._rewrite_postgres_uuid_placeholders(sql, (1,), (2,), "postgres") + assert "UUID[]" in rewritten + assert scalars == (1,) + assert arrays == (2,) + + +def test_repeated_uuid_bindings_preserve_sql_and_current_parameter_values() -> None: + """Repeated statement rewrites retain each call's current UUID binding.""" + sql = "SELECT $1 WHERE id = $2" + first = UUID("12345678-1234-5678-1234-567812345678") + second = UUID("87654321-4321-8765-4321-876543218765") + first_sql, first_parameters = prepare_postgres_uuid_bindings(sql, [first, 1], is_many=False, dialect="postgres") + second_sql, second_parameters = prepare_postgres_uuid_bindings(sql, [second, 2], is_many=False, dialect="postgres") + assert first_sql == second_sql == "SELECT CAST($1 AS UUID) WHERE id = $2" + assert first_parameters == [str(first), 1] + assert second_parameters == [str(second), 2] + + +def test_strict_coercion_types() -> None: + """_STRICT_COERCION_ALLOWED_TYPES contains NoneType bound once at module top level.""" + assert NoneType in adbc_core._STRICT_COERCION_ALLOWED_TYPES + + +def test_type_converter_caching() -> None: + """get_adbc_type_converter caches output converter instances.""" + conv1 = get_adbc_type_converter("postgres") + conv2 = get_adbc_type_converter("postgres") + assert conv1 is conv2 + + +def test_bytearray_memoryview_coercion() -> None: + """_BASE_TYPE_COERCION_MAP coerces bytearray and memoryview to bytes.""" + coercion = adbc_core._BASE_TYPE_COERCION_MAP + assert coercion[bytearray] is bytes + assert coercion[memoryview] is bytes + + +def test_session_context_driver() -> None: + """AdbcSessionContext uses provided driver_class without lazy import.""" + driver_instance = Mock() + mock_driver_cls = Mock(return_value=driver_instance) + mock_conn = Mock() + ctx = AdbcSessionContext( + acquire_connection=lambda: mock_conn, + release_connection=lambda _conn, **_kw: None, + statement_config=get_statement_config("postgres"), + driver_features={}, + prepare_driver=lambda d: d, + driver_class=cast("type[AdbcDriver]", mock_driver_cls), + ) + with ctx as driver: + assert driver is driver_instance + mock_driver_cls.assert_called_once() + + +def test_select_stream_chunk_size() -> None: + """AdbcSelectStreamSource buffers and returns chunks of specified chunk_size.""" + batch = pa.record_batch([pa.array([1, 2, 3, 4, 5])], names=["id"]) + reader = iter([batch]) + + class _MockDriver: + driver_features: dict[str, Any] = {} + _dialect_name = "sqlite" + + def with_cursor(self, _conn: Any) -> Any: + return Mock() + + def handle_database_exceptions(self) -> Any: + return Mock() + + def _check_pending_exception(self, _handler: Any) -> None: + pass + + source = adbc_driver_mod.AdbcSelectStreamSource(cast("Any", _MockDriver()), "SELECT 1", (), chunk_size=2) + source._reader = reader + chunk1 = source.fetch_chunk() + assert len(chunk1) == 2 + chunk2 = source.fetch_chunk() + assert len(chunk2) == 2 + chunk3 = source.fetch_chunk() + assert len(chunk3) == 1 + assert source.fetch_chunk() == [] + + +def test_select_to_arrow_lowercase_columns() -> None: + """select_to_arrow normalizes reader schema to lowercase when enable_lowercase_column_names is set.""" + table = pa.table({"UPPER_COL": [1, 2]}) + mock_cursor = Mock() + mock_cursor.execute.return_value = None + mock_cursor.fetch_record_batch.return_value = table.to_reader() + mock_conn = Mock() + mock_conn.cursor.return_value = mock_cursor + + driver = AdbcDriver( + mock_conn, + statement_config=get_statement_config("oracle"), + driver_features={"enable_lowercase_column_names": True}, + ) + result = driver.select_to_arrow("SELECT UPPER_COL FROM t", return_format="reader") + reader = result.data + assert reader.schema.names == ["upper_col"] + assert reader.read_all().to_pylist() == [{"upper_col": 1}, {"upper_col": 2}] + + +def test_adbc_exception_isolation() -> None: + """AdbcSelectStreamSource routes execution through exception handler with isolation check.""" + mock_conn = Mock() + mock_cursor = Mock() + mock_conn.cursor.return_value = mock_cursor + mock_cursor.execute.return_value = None + batch = pa.record_batch([pa.array([1])], names=["id"]) + + driver = AdbcDriver(mock_conn, statement_config=get_statement_config("sqlite")) + + source = adbc_driver_mod.AdbcSelectStreamSource(driver, "SELECT 1", ()) + with pytest.MonkeyPatch.context() as mp: + mp.setattr(adbc_driver_mod, "_fetch_record_batch", lambda _c: iter([batch])) + source.start() + assert source._reader is not None diff --git a/tests/unit/adapters/test_adbc/test_data_dictionary.py b/tests/unit/adapters/test_adbc/test_data_dictionary.py index eb0d006a6..dc2349b34 100644 --- a/tests/unit/adapters/test_adbc/test_data_dictionary.py +++ b/tests/unit/adapters/test_adbc/test_data_dictionary.py @@ -10,6 +10,7 @@ from adbc_driver_manager import NotSupportedError as AdbcNotSupportedError +from sqlspec.adapters.adbc import AdbcConfig from sqlspec.adapters.adbc.data_dictionary import ( AdbcDataDictionary, _arrow_type_to_sql, @@ -18,7 +19,7 @@ _normalize_native_statistics, _normalize_native_tables, ) -from sqlspec.data_dictionary import MetadataFidelity, MetadataSource, MetadataSupport +from sqlspec.data_dictionary import DDLResult, MetadataFidelity, MetadataSource, MetadataSupport from sqlspec.exceptions import OperationalError SQLITE_OBJECTS_PAYLOAD: list[dict[str, Any]] = [ @@ -200,9 +201,7 @@ def test_get_foreign_keys_empty_native_result_falls_back() -> None: "catalog_db_schemas": [ { "db_schema_name": "main", - "db_schema_tables": [ - {"table_name": "t2", "table_type": "table", "table_columns": [], "table_constraints": []} - ], + "db_schema_tables": [{"table_name": "t2", "table_type": "table", "table_columns": []}], } ], } @@ -216,6 +215,38 @@ def test_get_foreign_keys_empty_native_result_falls_back() -> None: driver.select.assert_called_once() +def test_empty_foreign_keys_not_failed() -> None: + """Tables with empty constraints list should return empty list without falling back.""" + driver = Mock() + driver.dialect = "duckdb" + driver.connection.adbc_get_objects.return_value = _make_reader([ + { + "catalog_name": "memory", + "catalog_db_schemas": [ + { + "db_schema_name": "main", + "db_schema_tables": [ + {"table_name": "t2", "table_type": "table", "table_columns": [], "table_constraints": []} + ], + } + ], + } + ]) + + result = AdbcDataDictionary().get_foreign_keys(driver, table="t2") + + assert result == [] + driver.select.assert_not_called() + + +def test_snowflake_indexes_fallback() -> None: + """Snowflake dialect returns empty indexes without raising SQLFileNotFoundError.""" + driver = Mock() + driver.dialect = "snowflake" + result = AdbcDataDictionary().get_indexes(driver, table="t1") + assert result == [] + + def test_get_columns_schema_wide_incomplete_falls_back() -> None: """Schema-wide native column discovery should fall back when type names are incomplete.""" driver = Mock() @@ -472,20 +503,19 @@ def test_adbc_get_objects_unique_check_constraints_are_lossy() -> None: assert all(item.attributes["is_lossy"] is True for item in constraint_items) -def test_adbc_ddl_is_unsupported_without_dialect_pack() -> None: - """ADBC transport metadata should not claim lossless DDL.""" - driver = Mock() - driver.dialect = "duckdb" - - result = AdbcDataDictionary().get_ddl( - driver, - "items", - schema="main", - object_type="view", - include_dependencies=False, - prefer_native=False, - redact=False, - ) +def test_adbc_view_ddl_is_explicitly_unsupported() -> None: + """ADBC transport metadata should not claim lossless view DDL.""" + config = AdbcConfig(connection_config={"driver_name": "sqlite"}) + with config.provide_session() as driver: + result = driver.data_dictionary.get_ddl( + driver, + "items", + schema="main", + object_type="view", + include_dependencies=False, + prefer_native=False, + redact=False, + ) assert result.identity.name == "items" assert result.identity.object_type == "view" @@ -497,6 +527,47 @@ def test_adbc_ddl_is_unsupported_without_dialect_pack() -> None: assert result.warnings +@pytest.mark.parametrize("target_dialect", [None, "postgres"]) +def test_adbc_table_and_schema_ddl_preserve_detected_dialect_and_quoted_names(target_dialect: str | None) -> None: + """DDL orchestration uses the active ADBC dialect for quoted catalog names and target rendering.""" + config = AdbcConfig(connection_config={"driver_name": "sqlite"}) + with config.provide_session() as driver: + driver.execute('CREATE TABLE "Mixed Items" ("Mixed ID" INTEGER PRIMARY KEY, "Label" TEXT)') + dictionary = driver.data_dictionary + table = dictionary.get_ddl( + driver, + '"Mixed Items"', + schema='"main"', + include_dependencies=False, + prefer_native=False, + target_dialect=target_dialect, + ) + schema = dictionary.get_schema_ddl( + driver, + schema='"main"', + include_domains=("tables",), + include_dependencies=False, + prefer_native=False, + target_dialect=target_dialect, + ) + dialect = target_dialect or "sqlite" + assert table.status == MetadataSupport.SUPPORTED + assert (table.identity.name, table.identity.schema, table.identity.dialect) == ("Mixed Items", "main", dialect) + assert isinstance(table.ddl, str) + assert '"Mixed Items"' in table.ddl + assert '"Mixed ID"' in table.ddl + assert len(schema.items) == 1 + item = schema.items[0] + assert isinstance(item, DDLResult) + assert (item.identity.name, item.identity.schema, item.identity.dialect, item.status) == ( + "Mixed Items", + "main", + dialect, + MetadataSupport.SUPPORTED, + ) + assert item.ddl == table.ddl + + def test_get_statistics_uses_native_statistic_names_when_available() -> None: """ADBC statistic-name metadata should replace built-in fallback labels.""" driver = Mock() diff --git a/tests/unit/adapters/test_adbc/test_extension_stores.py b/tests/unit/adapters/test_adbc/test_extension_stores.py index d328f1e24..c4b89a785 100644 --- a/tests/unit/adapters/test_adbc/test_extension_stores.py +++ b/tests/unit/adapters/test_adbc/test_extension_stores.py @@ -8,7 +8,8 @@ from sqlspec.adapters.adbc.config import AdbcConfig from sqlspec.adapters.adbc.events.store import AdbcEventQueueStore -from sqlspec.adapters.adbc.litestar.store import ADBCStore +from sqlspec.adapters.adbc.litestar.store import AdbcStore +from sqlspec.exceptions import IntegrityError __all__ = () @@ -25,7 +26,7 @@ def test_adbc_litestar_store_enterprise_dialect_primitives( dialect: str, expected_placeholder: str, expected_ts_expr: str, expected_ddl_fragment: str ) -> None: - store = ADBCStore(_make_adbc_config("sqlite:///:memory:")) + store = AdbcStore(_make_adbc_config("sqlite:///:memory:")) store._dialect = "mssql" if dialect == "tsql" else dialect assert store._get_param_placeholder(1) == expected_placeholder @@ -36,7 +37,7 @@ def test_adbc_litestar_store_enterprise_dialect_primitives( def test_adbc_litestar_store_normalize_tsql_dialect() -> None: driver_stub = SimpleNamespace(dialect="tsql") config_stub = SimpleNamespace(extension_config={"litestar": {}}, provide_session=lambda: nullcontext(driver_stub)) - store = ADBCStore(cast("Any", config_stub)) + store = AdbcStore(cast("Any", config_stub)) assert store._get_dialect() == "mssql" @@ -45,7 +46,7 @@ def test_adbc_litestar_store_normalize_tsql_dialect() -> None: [("oracle", "SQLCODE != -942"), ("db2", "SYSCAT.TABLES"), ("mssql", "IF OBJECT_ID")], ) def test_adbc_litestar_store_enterprise_drop_sql(dialect: str, expected_drop_fragment: str) -> None: - store = ADBCStore(_make_adbc_config("sqlite:///:memory:")) + store = AdbcStore(_make_adbc_config("sqlite:///:memory:")) store._dialect = dialect drop_sql = "\n".join(store._drop_table_sql()) assert expected_drop_fragment in drop_sql @@ -67,7 +68,7 @@ def commit(self) -> None: pass config_stub = SimpleNamespace(extension_config={"litestar": {}}, provide_session=lambda: nullcontext(_DriverStub())) - store = ADBCStore(cast("Any", config_stub)) + store = AdbcStore(cast("Any", config_stub)) store._dialect = dialect store._set("sess_1", b"payload", expires_in=None) @@ -98,5 +99,55 @@ def test_adbc_event_queue_store_enterprise_dialects( assert expected_drop_fragment in drop_sql +def test_adbc_litestar_store_concurrent_set_fallback() -> None: + """AdbcStore._set catches IntegrityError on INSERT race condition and falls back to UPDATE.""" + executed: list[tuple[str, tuple[Any, ...]]] = [] + + class _DriverStub: + def select_one(self, sql: str, *params: Any) -> dict[str, Any]: + executed.append((sql.strip(), params)) + return {"count": 0} + + def execute(self, sql: str, *params: Any) -> None: + executed.append((sql.strip(), params)) + if "INSERT INTO" in sql: + raise IntegrityError("duplicate key value violates unique constraint") + + def commit(self) -> None: + pass + + config_stub = SimpleNamespace(extension_config={"litestar": {}}, provide_session=lambda: nullcontext(_DriverStub())) + store = AdbcStore(cast("Any", config_stub)) + store._dialect = "oracle" + store._set("sess_1", b"payload", expires_in=None) + + assert any("UPDATE" in item[0] for item in executed) + + +def test_adbc_litestar_store_rollback_on_failure() -> None: + """Mutating operations on AdbcStore invoke rollback on execution error.""" + rolled_back = False + + class _DriverStub: + def execute(self, _sql: str, *_params: Any) -> None: + raise RuntimeError("disk failure") + + def commit(self) -> None: + pass + + def rollback(self) -> None: + nonlocal rolled_back + rolled_back = True + + config_stub = SimpleNamespace(extension_config={"litestar": {}}, provide_session=lambda: nullcontext(_DriverStub())) + store = AdbcStore(cast("Any", config_stub)) + store._dialect = "oracle" + + with pytest.raises(RuntimeError, match="disk failure"): + store._delete("sess_1") + + assert rolled_back is True + + def _make_adbc_config(dialect_uri: str) -> AdbcConfig: return AdbcConfig(connection_config={"uri": dialect_uri}) diff --git a/tests/unit/adapters/test_adbc/test_pool.py b/tests/unit/adapters/test_adbc/test_pool.py new file mode 100644 index 000000000..c395864ae --- /dev/null +++ b/tests/unit/adapters/test_adbc/test_pool.py @@ -0,0 +1,484 @@ +"""Exclusive native-shaped ADBC pool ownership contracts.""" + +import os +import threading +import time +from concurrent.futures import ThreadPoolExecutor +from unittest.mock import MagicMock, call + +import pytest + +from sqlspec.adapters.adbc.pool import AdbcConnectionPool +from sqlspec.exceptions import ConnectionTimeoutError, DatabaseConnectionError, ImproperConfigurationError + + +def test_adbc_pool_reuses_only_after_full_native_session_reset() -> None: + """A clean lease resets native session state before another borrower can reuse it.""" + seed, lease, replacement = MagicMock(), MagicMock(), MagicMock() + lease.autocommit = False + seed.adbc_clone.side_effect = [lease, replacement] + pool = AdbcConnectionPool(lambda: seed, max_size=1, timeout=0) + assert pool.acquire() is lease + pool.release(lease) + assert lease.mock_calls == [ + call.rollback(), + call.adbc_connection.set_autocommit(True), + call.cursor(), + call.cursor().execute("DISCARD ALL"), + call.cursor().close(), + call.adbc_connection.set_autocommit(False), + ] + assert lease.autocommit is False + assert pool.acquire() is lease + pool.release(lease, discard=True) + assert pool.acquire() is replacement + pool.release(replacement) + assert pool.close() is True + + +@pytest.mark.parametrize("stage", ["rollback", "enable", "cursor", "discard", "cursor_close", "restore"]) +@pytest.mark.parametrize("close_fails", [False, True]) +def test_adbc_pool_failed_reset_evicts_and_preserves_primary(stage: str, close_fails: bool) -> None: + """Each reset failure evicts; a secondary native-close failure stays retryable.""" + seed, lease, replacement = MagicMock(), MagicMock(), MagicMock() + primary = RuntimeError(f"reset {stage} failed") + seed.adbc_clone.side_effect = [lease, replacement] + cursor = lease.cursor.return_value + handle = lease.adbc_connection + if stage == "rollback": + lease.rollback.side_effect = primary + elif stage == "enable": + handle.set_autocommit.side_effect = [primary, None] + elif stage == "cursor": + lease.cursor.side_effect = primary + elif stage == "discard": + cursor.execute.side_effect = primary + elif stage == "cursor_close": + cursor.close.side_effect = primary + else: + handle.set_autocommit.side_effect = [None, primary] + if close_fails: + lease.close.side_effect = [KeyboardInterrupt("close interrupted"), None] + pool = AdbcConnectionPool(lambda: seed, max_size=1, timeout=0) + assert pool.acquire() is lease + with pytest.raises(RuntimeError) as caught: + pool.release(lease) + assert caught.value is primary + assert handle.set_autocommit.call_args == call(False) + if stage in {"discard", "cursor_close", "restore"}: + cursor.close.assert_called_once_with() + if close_fails: + with pytest.raises(ConnectionTimeoutError): + pool.acquire() + seed.close.assert_not_called() + else: + assert pool.acquire() is replacement + pool.release(replacement) + assert pool.close() is True + assert lease.close.call_count == (2 if close_fails else 1) + + +def test_adbc_pool_reset_primary_survives_all_secondary_cleanup_failures() -> None: + """Cursor, autocommit and native-close cleanup cannot replace a reset exception.""" + seed, lease = MagicMock(), MagicMock() + primary = KeyboardInterrupt("DISCARD interrupted") + cleanup = RuntimeError("cleanup failed") + seed.adbc_clone.return_value = lease + lease.cursor.return_value.execute.side_effect = primary + lease.cursor.return_value.close.side_effect = cleanup + lease.adbc_connection.set_autocommit.side_effect = [None, cleanup] + lease.close.side_effect = [cleanup, None] + pool = AdbcConnectionPool(lambda: seed, max_size=1, timeout=0) + assert pool.acquire() is lease + with pytest.raises(KeyboardInterrupt) as caught: + pool.release(lease) + assert caught.value is primary + lease.cursor.return_value.close.assert_called_once_with() + assert lease.adbc_connection.set_autocommit.call_args == call(False) + seed.close.assert_not_called() + assert pool.close() is True + + +@pytest.mark.parametrize("rollback_fails", [False, True]) +def test_adbc_pool_discard_rolls_back_and_closes_without_reset(rollback_fails: bool) -> None: + """Body-failure disposal attempts rollback and close without DISCARD ALL or reuse.""" + seed, lease, replacement = MagicMock(), MagicMock(), MagicMock() + primary = RuntimeError("rollback failed") + seed.adbc_clone.side_effect = [lease, replacement] + if rollback_fails: + lease.rollback.side_effect = primary + lease.close.side_effect = [KeyboardInterrupt("close interrupted"), None] + pool = AdbcConnectionPool(lambda: seed, max_size=1, timeout=0) + assert pool.acquire() is lease + if rollback_fails: + with pytest.raises(RuntimeError) as caught: + pool.release(lease, discard=True) + assert caught.value is primary + with pytest.raises(ConnectionTimeoutError): + pool.acquire() + else: + pool.release(lease, discard=True) + assert pool.acquire() is replacement + pool.release(replacement) + lease.rollback.assert_called_once_with() + lease.cursor.assert_not_called() + lease.adbc_connection.set_autocommit.assert_not_called() + assert pool.close() is True + + +@pytest.mark.parametrize("mode", ["reuse", "shutdown", "shutdown_close_failure"]) +def test_adbc_pool_reset_retains_capacity_and_obeys_shutdown(mode: str) -> None: + """A resetting lease stays owned until completion and never becomes idle after shutdown.""" + seed, lease = MagicMock(), MagicMock() + entered, proceed = threading.Event(), threading.Event() + cleanup_error = RuntimeError("shutdown close failed") + shutdown = mode != "reuse" + close_fails = mode == "shutdown_close_failure" + + def discard(_sql: str) -> None: + entered.set() + assert proceed.wait(3) + + seed.adbc_clone.return_value = lease + lease.cursor.return_value.execute.side_effect = discard + if close_fails: + lease.close.side_effect = [cleanup_error, None] + pool = AdbcConnectionPool(lambda: seed, max_size=1, timeout=0) + assert pool.acquire() is lease + with ThreadPoolExecutor(max_workers=1) as executor: + resetting = executor.submit(pool.release, lease) + assert entered.wait(3) + try: + with pytest.raises(ConnectionTimeoutError): + pool.acquire() + if shutdown: + assert pool.close() is False + seed.close.assert_not_called() + finally: + proceed.set() + if close_fails: + with pytest.raises(RuntimeError) as caught: + resetting.result(timeout=3) + assert caught.value is cleanup_error + assert pool.is_closed is False + seed.close.assert_not_called() + assert pool.close() is True + else: + resetting.result(timeout=3) + if shutdown: + assert pool.is_closed is True + with pytest.raises(DatabaseConnectionError): + pool.acquire() + else: + assert pool.acquire() is lease + pool.release(lease) + assert pool.close() is True + seed.adbc_clone.assert_called_once_with() + assert lease.close.call_count == (2 if close_fails else 1) + seed.close.assert_called_once_with() + + +def test_adbc_pool_exclusive_original_seed_leases_and_nested_timeout() -> None: + """Every borrower owns a distinct original-seed clone within the configured bound.""" + seed, first, second, third = MagicMock(), MagicMock(), MagicMock(), MagicMock() + seed.adbc_clone.side_effect = [first, second, third] + create = MagicMock(return_value=seed) + pool = AdbcConnectionPool(create, max_size=2, timeout=0) + create.assert_not_called() + assert pool.acquire() is first + assert pool.acquire() is second + with pytest.raises(ConnectionTimeoutError): + pool.acquire() + pool.release(first, discard=True) + assert pool.acquire() is third + for connection in (second, third): + pool.release(connection, discard=True) + create.assert_called_once_with() + assert seed.adbc_clone.call_count == 3 + first.close.assert_called_once_with() + assert pool.close() is True + assert pool.is_closed is True + seed.close.assert_called_once_with() + + +@pytest.mark.parametrize("max_size", [0, -1, True, 1.5]) +def test_adbc_pool_rejects_invalid_capacity(max_size: object) -> None: + """Capacity is a positive integer and rejects bool before native acquisition.""" + create = MagicMock() + with pytest.raises(ImproperConfigurationError): + AdbcConnectionPool(create, max_size=max_size) + create.assert_not_called() + + +@pytest.mark.parametrize("timeout", [-1, float("inf"), float("nan")]) +def test_adbc_pool_rejects_invalid_timeout(timeout: float) -> None: + """Timeout is finite and nonnegative.""" + with pytest.raises(ImproperConfigurationError): + AdbcConnectionPool(MagicMock(), timeout=timeout) + + +@pytest.mark.parametrize("failure_stage", ["create", "clone"]) +def test_adbc_pool_failed_acquisition_recovers_reserved_capacity(failure_stage: str) -> None: + """A failed native open or clone releases its capacity reservation.""" + seed, lease = MagicMock(), MagicMock() + failure = RuntimeError("native acquisition failed") + create = MagicMock(return_value=seed) + seed.adbc_clone.return_value = lease + if failure_stage == "create": + create.side_effect = [failure, seed] + else: + seed.adbc_clone.side_effect = [failure, lease] + pool = AdbcConnectionPool(create, max_size=1, timeout=0) + with pytest.raises(RuntimeError) as caught: + pool.acquire() + assert caught.value is failure + assert pool.acquire() is lease + pool.release(lease) + assert pool.close() is True + + +def test_adbc_pool_rejects_foreign_and_duplicate_release() -> None: + """Only the originating pool may release a currently active lease once.""" + seed, lease, foreign = MagicMock(), MagicMock(), MagicMock() + seed.adbc_clone.return_value = lease + pool = AdbcConnectionPool(lambda: seed) + assert pool.acquire() is lease + with pytest.raises(DatabaseConnectionError): + pool.release(foreign) + foreign.close.assert_not_called() + pool.release(lease) + with pytest.raises(DatabaseConnectionError): + pool.release(lease) + assert pool.close() is True + lease.close.assert_called_once_with() + + +def test_adbc_pool_shutdown_drains_active_lease_without_forcing_close() -> None: + """Terminal shutdown rejects new borrowers and leaves an executing lease alone.""" + seed, lease = MagicMock(), MagicMock() + seed.adbc_clone.return_value = lease + pool = AdbcConnectionPool(lambda: seed) + assert pool.acquire() is lease + assert pool.close() is False + assert pool.is_closed is False + lease.close.assert_not_called() + seed.close.assert_not_called() + with pytest.raises(DatabaseConnectionError): + pool.acquire() + pool.release(lease) + assert pool.is_closed is True + assert pool.close() is True + lease.close.assert_called_once_with() + seed.close.assert_called_once_with() + + +def test_adbc_pool_quarantined_close_retains_capacity_and_retries_seed_last() -> None: + """A failed close retains the native lease until a later shutdown retries it.""" + seed, lease = MagicMock(), MagicMock() + failure = RuntimeError("native close failed") + seed.adbc_clone.return_value = lease + lease.close.side_effect = [failure, None] + pool = AdbcConnectionPool(lambda: seed, max_size=1, timeout=0) + assert pool.acquire() is lease + with pytest.raises(RuntimeError) as caught: + pool.release(lease, discard=True) + assert caught.value is failure + with pytest.raises(ConnectionTimeoutError): + pool.acquire() + seed.close.assert_not_called() + assert pool.close() is True + assert lease.close.call_count == 2 + seed.close.assert_called_once_with() + + +def test_adbc_pool_failed_seed_close_is_retryable() -> None: + """CLOSED is published only after the original seed actually closes.""" + seed, lease = MagicMock(), MagicMock() + failure = RuntimeError("seed close failed") + seed.adbc_clone.return_value = lease + seed.close.side_effect = [failure, None] + pool = AdbcConnectionPool(lambda: seed) + assert pool.acquire() is lease + pool.release(lease) + with pytest.raises(RuntimeError) as caught: + pool.close() + assert caught.value is failure + assert pool.is_closed is False + assert pool.close() is True + assert seed.close.call_count == 2 + + +def test_adbc_pool_shutdown_during_clone_closes_clone_before_seed() -> None: + """Shutdown retains an in-flight reservation and closes its eventual clone first.""" + seed, lease = MagicMock(), MagicMock() + entered, proceed = threading.Event(), threading.Event() + closed: list[str] = [] + + def clone() -> MagicMock: + entered.set() + assert proceed.wait(3) + return lease + + seed.adbc_clone.side_effect = clone + seed.close.side_effect = lambda: closed.append("seed") + lease.close.side_effect = lambda: closed.append("lease") + pool = AdbcConnectionPool(lambda: seed, max_size=1) + with ThreadPoolExecutor(max_workers=2) as executor: + acquiring = executor.submit(pool.acquire) + assert entered.wait(3) + try: + assert executor.submit(pool.close).result(timeout=2) is False + assert closed == [] + finally: + proceed.set() + with pytest.raises(DatabaseConnectionError): + acquiring.result(timeout=3) + assert closed == ["lease", "seed"] + assert pool.is_closed is True + + +def test_adbc_pool_native_io_does_not_hold_bookkeeping_lock() -> None: + """Native open, clone and close permit another thread to inspect pool state.""" + seed, lease = MagicMock(), MagicMock() + with ThreadPoolExecutor(max_workers=1) as executor: + + def inspect_pool() -> None: + assert executor.submit(lambda: pool.is_closed).result(timeout=2) is False + + def create() -> MagicMock: + inspect_pool() + return seed + + def clone() -> MagicMock: + inspect_pool() + return lease + + seed.adbc_clone.side_effect = clone + seed.close.side_effect = inspect_pool + lease.close.side_effect = inspect_pool + lease.rollback.side_effect = inspect_pool + lease.adbc_connection.set_autocommit.side_effect = lambda _enabled: inspect_pool() + lease.cursor.return_value.execute.side_effect = lambda _sql: inspect_pool() + lease.cursor.return_value.close.side_effect = inspect_pool + pool = AdbcConnectionPool(create) + assert pool.acquire() is lease + pool.release(lease) + assert pool.close() is True + + +def test_adbc_pool_serializes_original_seed_cloning() -> None: + """Concurrent checkouts serialize original-seed cloning and yield distinct handles.""" + seed = MagicMock() + entered, proceed = threading.Event(), threading.Event() + start = threading.Barrier(3, timeout=3) + leases: list[MagicMock] = [] + native_guard = threading.Lock() + + def clone() -> MagicMock: + assert native_guard.acquire(blocking=False), "concurrent original-seed clone" + try: + entered.set() + assert proceed.wait(3) + lease = MagicMock() + lease.adbc_clone.side_effect = AssertionError("cloned a clone") + leases.append(lease) + return lease + finally: + native_guard.release() + + def acquire() -> object: + start.wait() + return pool.acquire() + + seed.adbc_clone.side_effect = clone + pool = AdbcConnectionPool(lambda: seed, max_size=2, timeout=0) + with ThreadPoolExecutor(max_workers=2) as executor: + first, second = executor.submit(acquire), executor.submit(acquire) + start.wait() + assert entered.wait(3) + proceed.set() + first_lease, second_lease = first.result(timeout=3), second.result(timeout=3) + assert first_lease is not second_lease + assert len(leases) == 2 + for lease in leases: + pool.release(lease) + assert pool.close() is True + + +def test_adbc_pool_closing_lease_retains_capacity_until_native_close_finishes() -> None: + """Native cleanup remains owned and counted while another thread is closing it.""" + seed, lease = MagicMock(), MagicMock() + entered, proceed = threading.Event(), threading.Event() + + def native_close() -> None: + entered.set() + assert proceed.wait(3) + + seed.adbc_clone.return_value = lease + lease.close.side_effect = native_close + pool = AdbcConnectionPool(lambda: seed, max_size=1, timeout=0) + assert pool.acquire() is lease + with ThreadPoolExecutor(max_workers=1) as executor: + releasing = executor.submit(pool.release, lease, discard=True) + assert entered.wait(3) + try: + with pytest.raises(ConnectionTimeoutError): + pool.acquire() + assert pool.close() is False + seed.close.assert_not_called() + finally: + proceed.set() + releasing.result(timeout=3) + assert pool.is_closed is True + seed.close.assert_called_once_with() + + +@pytest.mark.parametrize("timeout", [30.0, 1e100]) +def test_adbc_pool_shutdown_wakes_waiting_borrower(monkeypatch: pytest.MonkeyPatch, timeout: float) -> None: + """A borrower waiting for capacity wakes on shutdown rather than its long timeout.""" + seed, lease = MagicMock(), MagicMock() + waiting = threading.Event() + native_clock = time.monotonic + worker_id: list[int] = [] + clock_reads = 0 + + def clock() -> float: + nonlocal clock_reads + if worker_id and threading.get_ident() == worker_id[0]: + clock_reads += 1 + if clock_reads == 2: + waiting.set() + return native_clock() + + def acquire() -> object: + worker_id.append(threading.get_ident()) + return pool.acquire() + + seed.adbc_clone.return_value = lease + pool = AdbcConnectionPool(lambda: seed, max_size=1, timeout=timeout) + assert pool.acquire() is lease + monkeypatch.setattr(time, "monotonic", clock) + with ThreadPoolExecutor(max_workers=1) as executor: + borrower = executor.submit(acquire) + assert waiting.wait(3) + assert pool.close() is False + with pytest.raises(DatabaseConnectionError): + borrower.result(timeout=2) + pool.release(lease) + assert pool.is_closed is True + + +def test_adbc_pool_rejects_fork_before_inherited_locks(monkeypatch: pytest.MonkeyPatch) -> None: + """Every operation refuses a changed PID before touching a poisoned lock or native handle.""" + condition = MagicMock() + condition.__enter__.side_effect = AssertionError("inherited condition accessed") + monkeypatch.setattr(threading, "Condition", lambda *_args: condition) + create, lease = MagicMock(), MagicMock() + pool = AdbcConnectionPool(create) + child_pid = os.getpid() + 1 + monkeypatch.setattr(os, "getpid", lambda: child_pid) + for operation in (pool.acquire, lambda: pool.release(lease), pool.close, lambda: pool.is_closed, pool.check_pid): + with pytest.raises(DatabaseConnectionError, match="fork"): + operation() + create.assert_not_called() + lease.close.assert_not_called() diff --git a/tests/unit/adapters/test_adbc/test_profiling.py b/tests/unit/adapters/test_adbc/test_profiling.py index 231a56384..2900eb960 100644 --- a/tests/unit/adapters/test_adbc/test_profiling.py +++ b/tests/unit/adapters/test_adbc/test_profiling.py @@ -2,12 +2,109 @@ import subprocess import sys +import tempfile +from importlib import import_module from pathlib import Path +from unittest.mock import MagicMock + +import pytest from sqlspec.adapters.adbc import AdbcConfig +from sqlspec.exceptions import SQLSpecError from sqlspec.observability import ObservabilityConfig, StatementEvent +@pytest.mark.parametrize("stage", ["path", "configuration"]) +@pytest.mark.parametrize("borrowed", [False, True]) +@pytest.mark.parametrize("cleanup_fails", [False, True]) +def test_adbc_profiling_setup_failure_releases_owned_connection( + monkeypatch: pytest.MonkeyPatch, stage: str, borrowed: bool, cleanup_fails: bool +) -> None: + """Profiling failures preserve setup errors and close only owned connections.""" + connection = MagicMock() + primary = RuntimeError("profile setup failed") + removed: list[Path] = [] + original_unlink = Path.unlink + + def remove_profile(path: Path, missing_ok: bool = False) -> None: + if path.name.startswith("sqlspec_duckdb_profile_"): + removed.append(path) + if cleanup_fails: + raise KeyboardInterrupt("profile cleanup interrupted") + else: + original_unlink(path, missing_ok=missing_ok) + + def temporary_directory() -> str: + raise primary + + monkeypatch.setattr(import_module("adbc_driver_duckdb.dbapi"), "connect", lambda **_kwargs: connection) + monkeypatch.setattr(Path, "unlink", remove_profile) + if stage == "path": + monkeypatch.setattr(tempfile, "gettempdir", temporary_directory) + else: + connection.cursor.return_value.execute.side_effect = primary + if cleanup_fails: + connection.close.side_effect = SystemExit("connection cleanup interrupted") + config = AdbcConfig( + connection_config={"driver_name": "duckdb"}, + connection_instance=connection if borrowed else None, + driver_features={"enable_statement_profiling": True}, + ) + context = config.provide_session() + with pytest.raises(RuntimeError) as caught: + context.__enter__() + assert caught.value is primary + context.__exit__(None, None, None) + if borrowed: + connection.close.assert_not_called() + else: + connection.close.assert_called_once_with() + if stage == "configuration": + assert len(removed) == 1 + connection.cursor.return_value.close.assert_called_once_with() + + +@pytest.mark.parametrize("setup_fails", [False, True]) +def test_adbc_profiling_cleanup_attempts_both_paths_and_release_once( + monkeypatch: pytest.MonkeyPatch, setup_fails: bool +) -> None: + """A profile cleanup failure clears ownership and still attempts both path cleanups.""" + connection = MagicMock() + primary = KeyboardInterrupt("metadata interrupted") + cleanup_error = RuntimeError("profile cleanup failed") + removed: list[Path] = [] + original_unlink = Path.unlink + + def remove_profile(path: Path, missing_ok: bool = False) -> None: + if path.name.startswith("sqlspec_duckdb_profile_"): + removed.append(path) + raise cleanup_error + original_unlink(path, missing_ok=missing_ok) + + monkeypatch.setattr(import_module("adbc_driver_duckdb.dbapi"), "connect", lambda **_kwargs: connection) + monkeypatch.setattr(Path, "unlink", remove_profile) + if setup_fails: + connection.adbc_get_info.side_effect = primary + config = AdbcConfig( + connection_config={"driver_name": "duckdb"}, driver_features={"enable_statement_profiling": True} + ) + context = config.provide_session() + if setup_fails: + with pytest.raises(KeyboardInterrupt) as setup_caught: + context.__enter__() + assert setup_caught.value is primary + else: + context.__enter__() + with pytest.raises(RuntimeError) as cleanup_caught: + context.__exit__(None, None, None) + assert cleanup_caught.value is cleanup_error + assert len(removed) == 2 + assert removed[0] == removed[1] + context.__exit__(None, None, None) + assert len(removed) == 2 + connection.close.assert_called_once_with() + + def test_adbc_duckdb_enable_statement_profiling_and_execute_script_events() -> None: """AdbcConfig with driver_name='duckdb' and enable_statement_profiling=True populates StatementEvent.profile and emits per-statement script events.""" events: list[StatementEvent] = [] @@ -63,6 +160,24 @@ def test_adbc_duckdb_enable_statement_profiling_and_execute_script_events() -> N assert outer_event.script_ordinal is None +def test_adbc_duckdb_profiles_remain_scoped_after_execution_error() -> None: + """Each successful statement consumes its own profile without leaking a failed query's state.""" + events: list[StatementEvent] = [] + config = AdbcConfig( + connection_config={"driver_name": "duckdb"}, + driver_features={"enable_statement_profiling": True}, + observability_config=ObservabilityConfig(statement_observers=(events.append,)), + ) + with config.provide_session() as session: + assert session.select_value("SELECT 1 AS value") == 1 + with pytest.raises(SQLSpecError): + session.execute("SELECT missing_column FROM missing_table") + session.rollback() + assert session.select_value("SELECT 2 AS value") == 2 + assert [event.sql for event in events] == ["SELECT 1 AS value", "SELECT 2 AS value"] + assert all(isinstance(event.profile, dict) for event in events) + + def test_adbc_duckdb_statement_profiling_without_duckdb_package() -> None: """ADBC statement profiling with DuckDB backend works without the duckdb Python library.""" script = """ diff --git a/tests/unit/adapters/test_adbc/test_uuid_binding.py b/tests/unit/adapters/test_adbc/test_uuid_binding.py index 3d3e18d9c..eb0a6f11d 100644 --- a/tests/unit/adapters/test_adbc/test_uuid_binding.py +++ b/tests/unit/adapters/test_adbc/test_uuid_binding.py @@ -401,8 +401,9 @@ def test_batch_with_inconsistent_row_shapes_is_left_unchanged() -> None: def test_structural_rewrite_cache_is_bounded_and_keyed_by_dialect_and_ordinals() -> None: - rewrite = cast("Any", adbc_core)._rewrite_postgres_uuid_placeholders - rewrite.cache_clear() + rewrite = adbc_core._rewrite_postgres_uuid_placeholders + cache_state = adbc_core._REWRITE_CACHE_STATE + cache_state.clear() first = rewrite("SELECT $1, $2", (1,), (), "postgres") second = rewrite("SELECT $1, $2", (1,), (), "postgres") @@ -412,27 +413,28 @@ def test_structural_rewrite_cache_is_bounded_and_keyed_by_dialect_and_ordinals() assert first == second == ("SELECT CAST($1 AS UUID), $2", (1,), ()) assert third == ("SELECT $1, CAST($2 AS UUID)", (2,), ()) assert fourth == first - assert rewrite.cache_info().hits == 1 - assert rewrite.cache_info().misses == 3 - assert rewrite.cache_info().maxsize == 256 + assert cache_state.hits == 1 + assert cache_state.misses == 3 + assert cache_state.maxsize == 256 def test_structural_rewrite_cache_distinguishes_scalar_from_array_ordinals() -> None: """The same SQL and ordinal produce different rewrites depending on the ordinal's kind.""" - rewrite = cast("Any", adbc_core)._rewrite_postgres_uuid_placeholders - rewrite.cache_clear() + rewrite = adbc_core._rewrite_postgres_uuid_placeholders + cache_state = adbc_core._REWRITE_CACHE_STATE + cache_state.clear() scalar_only = rewrite("SELECT $1", (1,), (), "postgres") array_only = rewrite("SELECT $1", (), (1,), "postgres") assert scalar_only == ("SELECT CAST($1 AS UUID)", (1,), ()) assert array_only == ("SELECT CAST($1 AS UUID[])", (), (1,)) - assert rewrite.cache_info().misses == 2 - assert rewrite.cache_info().hits == 0 + assert cache_state.misses == 2 + assert cache_state.hits == 0 def test_structural_rewrite_uses_ast_and_ignores_literal_and_comment_placeholders() -> None: - rewrite = cast("Any", adbc_core)._rewrite_postgres_uuid_placeholders + rewrite = adbc_core._rewrite_postgres_uuid_placeholders rewritten, effective, effective_arrays = rewrite("SELECT '$1' AS literal, $1 -- $1\n", (1,), (), "postgres") @@ -465,7 +467,7 @@ def test_extension_dialect_operators_survive_uuid_ast_rewrite(dialect: str, sql: def test_structural_rewrite_reports_sqlglot_parse_errors() -> None: - rewrite = cast("Any", adbc_core)._rewrite_postgres_uuid_placeholders + rewrite = adbc_core._rewrite_postgres_uuid_placeholders with pytest.raises(SQLSpecError, match="Failed to parse PostgreSQL ADBC SQL for UUID parameter binding"): rewrite("SELECT (", (1,), (), "postgres") diff --git a/tests/unit/adapters/test_aiomysql/test_adk_store.py b/tests/unit/adapters/test_aiomysql/test_adk_store.py index ea75d6e65..fabfb7049 100644 --- a/tests/unit/adapters/test_aiomysql/test_adk_store.py +++ b/tests/unit/adapters/test_aiomysql/test_adk_store.py @@ -223,6 +223,7 @@ async def test_aiomysql_adk_memory_store_insert_and_search_json_handling() -> No cursor = MagicMock() cursor.rowcount = 1 cursor.execute = AsyncMock() + cursor.executemany = AsyncMock() cursor.description = [ ("id",), ("session_id",), @@ -284,7 +285,7 @@ async def test_aiomysql_adk_memory_store_insert_and_search_json_handling() -> No records = await store.search_entries("hello", "app", "user-1") assert inserted == 1 - insert_params = cursor.execute.call_args_list[0].args[1] + insert_params = cursor.executemany.call_args_list[0].args[1][0] assert insert_params[6] is None assert insert_params[11] is None assert len(records) == 1 @@ -292,6 +293,51 @@ async def test_aiomysql_adk_memory_store_insert_and_search_json_handling() -> No assert records[0]["metadata_json"] == {"source": "unit"} +async def test_aiomysql_insert_memory_entries_batches_with_executemany() -> None: + """insert_memory_entries batches multiple entries in a single cursor.executemany call.""" + from datetime import UTC, datetime + from unittest.mock import AsyncMock + + now = datetime.now(tz=UTC) + cursor = MagicMock() + cursor.rowcount = 2 + cursor.executemany = AsyncMock() + conn = MagicMock() + conn.cursor = AsyncMock(return_value=cursor) + conn.commit = AsyncMock() + conn.__aenter__ = AsyncMock(return_value=conn) + conn.__aexit__ = AsyncMock(return_value=None) + config = _mock_config() + config.provide_connection = lambda *_a, **_k: conn + store = AiomysqlADKMemoryStore(config) + + entries = [ + cast( + "Any", + { + "id": f"mem-{i}", + "session_id": "sess-1", + "app_name": "app", + "user_id": "user-1", + "event_id": f"evt-{i}", + "timestamp": now, + "content_json": {"idx": i}, + "content_text": f"text {i}", + "metadata_json": None, + "inserted_at": now, + }, + ) + for i in range(2) + ] + + inserted = await store.insert_memory_entries(entries) + + assert inserted == 2 + assert cursor.executemany.call_count == 1 + batch_params = cursor.executemany.call_args_list[0].args[1] + assert len(batch_params) == 2 + + async def test_aiomysql_stream_source_closes_cursor_when_execute_raises() -> None: """AiomysqlStreamSource.start() closes the cursor if execute() fails.""" from unittest.mock import AsyncMock diff --git a/tests/unit/adapters/test_aiomysql/test_config.py b/tests/unit/adapters/test_aiomysql/test_config.py index 71925992a..f3c4710ae 100644 --- a/tests/unit/adapters/test_aiomysql/test_config.py +++ b/tests/unit/adapters/test_aiomysql/test_config.py @@ -320,3 +320,31 @@ def test_aiomysql_config_with_dsn_applies_to_connection_kwargs() -> None: assert kwargs["password"] == "pass1" assert kwargs["host"] == "dbhost" assert kwargs["db"] == "production" + + +def test_aiomysql_session_factory_slots() -> None: + """_AiomysqlSessionFactory declares __slots__ for connection context tracking.""" + from sqlspec.adapters.aiomysql.config import _AiomysqlSessionFactory + + assert _AiomysqlSessionFactory.__slots__ == ("_contexts",) + + +async def test_aiomysql_create_connection_closes_socket_on_error(monkeypatch: pytest.MonkeyPatch) -> None: + """create_connection closes the standalone connection if _ensure_connection raises.""" + from unittest.mock import AsyncMock, MagicMock + + import aiomysql + + mock_conn = MagicMock() + mock_conn.close = MagicMock() + monkeypatch.setattr(aiomysql, "connect", AsyncMock(return_value=mock_conn)) + + config = AiomysqlConfig( + connection_config={"host": "localhost", "db": "test"}, + driver_features={"on_connection_create": AsyncMock(side_effect=RuntimeError("hook error"))}, + ) + + with pytest.raises(RuntimeError, match="hook error"): + await config.create_connection() + + assert mock_conn.close.called diff --git a/tests/unit/adapters/test_aiomysql/test_driver.py b/tests/unit/adapters/test_aiomysql/test_driver.py index 304d3f3ec..caefe1d50 100644 --- a/tests/unit/adapters/test_aiomysql/test_driver.py +++ b/tests/unit/adapters/test_aiomysql/test_driver.py @@ -97,3 +97,121 @@ def test_aiomysql_signature_namespace_includes_cursors() -> None: ns = config.get_signature_namespace() assert "AiomysqlDictCursor" in ns assert "AiomysqlRawCursor" in ns + + +async def test_aiomysql_execute_cache_hit_invokes_collect_rows_on_clean_query(monkeypatch: pytest.MonkeyPatch) -> None: + """Cached query without unescaped literal percent delegates to fast path and invokes collect_rows.""" + from unittest.mock import AsyncMock, MagicMock + + from sqlspec.core import OperationProfile, ParameterProfile, ProcessedState + from sqlspec.core.cache import CachedQuery + + cursor = MagicMock() + cursor.description = [("id", 3, None, None, None, None, None)] + cursor.fetchall = AsyncMock(return_value=[(1,)]) + cursor.execute = AsyncMock(return_value=None) + + conn = MagicMock() + conn.cursor = AsyncMock(return_value=cursor) + + driver = AiomysqlDriver(connection=cast("Any", conn)) + collect_called: list[bool] = [] + original_collect_rows = AiomysqlDriver.collect_rows + + def _spy_collect_rows(*args: Any, **kwargs: Any) -> Any: + collect_called.append(True) + return original_collect_rows(*args, **kwargs) + + monkeypatch.setattr(AiomysqlDriver, "collect_rows", _spy_collect_rows) + + cached = CachedQuery( + compiled_sql="SELECT 1", + parameter_profile=ParameterProfile(), + input_named_parameters=(), + applied_wrap_types=False, + parameter_casts={}, + operation_type="SELECT", + operation_profile=OperationProfile(returns_rows=True, modifies_rows=False), + param_count=0, + processed_state=ProcessedState(compiled_sql="SELECT 1", execution_parameters=[], operation_type="SELECT"), + ) + + result = await driver._execute_cache_hit("SELECT 1", (), cached) + assert bool(collect_called) is True + assert result.data == [(1,)] + + +async def test_aiomysql_dispatch_execute_uses_statement_statement_config() -> None: + """dispatch_execute compiles using statement.statement_config.""" + from unittest.mock import AsyncMock, MagicMock + + from sqlspec import SQL + from sqlspec.core import StatementConfig + + custom_config = StatementConfig(enable_caching=False) + statement = SQL("SELECT 1", statement_config=custom_config) + cursor = MagicMock() + cursor.description = [("1", 3, None, None, None, None, None)] + cursor.execute = AsyncMock() + cursor.fetchall = AsyncMock(return_value=[(1,)]) + cursor.rowcount = 1 + cursor.lastrowid = None + + driver = AiomysqlDriver(connection=MagicMock()) + result = await driver.dispatch_execute(cursor, statement) + assert cursor.execute.called + assert result.data_row_count == 1 + + +async def test_aiomysql_dispatch_select_stream_uses_statement_config() -> None: + """dispatch_select_stream compiles using statement.statement_config.""" + from unittest.mock import MagicMock + + from sqlspec import SQL + from sqlspec.core import StatementConfig + + custom_config = StatementConfig(enable_caching=False) + statement = SQL("SELECT 1", statement_config=custom_config) + driver = AiomysqlDriver(connection=MagicMock()) + stream = driver.dispatch_select_stream(statement, chunk_size=10) + assert stream is not None + + +async def test_aiomysql_dispatch_execute_script_delegates_to_super() -> None: + """dispatch_execute_script emits script statement events via super().""" + from unittest.mock import AsyncMock, MagicMock + + from sqlspec import SQL + + cursor = MagicMock() + cursor.execute = AsyncMock() + cursor.description = None + cursor.rowcount = 0 + cursor.lastrowid = None + + driver = AiomysqlDriver(connection=MagicMock()) + statement = SQL("SELECT 1; SELECT 2").as_script() + result = await driver.dispatch_execute_script(cursor, statement) + assert result.statement_count == 2 + assert result.successful_statements == 2 + + +async def test_aiomysql_transaction_methods_catch_pymysql_error() -> None: + """begin, commit, and rollback catch AiomysqlPymysqlError and raise mapped exception.""" + from unittest.mock import AsyncMock, MagicMock + + from sqlspec.adapters.aiomysql._typing import AiomysqlPymysqlError + from sqlspec.exceptions import SQLSpecError + + conn = MagicMock() + conn.begin = AsyncMock(side_effect=AiomysqlPymysqlError(1045, "Access denied")) + conn.commit = AsyncMock(side_effect=AiomysqlPymysqlError(1045, "Access denied")) + conn.rollback = AsyncMock(side_effect=AiomysqlPymysqlError(1045, "Access denied")) + + driver = AiomysqlDriver(connection=conn) + with pytest.raises(SQLSpecError): + await driver.begin() + with pytest.raises(SQLSpecError): + await driver.commit() + with pytest.raises(SQLSpecError): + await driver.rollback() diff --git a/tests/unit/adapters/test_aiomysql/test_litestar_store.py b/tests/unit/adapters/test_aiomysql/test_litestar_store.py index 38a61079c..8d7c74895 100644 --- a/tests/unit/adapters/test_aiomysql/test_litestar_store.py +++ b/tests/unit/adapters/test_aiomysql/test_litestar_store.py @@ -30,3 +30,38 @@ def test_aiomysql_litestar_store_honors_table_options() -> None: ddl = store._table_ddl() assert f"{DEFAULT_ENGINE_CLAUSE} COMMENT='litestar-session'" in ddl + + +def test_aiomysql_litestar_store_drop_table_sql_omits_drop_index() -> None: + """_drop_table_sql must emit only DROP TABLE IF EXISTS without preceding DROP INDEX.""" + store = AiomysqlStore(_mock_config()) + + statements = store._drop_table_sql() + + assert len(statements) == 1 + assert statements[0] == "DROP TABLE IF EXISTS litestar_session" + + +async def test_aiomysql_litestar_store_set_uses_mariadb_compatible_upsert() -> None: + """_set must use VALUES(col) in ON DUPLICATE KEY UPDATE for MySQL and MariaDB compatibility.""" + from unittest.mock import AsyncMock + + cursor = MagicMock() + cursor.execute = AsyncMock() + conn = MagicMock() + conn.cursor = AsyncMock(return_value=cursor) + conn.commit = AsyncMock() + conn.__aenter__ = AsyncMock(return_value=conn) + conn.__aexit__ = AsyncMock(return_value=None) + config = _mock_config() + config.provide_connection = lambda *_a, **_k: conn + + store = AiomysqlStore(config) + await store.set("session-1", b"payload", 300) + + assert cursor.execute.call_count == 1 + sql = cursor.execute.call_args_list[0].args[0] + assert "ON DUPLICATE KEY UPDATE" in sql + assert "data = VALUES(data)" in sql + assert "expires_at = VALUES(expires_at)" in sql + assert "AS new" not in sql diff --git a/tests/unit/adapters/test_aiomysql/test_load_from_arrow.py b/tests/unit/adapters/test_aiomysql/test_load_from_arrow.py index 6c2937303..8fb2715c3 100644 --- a/tests/unit/adapters/test_aiomysql/test_load_from_arrow.py +++ b/tests/unit/adapters/test_aiomysql/test_load_from_arrow.py @@ -122,3 +122,14 @@ async def test_load_from_arrow_binds_the_infile_path_as_a_parameter() -> None: assert "LOAD DATA LOCAL INFILE %s INTO TABLE" in statement assert parameters[0].endswith(".tsv") assert parameters[0] not in statement + + +async def test_load_from_arrow_percent_in_column_escapes_pyformat() -> None: + """Percent signs in table or column identifiers must be escaped for pyformat.""" + conn = _FakeConnection() + driver = _make_driver(conn, enable_local_infile=True) + + job = await driver.load_from_arrow("orders", pa.table({"discount%rate": [1, 2]})) + + assert job.telemetry["rows_processed"] == 2 + assert "(`discount%%rate`)" in conn._cursor.execute_calls[0] diff --git a/tests/unit/adapters/test_aiosqlite/test_adk_config.py b/tests/unit/adapters/test_aiosqlite/test_adk_config.py index 581196667..608cb5a3c 100644 --- a/tests/unit/adapters/test_aiosqlite/test_adk_config.py +++ b/tests/unit/adapters/test_aiosqlite/test_adk_config.py @@ -9,6 +9,7 @@ from sqlspec.adapters.aiosqlite.config import AiosqliteConfig from sqlspec.config import ADKConfig from sqlspec.exceptions import ImproperConfigurationError +from sqlspec.extensions.adk import StoredMemory class AsyncRecordingConnection: @@ -189,3 +190,96 @@ async def test_aiosqlite_list_sessions_rejects_invalid_options_before_connecting await store.list_sessions("app", **options) assert conn.calls == [] + + +async def test_aiosqlite_adk_recycled_connection_reruns_pragma_configuration() -> None: + """Verify WeakSet tracking allows recycled connections to re-run PRAGMAs.""" + import weakref + + config = AiosqliteConfig() + store = AiosqliteADKStore(config) + assert isinstance(store._configured_connections, weakref.WeakSet) + + conn1 = AsyncRecordingConnection() + await store._ensure_pragmas(conn1) + assert len(conn1.statements) == 4 + + await store._ensure_pragmas(conn1) + assert len(conn1.statements) == 4 + + conn2 = AsyncRecordingConnection() + await store._ensure_pragmas(conn2) + assert len(conn2.statements) == 4 + + +async def test_aiosqlite_adk_get_events_with_limit_binds_parameter() -> None: + """Verify get_events parameterizes LIMIT instead of string interpolation.""" + config = AiosqliteConfig(connection_config={"database": ":memory:"}) + store = AiosqliteADKStore(config) + calls: list[tuple[str, tuple[Any, ...]]] = [] + + class _MockEventsCursor: + async def fetchall(self) -> list[Any]: + return [] + + class _MockEventsConnection: + async def execute(self, sql: str, params: Any = None) -> _MockEventsCursor: + calls.append((sql, params)) + return _MockEventsCursor() + + class _AsyncContextManager: + async def __aenter__(self) -> Any: + return _MockEventsConnection() + + async def __aexit__(self, *_: Any) -> None: + return None + + config.provide_connection = MagicMock(return_value=_AsyncContextManager()) # type: ignore[method-assign] + + events = await store.get_events("app", "user", "session", limit=10) + assert events == [] + assert len(calls) >= 1 + sql, params = calls[-1] + assert "LIMIT ?" in sql + assert 10 in params + + +async def test_aiosqlite_adk_memory_store_persists_and_retrieves_float32_embeddings() -> None: + """Verify memory entries persist and retrieve float32 embeddings.""" + from datetime import UTC, datetime + from uuid import uuid4 + + config = AiosqliteConfig(connection_config={"database": ":memory:"}) + store = AiosqliteADKMemoryStore(config) + try: + await store.create_tables() + + now = datetime.now(UTC) + embedding = [0.1, 0.2, 0.3, 0.4] + record = StoredMemory( + id=str(uuid4()), + session_id="s1", + app_name="app", + user_id="u1", + scope="user", + event_id="evt-1", + author="user", + timestamp=now, + content_json={"text": "vector search test"}, + content_text="vector search test", + metadata_json=None, + inserted_at=now, + embedding=embedding, + ) + inserted = await store.insert_memory_entries([record]) + assert inserted == 1 + + results = await store.search_entries(query="vector", app_name="app", user_id="u1", embedding=embedding) + assert len(results) == 1 + assert results[0]["event_id"] == "evt-1" + assert results[0]["embedding"] is not None + assert len(results[0]["embedding"]) == 4 + for orig, rec in zip(embedding, results[0]["embedding"], strict=True): + assert abs(orig - rec) < 1e-6 + finally: + await config.close_pool() diff --git a/tests/unit/adapters/test_aiosqlite/test_config.py b/tests/unit/adapters/test_aiosqlite/test_config.py index a17c48a2a..09f2be0b2 100644 --- a/tests/unit/adapters/test_aiosqlite/test_config.py +++ b/tests/unit/adapters/test_aiosqlite/test_config.py @@ -313,3 +313,29 @@ async def test_create_connection_applies_the_configured_row_factory(tmp_path: Pa await connection.close() finally: await config.close_pool() + + +def test_filter_connection_params_with_hoisted_helper() -> None: + from sqlspec.adapters.aiosqlite.core import _filter_connection_params + + raw = {"database": ":memory:", "extra": {}, "pool_size": 10, "timeout": 30.0, "isolation_level": None} + filtered = _filter_connection_params(raw) + assert filtered == {"database": ":memory:", "timeout": 30.0, "isolation_level": None} + + +async def test_create_pool_preserves_zero_timeouts(tmp_path: Path) -> None: + config = AiosqliteConfig( + connection_config={ + "database": tmp_path / "zero.db", + "connect_timeout": 0.0, + "idle_timeout": 0.0, + "operation_timeout": 0.0, + } + ) + pool = await config._create_pool() + try: + assert pool._connect_timeout == 0.0 + assert pool._idle_timeout == 0.0 + assert pool._operation_timeout == 0.0 + finally: + await pool.close() diff --git a/tests/unit/adapters/test_aiosqlite/test_driver.py b/tests/unit/adapters/test_aiosqlite/test_driver.py index daeb6b842..c75014cfb 100644 --- a/tests/unit/adapters/test_aiosqlite/test_driver.py +++ b/tests/unit/adapters/test_aiosqlite/test_driver.py @@ -428,3 +428,59 @@ async def test_aiosqlite_pool_tracks_in_flight_creations_and_connection_registry finally: await pool.close() assert pool.size() == 0 + + +async def test_aiosqlite_cache_hit_preserves_ambient_exception_and_formats_wire_sql() -> None: + import sys + + connection = await aiosqlite.connect(":memory:") + try: + driver = AiosqliteDriver(connection=connection) + await driver.execute("SELECT 1 AS value") + + ambient_preserved = False + try: + raise RuntimeError("ambient error") + except RuntimeError as exc: + result = await driver.execute("SELECT 1 AS value") + ambient_preserved = sys.exc_info()[1] is exc + assert result.data == [(1,)] + assert ambient_preserved + finally: + await connection.close() + + +async def test_aiosqlite_cursor_lazy_allocation() -> None: + conn = await aiosqlite.connect(":memory:") + try: + driver = AiosqliteDriver(connection=conn) + result = await driver.execute("SELECT 1 AS probe") + assert result.data == [(1,)] + finally: + await conn.close() + + +async def test_aiosqlite_data_dictionary_version_cached_across_drivers() -> None: + conn1 = await aiosqlite.connect(":memory:") + conn2 = await aiosqlite.connect(":memory:") + try: + driver1 = AiosqliteDriver(connection=conn1) + driver2 = AiosqliteDriver(connection=conn2) + v1 = await driver1.data_dictionary.get_version(driver1) + v2 = await driver2.data_dictionary.get_version(driver2) + assert v1 is not None + assert v1 is v2 + finally: + await conn1.close() + await conn2.close() + + +async def test_aiosqlite_dispatch_execute_script_emits_events() -> None: + conn = await aiosqlite.connect(":memory:") + try: + driver = AiosqliteDriver(connection=conn) + result = await driver.execute_script("CREATE TABLE scr_t (x INT); INSERT INTO scr_t VALUES (1);") + assert result.total_statements == 2 + assert result.successful_statements == 2 + finally: + await conn.close() diff --git a/tests/unit/adapters/test_aiosqlite/test_exception_mapping.py b/tests/unit/adapters/test_aiosqlite/test_exception_mapping.py index a99e560f0..20c08fb60 100644 --- a/tests/unit/adapters/test_aiosqlite/test_exception_mapping.py +++ b/tests/unit/adapters/test_aiosqlite/test_exception_mapping.py @@ -5,7 +5,18 @@ import pytest from sqlspec.adapters.aiosqlite.core import create_mapped_exception -from sqlspec.exceptions import DeadlockError, OperationCancelledError, PermissionDeniedError, UniqueViolationError +from sqlspec.exceptions import ( + CheckViolationError, + DatabaseConnectionError, + DataError, + DeadlockError, + IntegrityError, + NotNullViolationError, + OperationalError, + OperationCancelledError, + PermissionDeniedError, + UniqueViolationError, +) class _SqliteError(sqlite3.OperationalError): @@ -128,3 +139,79 @@ def test_primary_key_identity_takes_precedence_over_message(table: str, code: in mapped = create_mapped_exception(error) assert isinstance(mapped, UniqueViolationError) assert mapped.__cause__ is error + + +def test_not_null_error_code_maps_to_not_null_violation() -> None: + err = _SqliteIntegrityError("NOT NULL constraint failed: t.val", 1299, None) + result = create_mapped_exception(err) + assert isinstance(result, NotNullViolationError) + assert result.__cause__ is err + + +def test_check_error_code_maps_to_check_violation() -> None: + err = _SqliteIntegrityError("CHECK constraint failed: val > 0", 275, None) + result = create_mapped_exception(err) + assert isinstance(result, CheckViolationError) + assert result.__cause__ is err + + +def test_trigger_abort_does_not_map_to_not_null_violation() -> None: + err = _SqliteIntegrityError("trigger aborted", 1811, None) + result = create_mapped_exception(err) + assert not isinstance(result, NotNullViolationError) + assert isinstance(result, IntegrityError) + assert result.__cause__ is err + + +def test_no_such_column_with_locked_name_maps_to_operational_error() -> None: + err = sqlite3.OperationalError("no such column: is_locked") + result = create_mapped_exception(err) + assert isinstance(result, OperationalError) + assert not isinstance(result, DeadlockError) + assert result.__cause__ is err + + +def test_programming_error_maps_to_data_error() -> None: + err = sqlite3.ProgrammingError("cannot operate on a closed database") + result = create_mapped_exception(err) + assert isinstance(result, DataError) + assert result.__cause__ is err + + +def test_interface_error_maps_to_operational_error() -> None: + err = sqlite3.InterfaceError("error binding parameter") + result = create_mapped_exception(err) + assert isinstance(result, OperationalError) + assert result.__cause__ is err + + +def test_auth_error_maps_to_permission_denied() -> None: + err = _SqliteError("not authorized", 23, "SQLITE_AUTH") + result = create_mapped_exception(err) + assert isinstance(result, PermissionDeniedError) + assert result.__cause__ is err + + +def test_range_and_toobig_errors_map_to_data_error() -> None: + err_range = _SqliteError("column index out of range", 25, "SQLITE_RANGE") + res_range = create_mapped_exception(err_range) + assert isinstance(res_range, DataError) + assert res_range.__cause__ is err_range + + err_toobig = _SqliteError("string or blob too big", 18, "SQLITE_TOOBIG") + res_toobig = create_mapped_exception(err_toobig) + assert isinstance(res_toobig, DataError) + assert res_toobig.__cause__ is err_toobig + + +def test_notadb_error_maps_to_database_connection_error() -> None: + err = _SqliteError("file is not a database", 26, "SQLITE_NOTADB") + result = create_mapped_exception(err) + assert isinstance(result, DatabaseConnectionError) + assert result.__cause__ is err + + +def test_create_mapped_exception_accepts_logger_parameter() -> None: + err = sqlite3.OperationalError("generic") + result = create_mapped_exception(err, logger=None) + assert isinstance(result, OperationalError) diff --git a/tests/unit/adapters/test_aiosqlite/test_pool.py b/tests/unit/adapters/test_aiosqlite/test_pool.py index 403f1db80..434473029 100644 --- a/tests/unit/adapters/test_aiosqlite/test_pool.py +++ b/tests/unit/adapters/test_aiosqlite/test_pool.py @@ -462,3 +462,23 @@ async def resolve() -> _FallbackConnection: return self._connection return resolve().__await__() + + +async def test_release_retires_connection_when_rollback_fails(monkeypatch: pytest.MonkeyPatch) -> None: + pool = AiosqliteConnectionPool({"database": ":memory:", "timeout": 30.0}) + pool_conn = await pool.acquire() + assert pool._queue.qsize() == 0 + + await pool_conn.connection.execute("CREATE TABLE test_tx (id INT)") + await pool_conn.connection.execute("INSERT INTO test_tx VALUES (1)") + assert getattr(pool_conn.connection, "in_transaction", False) + + async def _failing_end_transaction(_raw_conn: Any, *, commit: bool = False) -> None: + raise sqlite3.OperationalError("disk I/O error during rollback") + + monkeypatch.setattr(pool_module, "end_transaction", _failing_end_transaction) + + await pool.release(pool_conn) + assert pool._queue.qsize() == 0 + assert not pool_conn.is_healthy + await pool.close() diff --git a/tests/unit/adapters/test_arrow_odbc/test_adk_store.py b/tests/unit/adapters/test_arrow_odbc/test_adk_store.py index 9aa80febd..c1b52ef96 100644 --- a/tests/unit/adapters/test_arrow_odbc/test_adk_store.py +++ b/tests/unit/adapters/test_arrow_odbc/test_adk_store.py @@ -1,4 +1,3 @@ -# pyright: reportPrivateUsage=false """Unit tests for arrow-odbc ADK store data-dictionary existence checks.""" from typing import Any @@ -6,7 +5,8 @@ import pytest -from sqlspec.adapters.arrow_odbc.adk import ArrowOdbcADKStore +from sqlspec.adapters.arrow_odbc.adk import ArrowOdbcADKMemoryStore, ArrowOdbcADKStore +from sqlspec.exceptions import ImproperConfigurationError def _mock_config(adk_config: "dict[str, object] | None" = None) -> MagicMock: @@ -173,3 +173,23 @@ def test_arrow_odbc_list_sessions_rejects_invalid_options( store.list_sessions("app", **options) assert calls == [] + + +def test_arrow_odbc_adk_store_requires_mssql_or_db2() -> None: + """ArrowOdbcADKStore requires dialect mssql or db2.""" + config = MagicMock() + config.statement_config.dialect = "sqlite" + config.extension_config = {"adk": {}} + + with pytest.raises(ImproperConfigurationError, match="ArrowOdbcADKStore requires dialect"): + ArrowOdbcADKStore(config) + + +def test_arrow_odbc_adk_memory_store_rejects_embedding() -> None: + """ArrowOdbcADKMemoryStore rejects embeddings with NotImplementedError.""" + config = _mock_config({"enable_memory": True}) + config.statement_config.dialect = "mssql" + mem_store = ArrowOdbcADKMemoryStore(config) + + with pytest.raises(NotImplementedError, match="Vector embeddings are not supported"): + mem_store.search_entries("hello", "app", "user", embedding=[0.1, 0.2]) diff --git a/tests/unit/adapters/test_arrow_odbc/test_config.py b/tests/unit/adapters/test_arrow_odbc/test_config.py index 983d2b8d9..8518c7d27 100644 --- a/tests/unit/adapters/test_arrow_odbc/test_config.py +++ b/tests/unit/adapters/test_arrow_odbc/test_config.py @@ -179,9 +179,9 @@ def test_driver_pooling_defaults_to_off(monkeypatch: pytest.MonkeyPatch) -> None def test_enable_driver_pooling_is_a_declared_feature() -> None: assert ( - "enable_driver_pooling" + "enable_pooling" in ArrowOdbcConfig( - connection_config={"dsn": "analytics"}, driver_features={"enable_driver_pooling": False} + connection_config={"dsn": "analytics"}, driver_features={"enable_pooling": False} ).driver_features ) @@ -315,7 +315,13 @@ def broken_pooling() -> None: raise RuntimeError(msg) monkeypatch.setattr(arrow_odbc_config, "enable_odbc_connection_pooling", broken_pooling) - config = ArrowOdbcConfig(connection_config={"dsn": "analytics"}, driver_features={"enable_driver_pooling": True}) + config = ArrowOdbcConfig(connection_config={"dsn": "analytics"}, driver_features={"enable_pooling": True}) with pytest.raises(ImproperConfigurationError, match="Could not configure arrow-odbc connection"): config.create_connection() + + +def test_driver_features_chunk_size_default() -> None: + """ArrowOdbcConfig defaults chunk_size to 65536 matching documentation.""" + config = ArrowOdbcConfig() + assert config.driver_features["chunk_size"] == 65_536 diff --git a/tests/unit/adapters/test_arrow_odbc/test_db2_metadata.py b/tests/unit/adapters/test_arrow_odbc/test_db2_metadata.py index 35cc49781..8500b0a5e 100644 --- a/tests/unit/adapters/test_arrow_odbc/test_db2_metadata.py +++ b/tests/unit/adapters/test_arrow_odbc/test_db2_metadata.py @@ -10,6 +10,7 @@ from sqlspec.adapters.arrow_odbc import ArrowOdbcConfig, ArrowOdbcDriver from sqlspec.adapters.arrow_odbc.core import create_mapped_exception from sqlspec.adapters.arrow_odbc.type_converter import odbc_type_to_arrow +from sqlspec.data_dictionary import MetadataSource, MetadataSupport from sqlspec.exceptions import CheckViolationError, DeadlockError, QueryTimeoutError from tests.unit.adapters.test_arrow_odbc._db2_fakes import ( DB2_CONNECTION_STRING, @@ -133,3 +134,16 @@ def test_result_names_kept_when_lowercasing_is_off(driver_factory: Any) -> None: assert names == { key: ["ID", "CUSTOMER_NAME", "MixedCase"] for key in ("execute", "stream", "table", "reader", "reader_batches") } + + +def test_db2_metadata_capabilities_reports_supported_for_catalog_domains() -> None: + """Capabilities for views and schemas report SUPPORTED on dialects with bundled SQL packs.""" + connection = FakeArrowOdbcConnection() + driver = _db2_driver(connection) + caps = driver.data_dictionary.get_metadata_capabilities(driver) + schema_cap = caps.get("schemas") + views_cap = caps.get("views") + + assert schema_cap is not None and schema_cap.support == MetadataSupport.SUPPORTED + assert views_cap is not None and views_cap.support == MetadataSupport.SUPPORTED + assert schema_cap.source == MetadataSource.CATALOG diff --git a/tests/unit/adapters/test_arrow_odbc/test_driver.py b/tests/unit/adapters/test_arrow_odbc/test_driver.py index e76016fc7..843597d52 100644 --- a/tests/unit/adapters/test_arrow_odbc/test_driver.py +++ b/tests/unit/adapters/test_arrow_odbc/test_driver.py @@ -1277,3 +1277,35 @@ def insert_into_table(self, **kwargs: Any) -> None: arrow_result = ArrowResult(statement=SQL("SELECT 1"), data=table, rows_affected=2) driver.bulk_insert_arrow("items", arrow_result) assert connection.insert_calls == [("items", 65536, table)] + + +def test_ambient_exception_isolation_select_to_arrow() -> None: + """select_to_arrow executes cleanly inside an active exception block without ambient exception leakage.""" + connection = FakeConnection() + driver = ArrowOdbcDriver(cast("ArrowOdbcConnection", connection)) + + try: + raise ValueError("ambient error") + except ValueError: + result = driver.select_to_arrow("SELECT 1") + assert result.num_rows == 3 + + +def test_transaction_lifecycle_maps_generic_exception() -> None: + """begin, commit, and rollback map native exceptions with appropriate transaction failure messages.""" + connection = FakeConnection() + native_error = ArrowOdbcError.__new__(ArrowOdbcError) + native_error.args = ("native disk error",) + + def raise_native() -> None: + raise native_error + + setattr(connection, "commit", raise_native) + setattr(connection, "rollback", raise_native) + driver = ArrowOdbcDriver(cast("ArrowOdbcConnection", connection)) + + with pytest.raises(SQLSpecError, match="Failed to commit transaction"): + driver.commit() + + with pytest.raises(SQLSpecError, match="Failed to rollback transaction"): + driver.rollback() diff --git a/tests/unit/adapters/test_arrow_odbc/test_extensions.py b/tests/unit/adapters/test_arrow_odbc/test_extensions.py index 3f14edc9f..9e2b530e4 100644 --- a/tests/unit/adapters/test_arrow_odbc/test_extensions.py +++ b/tests/unit/adapters/test_arrow_odbc/test_extensions.py @@ -3,6 +3,8 @@ import ast from pathlib import Path +from sqlspec.adapters.arrow_odbc.core import quote_bracket_identifier, strip_bracket_quotes + def test_arrow_odbc_extensions_do_not_import_other_adapters() -> None: """Extension modules should stay inside the arrow-odbc adapter boundary.""" @@ -28,3 +30,14 @@ def test_arrow_odbc_extensions_do_not_import_other_adapters() -> None: violations.append((str(path.relative_to(adapter_root.parents[2])), module)) assert violations == [] + + +def test_bracket_quoting_and_stripping() -> None: + """Bracket identifier quoting and quote stripping round-trips correctly.""" + assert quote_bracket_identifier("users") == "[users]" + assert quote_bracket_identifier("user]name") == "[user]]name]" + assert strip_bracket_quotes("[users]") == "users" + assert strip_bracket_quotes("[user]]name]") == "user]name" + assert strip_bracket_quotes('"users"') == "users" + assert strip_bracket_quotes('""users""') == '"users"' + assert strip_bracket_quotes("plain") == "plain" diff --git a/tests/unit/adapters/test_arrow_odbc/test_tsql_stores.py b/tests/unit/adapters/test_arrow_odbc/test_tsql_stores.py index 8fc63b632..1aeb6ad89 100644 --- a/tests/unit/adapters/test_arrow_odbc/test_tsql_stores.py +++ b/tests/unit/adapters/test_arrow_odbc/test_tsql_stores.py @@ -10,6 +10,7 @@ from sqlspec.adapters.arrow_odbc import ArrowOdbcConfig from sqlspec.adapters.arrow_odbc.adk.store import ArrowOdbcADKMemoryStore, ArrowOdbcADKStore, _datetime_value +from sqlspec.adapters.arrow_odbc.core import atomic_store_write from sqlspec.adapters.arrow_odbc.events.store import ArrowOdbcEventQueueStore from sqlspec.adapters.arrow_odbc.litestar import ArrowOdbcStore from sqlspec.adapters.arrow_odbc.litestar.store import _TsqlSessionSql @@ -75,11 +76,12 @@ async def test_tsql_session_set_inline_inserts() -> None: await _session_store(connection).set("k", b"v", 60) assert normalized_calls(connection) == [ + ("BEGIN TRANSACTION", None), ("SELECT session_id FROM sess WHERE session_id = ?", ["k"]), ("INSERT INTO sess (session_id, data, expires_at) VALUES (?, ?, ?)", ["k", "dg==", ISO_TIMESTAMP]), ("DELETE FROM sess_chunks WHERE session_id = ?", ["k"]), + ("COMMIT TRANSACTION", None), ] - assert connection.commit_calls == 1 async def test_tsql_session_set_chunked_updates_existing() -> None: @@ -92,15 +94,17 @@ async def test_tsql_session_set_chunked_updates_existing() -> None: calls = normalized_calls(connection) assert [sql for sql, _ in calls] == [ + "BEGIN TRANSACTION", "SELECT session_id FROM sess WHERE session_id = ?", "UPDATE sess SET data = ?, expires_at = ?, updated_at = SYSUTCDATETIME() WHERE session_id = ?", "DELETE FROM sess_chunks WHERE session_id = ?", "INSERT INTO sess_chunks (session_id, chunk_index, data) VALUES (?, ?, ?)", "INSERT INTO sess_chunks (session_id, chunk_index, data) VALUES (?, ?, ?)", + "COMMIT TRANSACTION", ] - assert calls[1][1] == [None, None, "k"] - assert [parameters[:2] for _, parameters in calls[3:] if parameters] == [["k", "0"], ["k", "1"]] - assert [len(parameters[2] or "") for _, parameters in calls[3:] if parameters] == [3500, 500] + assert calls[2][1] == [None, None, "k"] + assert [parameters[:2] for _, parameters in calls[4:-1] if parameters] == [["k", "0"], ["k", "1"]] + assert [len(parameters[2] or "") for _, parameters in calls[4:-1] if parameters] == [3500, 500] async def test_tsql_session_get_renews_expiry() -> None: @@ -191,11 +195,13 @@ async def test_tsql_session_delete_expired() -> None: assert await _session_store(connection).delete_expired() == 2 assert normalized_calls(connection) == [ + ("BEGIN TRANSACTION", None), ( "SELECT COUNT(*) AS expired_count FROM sess WHERE expires_at IS NOT NULL AND expires_at < SYSUTCDATETIME()", None, ), ("DELETE FROM sess WHERE expires_at IS NOT NULL AND expires_at < SYSUTCDATETIME()", None), + ("COMMIT TRANSACTION", None), ] @@ -250,70 +256,137 @@ def test_tsql_event_store_create_and_drop_statements() -> None: "event_data": {"x": 1}, } _TSQL_ADK_SESSION_CALLS = [ + ("BEGIN TRANSACTION", None), ( - "INSERT INTO [dbo].[adk_session] ( id, app_name, user_id, [tenant_id], state, create_time, update_time ) VALUES (?, ?, ?, ?, ?, SYSUTCDATETIME(), SYSUTCDATETIME())", + ( + "INSERT INTO [dbo].[adk_session] ( id, app_name, user_id, [tenant_id], " + "state, create_time, update_time ) VALUES (?, ?, ?, ?, ?, SYSUTCDATETIME(), " + "SYSUTCDATETIME())" + ), ["s1", "app", "u", "7", '{"a":1}'], ), ( - "SELECT TOP 1 id, app_name, user_id, state, create_time, update_time FROM [dbo].[adk_session] WHERE app_name = ? AND user_id = ? AND id = ?", + ( + "SELECT TOP 1 id, app_name, user_id, state, create_time, update_time FROM " + "[dbo].[adk_session] WHERE app_name = ? AND user_id = ? AND id = ?" + ), ["app", "u", "s1"], ), + ("COMMIT TRANSACTION", None), ( "UPDATE [dbo].[adk_session] SET update_time = SYSUTCDATETIME() WHERE app_name = ? AND user_id = ? AND id = ?", ["app", "u", "s1"], ), ( - "SELECT TOP 1 id, app_name, user_id, state, create_time, update_time FROM [dbo].[adk_session] WHERE app_name = ? AND user_id = ? AND id = ?", + ( + "SELECT TOP 1 id, app_name, user_id, state, create_time, update_time FROM " + "[dbo].[adk_session] WHERE app_name = ? AND user_id = ? AND id = ?" + ), ["app", "u", "s1"], ), + ("BEGIN TRANSACTION", None), ( - "UPDATE [dbo].[adk_session] SET state = ?, update_time = SYSUTCDATETIME() WHERE app_name = ? AND user_id = ? AND id = ?", + ( + "UPDATE [dbo].[adk_session] SET state = ?, update_time = SYSUTCDATETIME() " + "WHERE app_name = ? AND user_id = ? AND id = ?" + ), ['{"a":2}', "app", "u", "s1"], ), + ("COMMIT TRANSACTION", None), ( - "SELECT id, app_name, user_id, state, create_time, update_time FROM [dbo].[adk_session] WHERE app_name = ? AND user_id = ? ORDER BY update_time DESC, id DESC OFFSET 10 ROWS FETCH NEXT 5 ROWS ONLY", + ( + "SELECT id, app_name, user_id, state, create_time, update_time FROM " + "[dbo].[adk_session] WHERE app_name = ? AND user_id = ? ORDER BY update_time " + "DESC, id DESC OFFSET 10 ROWS FETCH NEXT 5 ROWS ONLY" + ), ["app", "u"], ), + ("BEGIN TRANSACTION", None), ("DELETE FROM [dbo].[adk_session] WHERE app_name = ? AND user_id = ? AND id = ?", ["app", "u", "s1"]), + ("COMMIT TRANSACTION", None), + ("BEGIN TRANSACTION", None), ( - "INSERT INTO [dbo].[adk_event] ( id, app_name, user_id, session_id, invocation_id, timestamp, event_data ) VALUES (?, ?, ?, ?, ?, ?, ?)", + ( + "INSERT INTO [dbo].[adk_event] ( id, app_name, user_id, session_id, " + "invocation_id, timestamp, event_data ) VALUES (?, ?, ?, ?, ?, ?, ?)" + ), ["e1", "app", "u", "s1", "i1", "2026-01-02T03:04:05.678901", '{"x":1}'], ), + ("COMMIT TRANSACTION", None), + ("BEGIN TRANSACTION", None), ( - "UPDATE [dbo].[adk_session] SET state = ?, update_time = SYSUTCDATETIME() WHERE app_name = ? AND user_id = ? AND id = ?", + ( + "UPDATE [dbo].[adk_session] SET state = ?, update_time = SYSUTCDATETIME() " + "WHERE app_name = ? AND user_id = ? AND id = ?" + ), ['{"a":3}', "app", "u", "s1"], ), ( - "SELECT TOP 1 id, app_name, user_id, state, create_time, update_time FROM [dbo].[adk_session] WHERE app_name = ? AND user_id = ? AND id = ?", + ( + "SELECT TOP 1 id, app_name, user_id, state, create_time, update_time FROM " + "[dbo].[adk_session] WHERE app_name = ? AND user_id = ? AND id = ?" + ), ["app", "u", "s1"], ), ( - "INSERT INTO [dbo].[adk_event] ( id, app_name, user_id, session_id, invocation_id, timestamp, event_data ) VALUES (?, ?, ?, ?, ?, ?, ?)", + ( + "INSERT INTO [dbo].[adk_event] ( id, app_name, user_id, session_id, " + "invocation_id, timestamp, event_data ) VALUES (?, ?, ?, ?, ?, ?, ?)" + ), ["e1", "app", "u", "s1", "i1", "2026-01-02T03:04:05.678901", '{"x":1}'], ), ( - "MERGE INTO [dbo].[adk_app_state] WITH (HOLDLOCK) AS target USING (SELECT ? AS [app_name], ? AS state) AS source ON (target.[app_name] = source.[app_name]) WHEN MATCHED THEN UPDATE SET state = source.state, update_time = SYSUTCDATETIME() WHEN NOT MATCHED THEN INSERT ([app_name], [state], [update_time]) VALUES (source.[app_name], source.[state], SYSUTCDATETIME());", + ( + "MERGE INTO [dbo].[adk_app_state] WITH (HOLDLOCK) AS target USING (SELECT ? " + "AS [app_name], ? AS state) AS source ON (target.[app_name] = " + "source.[app_name]) WHEN MATCHED THEN UPDATE SET state = source.state, " + "update_time = SYSUTCDATETIME() WHEN NOT MATCHED THEN INSERT ([app_name], " + "[state], [update_time]) VALUES (source.[app_name], source.[state], " + "SYSUTCDATETIME());" + ), ["app", '{"g":1}'], ), ( - "MERGE INTO [dbo].[adk_user_state] WITH (HOLDLOCK) AS target USING (SELECT ? AS [app_name], ? AS [user_id], ? AS state) AS source ON (target.[app_name] = source.[app_name] AND target.[user_id] = source.[user_id]) WHEN MATCHED THEN UPDATE SET state = source.state, update_time = SYSUTCDATETIME() WHEN NOT MATCHED THEN INSERT ([app_name], [user_id], [state], [update_time]) VALUES (source.[app_name], source.[user_id], source.[state], SYSUTCDATETIME());", + ( + "MERGE INTO [dbo].[adk_user_state] WITH (HOLDLOCK) AS target USING (SELECT ? " + "AS [app_name], ? AS [user_id], ? AS state) AS source ON (target.[app_name] " + "= source.[app_name] AND target.[user_id] = source.[user_id]) WHEN MATCHED " + "THEN UPDATE SET state = source.state, update_time = SYSUTCDATETIME() WHEN " + "NOT MATCHED THEN INSERT ([app_name], [user_id], [state], [update_time]) " + "VALUES (source.[app_name], source.[user_id], source.[state], " + "SYSUTCDATETIME());" + ), ["app", "u", '{"h":2}'], ), + ("COMMIT TRANSACTION", None), ( - "SELECT TOP 4 id, app_name, user_id, session_id, invocation_id, timestamp, event_data FROM [dbo].[adk_event] WHERE app_name = ? AND user_id = ? AND session_id = ? AND timestamp > ? ORDER BY timestamp ASC", + ( + "SELECT TOP 4 id, app_name, user_id, session_id, invocation_id, timestamp, " + "event_data FROM [dbo].[adk_event] WHERE app_name = ? AND user_id = ? AND " + "session_id = ? AND timestamp > ? ORDER BY timestamp ASC" + ), ["app", "u", "s1", "2026-01-02T03:04:05.678901"], ), ( - "SELECT id, app_name, user_id, session_id, invocation_id, timestamp, event_data FROM [dbo].[adk_event] WHERE app_name = ? AND user_id = ? AND session_id = ? ORDER BY timestamp ASC", + ( + "SELECT id, app_name, user_id, session_id, invocation_id, timestamp, " + "event_data FROM [dbo].[adk_event] WHERE app_name = ? AND user_id = ? AND " + "session_id = ? ORDER BY timestamp ASC" + ), ["app", "u", "s1"], ), + ("BEGIN TRANSACTION", None), ( "SELECT COUNT(*) AS row_count FROM [dbo].[adk_event] WHERE timestamp < ? AND app_name = ?", ["2026-01-02T03:04:05.678901", "app"], ), ("DELETE FROM [dbo].[adk_event] WHERE timestamp < ? AND app_name = ?", ["2026-01-02T03:04:05.678901", "app"]), + ("COMMIT TRANSACTION", None), + ("BEGIN TRANSACTION", None), ("SELECT COUNT(*) AS row_count FROM [dbo].[adk_session] WHERE update_time < ?", ["2026-01-02T03:04:05.678901"]), ("DELETE FROM [dbo].[adk_session] WHERE update_time < ?", ["2026-01-02T03:04:05.678901"]), + ("COMMIT TRANSACTION", None), + ("BEGIN TRANSACTION", None), ( "SELECT COUNT(*) AS row_count FROM [dbo].[adk_user_state] WHERE update_time < ? AND app_name = ?", ["2026-01-02T03:04:05.678901", "app"], @@ -322,21 +395,48 @@ def test_tsql_event_store_create_and_drop_statements() -> None: "DELETE FROM [dbo].[adk_user_state] WHERE update_time < ? AND app_name = ?", ["2026-01-02T03:04:05.678901", "app"], ), + ("COMMIT TRANSACTION", None), ("SELECT TOP 1 state FROM [dbo].[adk_app_state] WHERE app_name = ?", ["app"]), ("SELECT TOP 1 state FROM [dbo].[adk_user_state] WHERE app_name = ? AND user_id = ?", ["app", "u"]), + ("BEGIN TRANSACTION", None), ( - "MERGE INTO [dbo].[adk_app_state] WITH (HOLDLOCK) AS target USING (SELECT ? AS [app_name], ? AS state) AS source ON (target.[app_name] = source.[app_name]) WHEN MATCHED THEN UPDATE SET state = source.state, update_time = SYSUTCDATETIME() WHEN NOT MATCHED THEN INSERT ([app_name], [state], [update_time]) VALUES (source.[app_name], source.[state], SYSUTCDATETIME());", + ( + "MERGE INTO [dbo].[adk_app_state] WITH (HOLDLOCK) AS target USING (SELECT ? " + "AS [app_name], ? AS state) AS source ON (target.[app_name] = " + "source.[app_name]) WHEN MATCHED THEN UPDATE SET state = source.state, " + "update_time = SYSUTCDATETIME() WHEN NOT MATCHED THEN INSERT ([app_name], " + "[state], [update_time]) VALUES (source.[app_name], source.[state], " + "SYSUTCDATETIME());" + ), ["app", '{"g":2}'], ), + ("COMMIT TRANSACTION", None), + ("BEGIN TRANSACTION", None), ( - "MERGE INTO [dbo].[adk_user_state] WITH (HOLDLOCK) AS target USING (SELECT ? AS [app_name], ? AS [user_id], ? AS state) AS source ON (target.[app_name] = source.[app_name] AND target.[user_id] = source.[user_id]) WHEN MATCHED THEN UPDATE SET state = source.state, update_time = SYSUTCDATETIME() WHEN NOT MATCHED THEN INSERT ([app_name], [user_id], [state], [update_time]) VALUES (source.[app_name], source.[user_id], source.[state], SYSUTCDATETIME());", + ( + "MERGE INTO [dbo].[adk_user_state] WITH (HOLDLOCK) AS target USING (SELECT ? " + "AS [app_name], ? AS [user_id], ? AS state) AS source ON (target.[app_name] " + "= source.[app_name] AND target.[user_id] = source.[user_id]) WHEN MATCHED " + "THEN UPDATE SET state = source.state, update_time = SYSUTCDATETIME() WHEN " + "NOT MATCHED THEN INSERT ([app_name], [user_id], [state], [update_time]) " + "VALUES (source.[app_name], source.[user_id], source.[state], " + "SYSUTCDATETIME());" + ), ["app", "u", '{"h":3}'], ), + ("COMMIT TRANSACTION", None), ("SELECT TOP 1 value FROM [dbo].[adk_internal_metadata] WHERE [key] = ?", ["k"]), + ("BEGIN TRANSACTION", None), ( - "MERGE INTO [dbo].[adk_internal_metadata] WITH (HOLDLOCK) AS target USING (SELECT ? AS [key], ? AS value) AS source ON (target.[key] = source.[key]) WHEN MATCHED THEN UPDATE SET value = source.value WHEN NOT MATCHED THEN INSERT ([key], value) VALUES (source.[key], source.value);", + ( + "MERGE INTO [dbo].[adk_internal_metadata] WITH (HOLDLOCK) AS target USING " + "(SELECT ? AS [key], ? AS value) AS source ON (target.[key] = source.[key]) " + "WHEN MATCHED THEN UPDATE SET value = source.value WHEN NOT MATCHED THEN " + "INSERT ([key], value) VALUES (source.[key], source.value);" + ), ["k", "v"], ), + ("COMMIT TRANSACTION", None), ] _TSQL_ADK_DDL = [ "CREATE TABLE [dbo].[adk_session] ( row_id UNIQUEIDENTIFIER NOT NULL CONSTRAINT [df_adk_session_row_id] DEFAULT NEWSEQUENTIALID(), id NVARCHAR(128) NOT NULL, app_name NVARCHAR(128) NOT NULL, user_id NVARCHAR(128) NOT NULL, tenant_id INTEGER, state NVARCHAR(MAX) NOT NULL, create_time DATETIME2(6) NOT NULL CONSTRAINT [df_adk_session_create_time] DEFAULT SYSUTCDATETIME(), update_time DATETIME2(6) NOT NULL CONSTRAINT [df_adk_session_update_time] DEFAULT SYSUTCDATETIME(), CONSTRAINT [pk_adk_session_row_id] PRIMARY KEY (row_id), CONSTRAINT [uq_adk_session_id] UNIQUE (id) )", @@ -415,9 +515,14 @@ def test_tsql_adk_missing_table_reads_empty(monkeypatch: pytest.MonkeyPatch) -> _TSQL_ADK_MEMORY_CALLS = [ + ("BEGIN TRANSACTION", None), ("SELECT TOP 1 id FROM [dbo].[adk_memory] WHERE event_id = ?", ["e1"]), ( - "INSERT INTO [dbo].[adk_memory] ( id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json, inserted_at, [tenant_id] ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + ( + "INSERT INTO [dbo].[adk_memory] ( id, session_id, app_name, user_id, scope, " + "event_id, author, timestamp, content_json, content_text, metadata_json, " + "inserted_at, [tenant_id] ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)" + ), [ "m1", "s1", @@ -434,16 +539,31 @@ def test_tsql_adk_missing_table_reads_empty(monkeypatch: pytest.MonkeyPatch) -> "7", ], ), + ("COMMIT TRANSACTION", None), ( - "SELECT id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json, inserted_at FROM [dbo].[adk_memory] WHERE app_name = ? AND ((scope = 'user' AND user_id = ?) OR scope = 'app') AND content_text LIKE ? ORDER BY timestamp DESC OFFSET 0 ROWS FETCH NEXT 3 ROWS ONLY", + ( + "SELECT id, session_id, app_name, user_id, scope, event_id, author, " + "timestamp, content_json, content_text, metadata_json, inserted_at FROM " + "[dbo].[adk_memory] WHERE app_name = ? AND ((scope = 'user' AND user_id = ?) " + "OR scope = 'app') AND content_text LIKE ? ORDER BY timestamp DESC OFFSET 0 " + "ROWS FETCH NEXT 3 ROWS ONLY" + ), ["app", "u", "%hello%"], ), ( - "SELECT id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json, inserted_at FROM [dbo].[adk_memory] WHERE app_name = ? AND scope = 'app' AND content_text LIKE ? ORDER BY timestamp DESC OFFSET 0 ROWS FETCH NEXT 20 ROWS ONLY", + ( + "SELECT id, session_id, app_name, user_id, scope, event_id, author, " + "timestamp, content_json, content_text, metadata_json, inserted_at FROM " + "[dbo].[adk_memory] WHERE app_name = ? AND scope = 'app' AND content_text " + "LIKE ? ORDER BY timestamp DESC OFFSET 0 ROWS FETCH NEXT 20 ROWS ONLY" + ), ["app", "%hello%"], ), + ("BEGIN TRANSACTION", None), ("SELECT COUNT(*) AS row_count FROM [dbo].[adk_memory] WHERE session_id = ?", ["s1"]), ("DELETE FROM [dbo].[adk_memory] WHERE session_id = ?", ["s1"]), + ("COMMIT TRANSACTION", None), + ("BEGIN TRANSACTION", None), ( "SELECT COUNT(*) AS row_count FROM [dbo].[adk_memory] WHERE inserted_at < ? AND app_name = ? AND scope = ?", [ISO_TIMESTAMP, "app", "user"], @@ -452,6 +572,7 @@ def test_tsql_adk_missing_table_reads_empty(monkeypatch: pytest.MonkeyPatch) -> "DELETE FROM [dbo].[adk_memory] WHERE inserted_at < ? AND app_name = ? AND scope = ?", [ISO_TIMESTAMP, "app", "user"], ), + ("COMMIT TRANSACTION", None), ] _TSQL_ADK_MEMORY_DDL = "CREATE TABLE [dbo].[adk_memory] ( id NVARCHAR(128) NOT NULL, session_id NVARCHAR(128) NOT NULL, app_name NVARCHAR(128) NOT NULL, user_id NVARCHAR(128) NOT NULL, scope NVARCHAR(16) NOT NULL DEFAULT 'user', event_id NVARCHAR(128) NOT NULL, author NVARCHAR(256) NULL, timestamp DATETIME2(6) NOT NULL, content_json NVARCHAR(MAX) NOT NULL, content_text NVARCHAR(MAX) NOT NULL, metadata_json NVARCHAR(MAX) NULL, inserted_at DATETIME2(6) NOT NULL, tenant_id INTEGER, CONSTRAINT [pk_adk_memory_id] PRIMARY KEY (id), CONSTRAINT [uq_adk_memory_event_id] UNIQUE (event_id) )" _TSQL_ADK_MEMORY_DROPS = ["DROP TABLE IF EXISTS [dbo].[adk_memory]"] @@ -505,10 +626,11 @@ def test_tsql_adk_memory_create_tables_and_insert_without_owner() -> None: "CREATE INDEX [idx_adk_memory_scope] ON [dbo].[adk_memory]", "CREATE INDEX [idx_adk_memory_session] ON [dbo].[adk_memory]", ] - assert statements[-1] == ( + assert statements[-2] == ( "INSERT INTO [dbo].[adk_memory] ( id, session_id, app_name, user_id, scope, event_id, author, timestamp, " "content_json, content_text, metadata_json, inserted_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)" ) + assert statements[-1] == "COMMIT TRANSACTION" async def test_tsql_session_get_parses_string_expiry_and_deletes_expired_in_same_session( @@ -556,7 +678,7 @@ def test_tsql_adk_append_event_and_update_state_rolls_back_when_session_missing( with pytest.raises(ValueError, match="Session s1 not found"): store.append_event_and_update_state(_ADK_EVENT, "app", "u", "s1", {"a": 3}) - assert connection.rollback_calls == 1 + assert any(sql == "ROLLBACK TRANSACTION" for sql, _ in normalized_calls(connection)) assert connection.commit_calls == 0 @@ -586,22 +708,22 @@ def test_schema_qualified_tables_across_tsql_stores() -> None: extension_config={"events": {"queue_table": "custom.app_events"}}, ) adk_store = ArrowOdbcADKStore(config) - adk_store._session_table = "custom.adk_session" # pyright: ignore[reportPrivateUsage] + adk_store._session_table = "custom.adk_session" memory_store = ArrowOdbcADKMemoryStore(config) - memory_store._memory_table = "custom.adk_memory" # pyright: ignore[reportPrivateUsage] + memory_store._memory_table = "custom.adk_memory" litestar_store = ArrowOdbcStore(config) - litestar_store._sql = _TsqlSessionSql("custom.sess") # pyright: ignore[reportPrivateUsage] + litestar_store._sql = _TsqlSessionSql("custom.sess") event_store = ArrowOdbcEventQueueStore(config) - session_ddl = adk_store._sessions_table_ddl() # pyright: ignore[reportPrivateUsage] + session_ddl = adk_store._sessions_table_ddl() assert "[custom].[adk_session]" in session_ddl assert "[pk_adk_session_row_id]" in session_ddl - memory_ddl = memory_store._memory_table_ddl() # pyright: ignore[reportPrivateUsage] + memory_ddl = memory_store._memory_table_ddl() assert "[custom].[adk_memory]" in memory_ddl assert "[pk_adk_memory_id]" in memory_ddl - litestar_ddl = litestar_store._table_ddl() # pyright: ignore[reportPrivateUsage] + litestar_ddl = litestar_store._table_ddl() assert "SCHEMA_ID(N'custom')" in litestar_ddl assert "IX_sess_expires_at" in litestar_ddl @@ -616,3 +738,38 @@ def test_datetime_value_accepts_date_and_rejects_invalid_types() -> None: assert _datetime_value(date(2026, 5, 10)) == datetime(2026, 5, 10, 0, 0, tzinfo=UTC) with pytest.raises(TypeError, match="Unsupported datetime value"): _datetime_value(12345) + + +def test_atomic_write_rolls_back_on_failure_with_default_autocommit() -> None: + """_atomic_write rolls back on failure even with default autocommit=True.""" + + class MockDriver: + def __init__(self) -> None: + self._connection_autocommit = True + self._dialect = "mssql" + self.begun = False + self.rolled_back = False + self.committed = False + + def _connection_in_transaction(self) -> bool: + return self.begun + + def begin(self) -> None: + self.begun = True + + def rollback(self) -> None: + self.rolled_back = True + self.begun = False + + def commit(self) -> None: + self.committed = True + self.begun = False + + driver = MockDriver() + with pytest.raises(RuntimeError, match="second statement failed"): + with atomic_store_write(driver): + assert driver.begun is True + raise RuntimeError("second statement failed") + + assert driver.rolled_back is True + assert driver.committed is False diff --git a/tests/unit/adapters/test_asyncmy/test_adk_store.py b/tests/unit/adapters/test_asyncmy/test_adk_store.py index e3b2d6759..d60897090 100644 --- a/tests/unit/adapters/test_asyncmy/test_adk_store.py +++ b/tests/unit/adapters/test_asyncmy/test_adk_store.py @@ -226,6 +226,7 @@ async def test_asyncmy_adk_memory_store_insert_and_search_json_handling() -> Non cursor = MagicMock() cursor.rowcount = 1 cursor.execute = AsyncMock() + cursor.executemany = AsyncMock() cursor.description = [ ("id",), ("session_id",), @@ -288,7 +289,7 @@ async def test_asyncmy_adk_memory_store_insert_and_search_json_handling() -> Non records = await store.search_entries("hello", "app", "user-1") assert inserted == 1 - insert_params = cursor.execute.call_args_list[0].args[1] + insert_params = cursor.executemany.call_args_list[0].args[1][0] assert insert_params[6] is None assert insert_params[11] is None assert len(records) == 1 @@ -296,6 +297,84 @@ async def test_asyncmy_adk_memory_store_insert_and_search_json_handling() -> Non assert records[0]["metadata_json"] == {"source": "unit"} +async def test_asyncmy_insert_memory_entries_batches_with_executemany() -> None: + """insert_memory_entries batches multiple entries in a single cursor.executemany call.""" + from datetime import UTC, datetime + from unittest.mock import AsyncMock + + now = datetime.now(tz=UTC) + cursor = MagicMock() + cursor.rowcount = 2 + cursor.executemany = AsyncMock() + conn = MagicMock() + conn.cursor = MagicMock(return_value=cursor) + conn.commit = AsyncMock() + conn.__aenter__ = AsyncMock(return_value=conn) + conn.__aexit__ = AsyncMock(return_value=None) + config = _mock_config() + config.provide_connection = lambda *_a, **_k: conn + store = AsyncmyADKMemoryStore(config) + + entries = [ + cast( + "Any", + { + "id": f"mem-{i}", + "session_id": "sess-1", + "app_name": "app", + "user_id": "user-1", + "event_id": f"evt-{i}", + "timestamp": now, + "content_json": {"idx": i}, + "content_text": f"text {i}", + "metadata_json": None, + "inserted_at": now, + }, + ) + for i in range(2) + ] + + inserted = await store.insert_memory_entries(entries) + + assert inserted == 2 + assert cursor.executemany.call_count == 1 + batch_params = cursor.executemany.call_args_list[0].args[1] + assert len(batch_params) == 2 + + +async def test_asyncmy_search_entries_escapes_like_wildcards() -> None: + """search_entries escapes %, _, and \\ wildcards and appends ESCAPE clause.""" + from unittest.mock import AsyncMock + + cursor = MagicMock() + cursor.fetchall = AsyncMock(return_value=[]) + cursor.description = [] + cursor.execute = AsyncMock() + conn = MagicMock() + conn.cursor = MagicMock(return_value=cursor) + conn.__aenter__ = AsyncMock(return_value=conn) + conn.__aexit__ = AsyncMock(return_value=None) + config = _mock_config() + config.provide_connection = lambda *_a, **_k: conn + store = AsyncmyADKMemoryStore(config) + + await store.search_entries("100%_val\\id", "app", "user-1") + + assert cursor.execute.call_count == 1 + executed_sql = cursor.execute.call_args_list[0].args[0] + executed_params = cursor.execute.call_args_list[0].args[1] + assert "LIKE %s ESCAPE '\\\\'" in executed_sql + assert "%100\\%\\_val\\\\id%" in executed_params + + +async def test_asyncmy_memory_table_ddl_uses_longtext() -> None: + """_memory_table_ddl uses LONGTEXT for content_text.""" + config = _mock_config() + store = AsyncmyADKMemoryStore(config) + ddl = await store._memory_table_ddl() + assert "content_text LONGTEXT NOT NULL" in ddl + + async def test_asyncmy_stream_source_closes_cursor_when_execute_raises() -> None: """AsyncmyStreamSource.start() closes the cursor if execute() fails.""" from unittest.mock import AsyncMock diff --git a/tests/unit/adapters/test_asyncmy/test_config.py b/tests/unit/adapters/test_asyncmy/test_config.py index 5d091f1d2..35cdf6547 100644 --- a/tests/unit/adapters/test_asyncmy/test_config.py +++ b/tests/unit/adapters/test_asyncmy/test_config.py @@ -329,3 +329,73 @@ def test_asyncmy_config_with_dsn_applies_to_connection_kwargs() -> None: assert conn_kwargs["password"] == "pass1" assert conn_kwargs["host"] == "dbhost" assert conn_kwargs["database"] == "production" + + +def test_asyncmy_config_driver_features_initialization() -> None: + """AsyncmyConfig initializes driver_features without redundant copying.""" + config = AsyncmyConfig(connection_config={"local_infile": True}, driver_features={"json_serializer": str}) + assert config.driver_features["json_serializer"] is str + assert config.driver_features["enable_local_infile_bulk_load"] is True + + +async def test_asyncmy_create_connection_closes_socket_on_error(monkeypatch: pytest.MonkeyPatch) -> None: + """create_connection closes the standalone connection if _ensure_connection raises.""" + from unittest.mock import AsyncMock, MagicMock + + import sqlspec.adapters.asyncmy.config as config_mod + + mock_conn = MagicMock() + mock_conn.close = MagicMock() + monkeypatch.setattr(config_mod.asyncmy, "connect", AsyncMock(return_value=mock_conn)) + + config = AsyncmyConfig( + connection_config={"host": "localhost", "database": "test"}, + driver_features={"on_connection_create": AsyncMock(side_effect=RuntimeError("hook error"))}, + ) + + with pytest.raises(RuntimeError, match="hook error"): + await config.create_connection() + + assert mock_conn.close.called + + +def test_asyncmy_config_capability_flags() -> None: + """AsyncmyConfig declares accurate native Arrow/Parquet capabilities as False.""" + assert AsyncmyConfig.supports_native_arrow_export is False + assert AsyncmyConfig.supports_native_parquet_export is False + assert AsyncmyConfig.supports_native_arrow_import is False + assert AsyncmyConfig.supports_native_parquet_import is False + + +def test_conversion_capabilities_require_pyarrow(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(AsyncmyConfig, "_dependency_available", staticmethod(lambda _: False)) + + capabilities = AsyncmyConfig().storage_capabilities() + + assert capabilities["arrow_import_enabled"] is False + assert capabilities["arrow_export_enabled"] is False + assert capabilities["parquet_import_enabled"] is False + assert capabilities["parquet_export_enabled"] is False + + +async def test_asyncmy_session_context_initialization() -> None: + """AsyncmySessionContext enters and exits using AsyncmyDriver without circular imports.""" + from unittest.mock import AsyncMock, MagicMock + + from sqlspec.adapters.asyncmy.config import AsyncmySessionContext + from sqlspec.adapters.asyncmy.core import default_statement_config + + mock_conn = MagicMock() + acquire = AsyncMock(return_value=mock_conn) + release = AsyncMock() + + ctx = AsyncmySessionContext( + acquire_connection=acquire, + release_connection=release, + statement_config=default_statement_config, + driver_features={}, + prepare_driver=lambda d: d, + ) + async with ctx as driver: + assert driver.connection is mock_conn + assert release.called diff --git a/tests/unit/adapters/test_asyncmy/test_driver.py b/tests/unit/adapters/test_asyncmy/test_driver.py index 454265359..4a4541ea8 100644 --- a/tests/unit/adapters/test_asyncmy/test_driver.py +++ b/tests/unit/adapters/test_asyncmy/test_driver.py @@ -98,3 +98,137 @@ def acquire(self) -> _AcquireCtx: assert cast("object", acquired) is conn assert conn.rolled_back is True assert pool.released == [conn] + + +async def test_asyncmy_execute_cache_hit_invokes_collect_rows_on_clean_query(monkeypatch: pytest.MonkeyPatch) -> None: + """Cached query without unescaped literal percent delegates to fast path and invokes collect_rows.""" + from unittest.mock import AsyncMock, MagicMock + + from sqlspec.core import OperationProfile, ParameterProfile, ProcessedState + from sqlspec.core.cache import CachedQuery + + cursor = MagicMock() + cursor.description = [("id", 3, None, None, None, None, None)] + cursor.fetchall = AsyncMock(return_value=[(1,)]) + cursor.execute = AsyncMock(return_value=None) + + conn = MagicMock() + conn.cursor = MagicMock(return_value=cursor) + + driver = AsyncmyDriver(connection=cast("Any", conn)) + collect_called: list[bool] = [] + original_collect_rows = AsyncmyDriver.collect_rows + + def _spy_collect_rows(*args: Any, **kwargs: Any) -> Any: + collect_called.append(True) + return original_collect_rows(*args, **kwargs) + + monkeypatch.setattr(AsyncmyDriver, "collect_rows", _spy_collect_rows) + + cached = CachedQuery( + compiled_sql="SELECT 1", + parameter_profile=ParameterProfile(), + input_named_parameters=(), + applied_wrap_types=False, + parameter_casts={}, + operation_type="SELECT", + operation_profile=OperationProfile(returns_rows=True, modifies_rows=False), + param_count=0, + processed_state=ProcessedState(compiled_sql="SELECT 1", execution_parameters=[], operation_type="SELECT"), + ) + + result = await driver._execute_cache_hit("SELECT 1", (), cached) + assert bool(collect_called) is True + assert result.data == [(1,)] + + +def test_asyncmy_mapped_exceptions_without_sqlstate() -> None: + """Explicit MySQL errnos map correctly even when exception object lacks sqlstate.""" + from sqlspec.adapters.asyncmy.core import create_mapped_exception + from sqlspec.exceptions import DatabaseConnectionError, PermissionDeniedError, SQLParsingError + + class _FakeErr(Exception): + pass + + assert isinstance(create_mapped_exception(_FakeErr(1054, "Unknown column")), SQLParsingError) + assert isinstance(create_mapped_exception(_FakeErr(1146, "Table doesn't exist")), SQLParsingError) + assert isinstance(create_mapped_exception(_FakeErr(1040, "Too many connections")), DatabaseConnectionError) + assert isinstance( + create_mapped_exception(_FakeErr(1290, "The MySQL server is running with the --read-only option")), + PermissionDeniedError, + ) + assert isinstance(create_mapped_exception(_FakeErr(1836, "Running in read-only mode")), PermissionDeniedError) + + +async def test_asyncmy_dispatch_execute_uses_statement_statement_config() -> None: + """dispatch_execute compiles using statement.statement_config.""" + from unittest.mock import AsyncMock, MagicMock + + from sqlspec import SQL + from sqlspec.core import StatementConfig + + custom_config = StatementConfig(enable_caching=False) + statement = SQL("SELECT 1", statement_config=custom_config) + cursor = MagicMock() + cursor.description = [("1", 3, None, None, None, None, None)] + cursor.execute = AsyncMock() + cursor.fetchall = AsyncMock(return_value=[(1,)]) + cursor.rowcount = 1 + cursor.lastrowid = None + + driver = AsyncmyDriver(connection=MagicMock()) + result = await driver.dispatch_execute(cursor, statement) + assert cursor.execute.called + assert result.data_row_count == 1 + + +async def test_asyncmy_dispatch_select_stream_uses_statement_config() -> None: + """dispatch_select_stream compiles using statement.statement_config.""" + from unittest.mock import MagicMock + + from sqlspec import SQL + from sqlspec.core import StatementConfig + + custom_config = StatementConfig(enable_caching=False) + statement = SQL("SELECT 1", statement_config=custom_config) + driver = AsyncmyDriver(connection=MagicMock()) + stream = driver.dispatch_select_stream(statement, chunk_size=10) + assert stream is not None + + +async def test_asyncmy_dispatch_execute_script_delegates_to_super() -> None: + """dispatch_execute_script emits script statement events via super().""" + from unittest.mock import AsyncMock, MagicMock + + from sqlspec import SQL + + cursor = MagicMock() + cursor.execute = AsyncMock() + cursor.description = None + cursor.rowcount = 0 + cursor.lastrowid = None + + driver = AsyncmyDriver(connection=MagicMock()) + statement = SQL("SELECT 1; SELECT 2").as_script() + result = await driver.dispatch_execute_script(cursor, statement) + assert result.statement_count == 2 + assert result.successful_statements == 2 + + +async def test_asyncmy_stream_source_close_closes_connection_on_cursor_close_error() -> None: + """AsyncmyStreamSource.close closes the connection if cursor.close raises.""" + from unittest.mock import AsyncMock, MagicMock + + from sqlspec.adapters.asyncmy.core import AsyncmyStreamSource + + cursor = MagicMock() + cursor.close = AsyncMock(side_effect=RuntimeError("close boom")) + connection = MagicMock() + driver = MagicMock() + driver.connection = connection + + source = AsyncmyStreamSource(driver, "SELECT 1", (), 10, set()) + source._cursor = cursor + + await source.close() + assert connection.close.called diff --git a/tests/unit/adapters/test_asyncmy/test_litestar_store.py b/tests/unit/adapters/test_asyncmy/test_litestar_store.py index f47c370de..b44d9a91d 100644 --- a/tests/unit/adapters/test_asyncmy/test_litestar_store.py +++ b/tests/unit/adapters/test_asyncmy/test_litestar_store.py @@ -30,3 +30,38 @@ def test_asyncmy_litestar_store_honors_table_options() -> None: ddl = store._table_ddl() assert f"{DEFAULT_ENGINE_CLAUSE} COMMENT='litestar-session'" in ddl + + +def test_asyncmy_litestar_store_drop_table_sql_omits_drop_index() -> None: + """_drop_table_sql must emit only DROP TABLE IF EXISTS without preceding DROP INDEX.""" + store = AsyncmyStore(_mock_config()) + + statements = store._drop_table_sql() + + assert len(statements) == 1 + assert statements[0] == "DROP TABLE IF EXISTS litestar_session" + + +async def test_asyncmy_litestar_store_set_uses_mariadb_compatible_upsert() -> None: + """_set must use VALUES(col) in ON DUPLICATE KEY UPDATE for MySQL and MariaDB compatibility.""" + from unittest.mock import AsyncMock + + cursor = MagicMock() + cursor.execute = AsyncMock() + conn = MagicMock() + conn.cursor = MagicMock(return_value=cursor) + conn.commit = AsyncMock() + conn.__aenter__ = AsyncMock(return_value=conn) + conn.__aexit__ = AsyncMock(return_value=None) + config = _mock_config() + config.provide_connection = lambda *_a, **_k: conn + + store = AsyncmyStore(config) + await store.set("session-1", b"payload", 300) + + assert cursor.execute.call_count == 1 + sql = cursor.execute.call_args_list[0].args[0] + assert "ON DUPLICATE KEY UPDATE" in sql + assert "data = VALUES(data)" in sql + assert "expires_at = VALUES(expires_at)" in sql + assert "AS new" not in sql diff --git a/tests/unit/adapters/test_asyncmy/test_load_from_arrow.py b/tests/unit/adapters/test_asyncmy/test_load_from_arrow.py index 477272e68..621c70542 100644 --- a/tests/unit/adapters/test_asyncmy/test_load_from_arrow.py +++ b/tests/unit/adapters/test_asyncmy/test_load_from_arrow.py @@ -1,8 +1,13 @@ """Asyncmy load_from_arrow ingest paths.""" import asyncio +import tempfile +from collections.abc import Generator +from contextlib import contextmanager from datetime import timedelta from pathlib import Path +from threading import Event +from types import SimpleNamespace from typing import Any, cast import pyarrow as pa @@ -11,6 +16,7 @@ from sqlspec.adapters.asyncmy.config import AsyncmyConfig from sqlspec.adapters.asyncmy.driver import AsyncmyDriver +from sqlspec.exceptions import StorageCapabilityError _CAPS: dict[str, Any] = { "arrow_export_enabled": True, @@ -70,9 +76,7 @@ async def test_load_from_arrow_explicit_bulk_disable_uses_executemany(connection config = AsyncmyConfig( connection_config={"local_infile": connection_opt_in}, driver_features={"enable_local_infile_bulk_load": False} ) - driver = AsyncmyDriver( - connection=cast("Any", conn), driver_features={**config.driver_features, "storage_capabilities": _CAPS} - ) + driver = AsyncmyDriver(connection=cast("Any", conn), driver_features=config.driver_features) job = await driver.load_from_arrow("orders", pa.table({"id": [1, 2], "name": ["a", "b"]})) @@ -93,6 +97,19 @@ async def test_load_from_arrow_overwrite_truncates_first() -> None: assert conn._cursor.executemany_calls[0][0].startswith("INSERT INTO") +async def test_config_preserves_explicit_arrow_import_disable() -> None: + conn = _FakeConnection() + config = AsyncmyConfig( + driver_features=cast("Any", {"storage_capabilities": {**_CAPS, "arrow_import_enabled": False}}) + ) + driver = AsyncmyDriver(connection=cast("Any", conn), driver_features=config.driver_features) + + with pytest.raises(StorageCapabilityError, match="Arrow import"): + await driver.load_from_arrow("orders", pa.table({"id": [1]})) + + assert conn._cursor.executemany_calls == [] + + async def test_load_from_storage_reads_parquet_and_delegates(tmp_path: Path) -> None: parquet_path = tmp_path / "data.parquet" pq.write_table(pa.table({"id": [1, 2], "name": ["a", "b"]}), parquet_path) @@ -142,6 +159,66 @@ async def test_local_infile_failure_discards_connection_and_payload(failure: Bas assert not conn._cursor.payload_path.parent.exists() +@pytest.mark.parametrize("cancel_count", [1, 2]) +@pytest.mark.parametrize("write_fails", [False, True]) +async def test_local_infile_cancelled_writer_removes_payload_after_thread_finishes( + cancel_count: int, write_fails: bool, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Cancellation preserves ownership until the real file writer stops using the payload.""" + started = Event() + release = Event() + finished = Event() + original = tempfile.NamedTemporaryFile + monkeypatch.setattr(tempfile, "tempdir", str(tmp_path)) + + @contextmanager + def blocking_payload(*args: Any, **kwargs: Any) -> Generator[SimpleNamespace, None, None]: + try: + with original(*args, **kwargs) as payload: + + def write(data: bytes) -> int: + written = payload.write(data) + payload.flush() + started.set() + if not release.wait(timeout=5): + msg = "The payload writer was never released." + raise TimeoutError(msg) + if write_fails: + msg = "Payload write failed after cancellation." + raise OSError(msg) + return written + + yield SimpleNamespace(write=write, name=payload.name) + finally: + finished.set() + + monkeypatch.setattr(tempfile, "NamedTemporaryFile", blocking_payload) + conn = _FakeConnection() + driver = AsyncmyDriver( + connection=cast("Any", conn), + driver_features={"storage_capabilities": _CAPS, "enable_local_infile_bulk_load": True}, + ) + task = asyncio.create_task(driver.load_from_arrow("orders", pa.table({"id": [1]}))) + try: + assert await asyncio.to_thread(started.wait, 5) + for _ in range(cancel_count): + task.cancel() + await asyncio.sleep(0) + paths = await asyncio.to_thread(lambda: list(tmp_path.rglob("*.tsv"))) + assert len(paths) == 1 + assert await asyncio.to_thread(paths[0].read_bytes) == b"1\n" + release.set() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(asyncio.shield(task), timeout=5) + assert await asyncio.to_thread(finished.wait, 5) + assert await asyncio.to_thread(lambda: list(tmp_path.iterdir())) == [] + assert conn._cursor.execute_calls == [] + finally: + release.set() + await asyncio.gather(task, return_exceptions=True) + await asyncio.to_thread(finished.wait, 5) + + async def test_local_infile_empty_table_does_not_send_payload() -> None: conn = _FakeConnection() driver = AsyncmyDriver( diff --git a/tests/unit/adapters/test_asyncpg/test_adk_store.py b/tests/unit/adapters/test_asyncpg/test_adk_store.py index 25a519155..fbba4de18 100644 --- a/tests/unit/adapters/test_asyncpg/test_adk_store.py +++ b/tests/unit/adapters/test_asyncpg/test_adk_store.py @@ -312,3 +312,12 @@ async def test_asyncpg_list_sessions_rejects_invalid_options_before_connecting(o await store.list_sessions("app", **options) assert conn.calls == [] + + +async def test_asyncpg_adk_scann_quantizer_validation() -> None: + """Invalid SCaNN quantizer must raise ImproperConfigurationError.""" + settings = AsyncpgADKConfig(vector_index_type="scann", scann_quantizer="INVALID") + config = AsyncpgConfig(extension_config={"adk": settings}) + store = AsyncpgADKMemoryStore(config) + with pytest.raises(ImproperConfigurationError, match="Invalid SCaNN quantizer"): + await store._memory_table_ddl() diff --git a/tests/unit/adapters/test_asyncpg/test_config.py b/tests/unit/adapters/test_asyncpg/test_config.py index b1d0cc43a..05390c655 100644 --- a/tests/unit/adapters/test_asyncpg/test_config.py +++ b/tests/unit/adapters/test_asyncpg/test_config.py @@ -14,14 +14,10 @@ AsyncpgDriverFeatures, AsyncpgPoolConfig, ) -from sqlspec.adapters.asyncpg.core import ( - build_connection_config, - build_postgres_extension_probe_names, - build_statement_config, - resolve_postgres_extension_state, -) +from sqlspec.adapters.asyncpg.core import build_connection_config, build_statement_config from sqlspec.adapters.asyncpg.driver import AsyncpgDriver from sqlspec.core import StatementConfig, StatementStack +from sqlspec.core.config_runtime import build_postgres_extension_probe_names, resolve_postgres_extension_state from sqlspec.exceptions import ImproperConfigurationError @@ -322,6 +318,10 @@ def test_asyncpg_config_normalizes_aliases() -> None: def test_asyncpg_driver_features_contract_retains_configured_keys() -> None: """AsyncpgConfig retains on_connection_create and type_codecs and populates event defaults.""" assert "connection_instance" not in AsyncpgDriverFeatures.__annotations__ + assert "enable_pgcopy_binary" in AsyncpgDriverFeatures.__annotations__ + assert AsyncpgConfig.supports_native_arrow_export is False + assert AsyncpgConfig.supports_native_parquet_export is False + assert AsyncpgConfig.supports_native_parquet_import is False async def hook(_: object) -> None: return None @@ -403,6 +403,9 @@ async def fake_create_pool(*, init: Any = None, **kwargs: Any) -> object: monkeypatch.setattr(config_mod, "asyncpg_create_pool", fake_create_pool) + assert "pgbouncer" not in AsyncpgDriverFeatures.__annotations__ + assert "pgbouncer_mode" not in AsyncpgDriverFeatures.__annotations__ + assert "pgbouncer" not in AsyncpgConnectionConfig.__annotations__ config = AsyncpgConfig( connection_config={"dsn": "postgresql://localhost/test"}, driver_features={"pooler_mode": "transaction"} ) @@ -410,18 +413,6 @@ async def fake_create_pool(*, init: Any = None, **kwargs: Any) -> object: await config._create_pool() # pyright: ignore[reportPrivateUsage] assert captured_connect_kwargs["statement_cache_size"] == 0 - alias_config = AsyncpgConfig( - connection_config={"dsn": "postgresql://localhost/test"}, driver_features={"pgbouncer": True} - ) - assert alias_config.driver_features["pooler_mode"] == "transaction" - assert alias_config.connection_config["statement_cache_size"] == 0 - - mode_alias_config = AsyncpgConfig( - connection_config={"dsn": "postgresql://localhost/test"}, driver_features={"pgbouncer_mode": True} - ) - assert mode_alias_config.driver_features["pooler_mode"] == "transaction" - assert mode_alias_config.connection_config["statement_cache_size"] == 0 - conn = AsyncMock() conn.is_in_transaction = lambda: True conn.execute.return_value = "UPDATE 1" @@ -486,12 +477,12 @@ async def test_asyncpg_columnar_engine_session_vs_transaction() -> None: "google_columnar_engine.enable_columnar_scan": "on" } - auto_pgbouncer_config = AsyncpgConfig( + tx_columnar_config = AsyncpgConfig( connection_config={"dsn": "postgresql://localhost/test"}, - driver_features={"pooler_mode": "auto", "pgbouncer": True, "enable_alloydb_columnar": True}, + driver_features={"pooler_mode": "transaction", "enable_alloydb_columnar": True}, ) - assert auto_pgbouncer_config.driver_features["pooler_mode"] == "transaction" - assert "server_settings" not in auto_pgbouncer_config.connection_config + assert tx_columnar_config.driver_features["pooler_mode"] == "transaction" + assert "server_settings" not in tx_columnar_config.connection_config tx_config = AsyncpgConfig( connection_config={"dsn": "postgresql://localhost/test"}, diff --git a/tests/unit/adapters/test_asyncpg/test_load_from_records.py b/tests/unit/adapters/test_asyncpg/test_load_from_records.py index d5e53fcbd..6243b537b 100644 --- a/tests/unit/adapters/test_asyncpg/test_load_from_records.py +++ b/tests/unit/adapters/test_asyncpg/test_load_from_records.py @@ -166,3 +166,12 @@ def fail_if_called(*args: object) -> None: await driver.load_from_records("events", [(1,)], columns=["id"]) connection.copy_records_to_table.assert_awaited_once() + + +async def test_load_from_records_empty_records_with_overwrite(connection: AsyncMock) -> None: + """Empty records with overwrite=True should truncate the table and return 0-row job.""" + driver = AsyncpgDriver(connection, driver_features=DRIVER_FEATURES) # type: ignore[arg-type] + job = await driver.load_from_records("events", [], overwrite=True) + assert job.telemetry["rows_processed"] == 0 + connection.execute.assert_awaited_once() + assert "TRUNCATE TABLE" in connection.execute.await_args[0][0] diff --git a/tests/unit/adapters/test_asyncpg/test_transaction_state.py b/tests/unit/adapters/test_asyncpg/test_transaction_state.py index 6ac4c9341..88fb20412 100644 --- a/tests/unit/adapters/test_asyncpg/test_transaction_state.py +++ b/tests/unit/adapters/test_asyncpg/test_transaction_state.py @@ -1,8 +1,10 @@ """asyncpg transaction control uses the native transaction handle.""" import inspect +from asyncio import CancelledError from collections.abc import AsyncIterator from contextlib import asynccontextmanager +from pathlib import Path from typing import Any, cast from unittest.mock import AsyncMock, MagicMock @@ -10,16 +12,32 @@ import pytest from sqlspec.adapters.asyncpg.config import AsyncpgConfig -from sqlspec.adapters.asyncpg.core import AsyncpgStreamSource +from sqlspec.adapters.asyncpg.core import AsyncpgStreamSource, build_postgres_storage_params_sql from sqlspec.adapters.asyncpg.driver import AsyncpgDriver from sqlspec.adapters.asyncpg.events.backend import AsyncpgEventsBackend, AsyncpgHybridEventsBackend from sqlspec.adapters.asyncpg.litestar.store import AsyncpgStore from sqlspec.core import SQL, StatementConfig, StatementStack -from sqlspec.exceptions import DeadlockError +from sqlspec.exceptions import DatabaseConnectionError, DeadlockError, QueryTimeoutError pytestmark = pytest.mark.anyio +async def test_storage_bridge_writes_client_converted_rows_without_native_export(tmp_path: Path) -> None: + """Client-side storage export works without claiming native Arrow export.""" + parquet = pytest.importorskip("pyarrow.parquet") + connection = AsyncMock() + connection.fetch.return_value = [{"name": "Ada", "value": 1}] + config = AsyncpgConfig() + driver = AsyncpgDriver(cast("Any", connection), driver_features=config.driver_features) + target = tmp_path / "asyncpg.parquet" + + job = await driver.select_to_storage("SELECT name, value FROM storage_rows", target, format_hint="parquet") + + assert config.storage_capabilities()["arrow_export_enabled"] is False + assert job.telemetry["rows_processed"] == 1 + assert parquet.read_table(target).to_pylist() == [{"name": "Ada", "value": 1}] + + class _FakeTransaction: def __init__(self, record: "list[str]", nested: bool = False) -> None: self._record = record @@ -114,7 +132,8 @@ async def test_a_stream_can_open_a_transaction_after_driver_begin() -> None: async def test_stream_source_reuses_outer_transaction_without_savepoint() -> None: """AsyncpgStreamSource should not start a nested savepoint when already in a transaction.""" connection = _FakeConnection() - connection.cursor = AsyncMock(return_value=MagicMock(fetch=AsyncMock(return_value=[]))) # type: ignore[attr-defined] + cursor = MagicMock(_portal_name="portal", _close_portal=AsyncMock(), fetch=AsyncMock(return_value=[])) + connection.cursor = AsyncMock(return_value=cursor) # type: ignore[attr-defined] driver = AsyncpgDriver(cast("Any", connection)) await driver.begin() @@ -123,10 +142,11 @@ async def test_stream_source_reuses_outer_transaction_without_savepoint() -> Non await source.close() assert connection.transaction_calls == ["start"] + cursor._close_portal.assert_awaited_once_with(None) -async def test_stream_source_suppresses_rollback_error_on_cursor_failure() -> None: - """AsyncpgStreamSource should suppress rollback errors when cursor creation fails.""" +async def test_stream_start_failure_retires_connection_when_rollback_fails() -> None: + """A failed startup rollback retires the socket while preserving the cursor failure.""" tx = MagicMock() tx.start = AsyncMock() tx.rollback = AsyncMock(side_effect=RuntimeError("connection lost")) @@ -134,12 +154,53 @@ async def test_stream_source_suppresses_rollback_error_on_cursor_failure() -> No conn.transaction = MagicMock(return_value=tx) conn.is_in_transaction = MagicMock(return_value=False) conn.cursor = AsyncMock(side_effect=ValueError("cursor failed")) + conn.close = AsyncMock() driver = AsyncpgDriver(cast("Any", conn)) source = AsyncpgStreamSource(driver, "SELECT 1", (), 10) with pytest.raises(ValueError, match="cursor failed"): await source.start() assert source._transaction is None # pyright: ignore[reportPrivateUsage] + conn.close.assert_awaited_once() + + +@pytest.mark.parametrize("owns_transaction", [False, True]) +@pytest.mark.parametrize("error", [False, True]) +async def test_stream_closes_portal_before_owned_transaction(owns_transaction: bool, error: bool) -> None: + """The native portal closes before transaction control without ending a caller transaction.""" + calls: list[str] = [] + connection = MagicMock() + driver = AsyncpgDriver(cast("Any", connection)) + source = AsyncpgStreamSource(driver, "SELECT 1", (), 10, 4.0) + cursor = MagicMock() + cursor._portal_name = "portal" + cursor._close_portal = AsyncMock(side_effect=lambda _: calls.append("portal")) + source._cursor = cursor # pyright: ignore[reportPrivateUsage] + if owns_transaction: + source._transaction = MagicMock( # pyright: ignore[reportPrivateUsage] + commit=AsyncMock(side_effect=lambda: calls.append("commit")), + rollback=AsyncMock(side_effect=lambda: calls.append("rollback")), + ) + + await source.close(error=error) + + assert calls == (["portal", "rollback" if error else "commit"] if owns_transaction else ["portal"]) + cursor._close_portal.assert_awaited_once_with(4.0) + + +@pytest.mark.parametrize("failure", [RuntimeError("portal close failed"), CancelledError()]) +async def test_stream_cleanup_failure_retires_connection(failure: BaseException) -> None: + """Failed or cancelled portal cleanup cannot return a dirty socket to the pool.""" + connection = MagicMock(close=AsyncMock()) + driver = AsyncpgDriver(cast("Any", connection)) + source = AsyncpgStreamSource(driver, "SELECT 1", (), 10) + source._cursor = MagicMock( # pyright: ignore[reportPrivateUsage] + _portal_name="portal", _close_portal=AsyncMock(side_effect=failure) + ) + + with pytest.raises(type(failure)): + await source.close(error=True) + connection.close.assert_awaited_once() async def test_begin_resets_stale_transaction_and_maps_exceptions() -> None: @@ -292,3 +353,187 @@ async def test_create_connection_drops_every_pool_only_setting(monkeypatch: pyte accepted = set(inspect.signature(asyncpg.connect).parameters) assert set(connect.call_args.kwargs) <= accepted + + +async def test_asyncpg_begin_rollback_on_set_local_failure() -> None: + """A failure during SET LOCAL in begin() must rollback the native transaction and raise mapped error.""" + conn = MagicMock() + tx = AsyncMock() + conn.transaction.return_value = tx + conn.execute = AsyncMock(side_effect=RuntimeError("columnar scan configuration error")) + conn.is_in_transaction.return_value = False + + driver = AsyncpgDriver( + cast("Any", conn), driver_features={"enable_alloydb_columnar": True, "pooler_mode": "transaction"} + ) + with pytest.raises(Exception): + await driver.begin() + assert driver._transaction is None + tx.rollback.assert_awaited_once() + + +async def test_asyncpg_dispatch_execute_script_with_parameters_and_observers() -> None: + """dispatch_execute_script executes sub-statements, emits observer events, and returns cursor.""" + conn = MagicMock() + conn.execute = AsyncMock(return_value="CREATE 1") + driver = AsyncpgDriver(cast("Any", conn)) + + stmt = SQL("CREATE TABLE t1 (id int); CREATE TABLE t2 (id int);") + result = await driver.dispatch_execute_script(cast("Any", conn), stmt) + + assert result.is_script_result is True + assert result.statement_count == 2 + assert result.successful_statements == 2 + assert result.cursor_result is conn + + +async def test_asyncpg_execute_stack_in_outer_transaction_does_not_commit_or_rollback() -> None: + """execute_stack with continue_on_error inside an outer transaction must not commit or rollback.""" + conn = MagicMock() + conn.is_in_transaction = MagicMock(return_value=True) + + async def execute(sql: str) -> str: + if sql == "UPDATE t SET x = 2": + raise RuntimeError("failed op") + return "UPDATE 1" if sql.startswith("UPDATE") else "" + + conn.execute = AsyncMock(side_effect=execute) + + driver = AsyncpgDriver(cast("Any", conn)) + driver._transaction = MagicMock() + + stack = ( + StatementStack() + .push_execute("UPDATE t SET x = 1") + .push_execute("UPDATE t SET x = 2") + .push_execute("UPDATE t SET x = 3") + ) + + results = await driver.execute_stack(stack, continue_on_error=True) + assert len(results) == 3 + assert results[0].result is not None + assert results[1].error is not None + assert results[2].rows_affected == 1 + assert driver._transaction is not None + commands = [call.args[0] for call in conn.execute.await_args_list] + savepoints = [command.removeprefix("SAVEPOINT ") for command in commands if command.startswith("SAVEPOINT ")] + assert len(savepoints) == 3 + assert commands == [ + f"SAVEPOINT {savepoints[0]}", + "UPDATE t SET x = 1", + f"RELEASE SAVEPOINT {savepoints[0]}", + f"SAVEPOINT {savepoints[1]}", + "UPDATE t SET x = 2", + f"ROLLBACK TO SAVEPOINT {savepoints[1]}", + f"RELEASE SAVEPOINT {savepoints[1]}", + f"SAVEPOINT {savepoints[2]}", + "UPDATE t SET x = 3", + f"RELEASE SAVEPOINT {savepoints[2]}", + ] + + +async def test_asyncpg_exception_mapping_timeout_and_interface_errors() -> None: + """AsyncpgExceptionHandler and create_mapped_exception map TimeoutError and InterfaceError/OSError.""" + driver = AsyncpgDriver(cast("Any", MagicMock())) + + handler = driver.handle_database_exceptions() + with pytest.raises(QueryTimeoutError): + async with handler: + raise TimeoutError("query timed out") + driver._check_pending_exception(handler) + + handler2 = driver.handle_database_exceptions() + with pytest.raises(DatabaseConnectionError): + async with handler2: + raise OSError("connection reset by peer") + driver._check_pending_exception(handler2) + + +async def test_asyncpg_truncate_and_copy_map_errors() -> None: + """Truncate and copy operations map errors rather than leaking unmapped exceptions.""" + conn = MagicMock() + conn.execute = AsyncMock(side_effect=RuntimeError("truncate table failed")) + driver = AsyncpgDriver(cast("Any", conn)) + + with pytest.raises(Exception) as exc_info: + await driver.load_from_records("users", [{"id": 1}], overwrite=True) + assert not isinstance(exc_info.value, RuntimeError) or isinstance(exc_info.value, DatabaseConnectionError) + + +def test_asyncpg_build_postgres_storage_params_sql() -> None: + """build_postgres_storage_params_sql consolidates storage and autovacuum clauses.""" + settings = {"fillfactor": 85, "autovacuum_vacuum_scale_factor": 0.05, "autovacuum_analyze_scale_factor": 0.02} + clause = build_postgres_storage_params_sql(settings, include_autovacuum=True) + assert "WITH (" in clause + assert "fillfactor = 85" in clause + assert "autovacuum_vacuum_scale_factor = 0.05" in clause + assert "autovacuum_analyze_scale_factor = 0.02" in clause + + +def test_asyncpg_events_validate_listener_pool_capacity_skips_when_direct_dsn() -> None: + """_validate_listener_pool_capacity should not enforce pool size when direct_dsn is configured.""" + from sqlspec.adapters.asyncpg.events.backend import _validate_listener_pool_capacity + + config = AsyncpgConfig( + connection_config={"dsn": "postgresql://localhost/test", "max_size": 1}, + driver_features={"direct_dsn": "postgresql://localhost/direct"}, + ) + _validate_listener_pool_capacity(config) + + +def test_asyncpg_events_parse_notify_payload_handles_malformed_json() -> None: + """Malformed NOTIFY payload does not raise and returns None safely.""" + from sqlspec.adapters.asyncpg.events.backend import ( + _extract_batch_size, + _extract_event_id, + _extract_marker_id, + _parse_notify_payload, + ) + + bad_payload = "not valid json {{" + assert _parse_notify_payload(bad_payload) is None + assert _extract_event_id(bad_payload) is None + assert _extract_batch_size(bad_payload) == 0 + assert _extract_marker_id(bad_payload) is None + + good_payload = '{"event_id": "e1", "batch_size": 5, "marker_id": "m1"}' + parsed = _parse_notify_payload(good_payload) + assert parsed is not None + assert _extract_event_id(parsed) == "e1" + assert _extract_batch_size(parsed) == 5 + assert _extract_marker_id(parsed) == "m1" + + +def test_asyncpg_events_resolve_listener_config_allows_mapping_direct_dsn() -> None: + """_resolve_listener_config accepts Mapping for direct_dsn without raising.""" + from sqlspec.adapters.asyncpg.events._hub import AsyncpgListenerHub + + config = AsyncpgConfig( + connection_config={"dsn": "postgresql://localhost/test"}, + driver_features={"pooler_mode": "transaction", "direct_dsn": {"host": "localhost", "database": "direct_db"}}, + ) + hub = AsyncpgListenerHub(config) + resolved = hub._resolve_listener_config() + assert resolved is not None + + +async def test_asyncpg_cache_hit_with_cursor() -> None: + """AsyncpgDriver executes cached queries directly on connection via _execute_cache_hit_with_cursor.""" + conn = MagicMock() + conn.fetch = AsyncMock(return_value=[{"val": 42}]) + conn.execute = AsyncMock(return_value="UPDATE 1") + + driver = AsyncpgDriver(cast("Any", conn)) + res1 = await driver.execute("SELECT 42") + assert res1.data == [{"val": 42}] + + res2 = await driver.execute("SELECT 42") + assert res2.data == [{"val": 42}] + assert conn.fetch.call_count == 2 + + upd1 = await driver.execute("UPDATE t SET x = 1") + assert upd1.rows_affected == 1 + + upd2 = await driver.execute("UPDATE t SET x = 1") + assert upd2.rows_affected == 1 + assert conn.execute.call_count == 2 diff --git a/tests/unit/adapters/test_asyncpg/test_type_handlers.py b/tests/unit/adapters/test_asyncpg/test_type_handlers.py index 6b360b625..79e24e2a4 100644 --- a/tests/unit/adapters/test_asyncpg/test_type_handlers.py +++ b/tests/unit/adapters/test_asyncpg/test_type_handlers.py @@ -6,8 +6,14 @@ import pytest from sqlspec.adapters.asyncpg.config import register_json_codecs, register_pgvector_support -from sqlspec.adapters.asyncpg.core import create_mapped_exception, invoke_prepared_statement -from sqlspec.exceptions import OperationCancelledError, PermissionDeniedError, QueryTimeoutError, UniqueViolationError +from sqlspec.adapters.asyncpg.core import create_mapped_exception +from sqlspec.exceptions import ( + DatabaseConnectionError, + OperationCancelledError, + PermissionDeniedError, + QueryTimeoutError, + UniqueViolationError, +) from sqlspec.utils.serializers import from_json, to_json @@ -42,31 +48,34 @@ async def test_register_json_codecs_success() -> None: assert json_call.kwargs["encoder"]({"value": 1}) == b'{"value":1}' assert json_call.kwargs["encoder"]("console") == b'"console"' + assert json_call.kwargs["encoder"]("123") == b"123" + assert json_call.kwargs["encoder"]("true") == b"true" + assert json_call.kwargs["encoder"]("null") == b"null" + assert json_call.kwargs["encoder"]('"text"') == b'"text"' assert json_call.kwargs["encoder"]('{"pre":"serialized"}') == b'{"pre":"serialized"}' assert json_call.kwargs["encoder"](b'\x01{"value":1}') == b'{"value":1}' assert json_call.kwargs["decoder"](b'{"value":1}') == {"value": 1} assert jsonb_call.kwargs["encoder"]({"value": 1}) == b'\x01{"value":1}' assert jsonb_call.kwargs["encoder"]("console") == b'\x01"console"' + assert jsonb_call.kwargs["encoder"]("123") == b"\x01123" + assert jsonb_call.kwargs["encoder"]("true") == b"\x01true" + assert jsonb_call.kwargs["encoder"]("null") == b"\x01null" + assert jsonb_call.kwargs["encoder"]('"text"') == b'\x01"text"' assert jsonb_call.kwargs["encoder"]('{"pre":"serialized"}') == b'\x01{"pre":"serialized"}' assert jsonb_call.kwargs["encoder"](b'\x01{"value":1}') == b'\x01{"value":1}' assert jsonb_call.kwargs["decoder"](b'\x01{"value":1}') == {"value": 1} @pytest.mark.anyio -async def test_invoke_prepared_statement_unpacks_dict_parameters_positionally() -> None: - """Dict parameters passed to invoke_prepared_statement should unpack values positionally.""" - prepared = MagicMock() - prepared.fetch = AsyncMock(return_value=[{"id": 1}]) - prepared.get_statusmsg = MagicMock(return_value="UPDATE 2") - - rows = await invoke_prepared_statement(prepared, {"p1": "alpha", "p2": 42}, fetch=True, timeout=1.5) - assert rows == [{"id": 1}] - prepared.fetch.assert_awaited_once_with("alpha", 42, timeout=1.5) - - prepared.fetch.reset_mock() - status = await invoke_prepared_statement(prepared, {"p1": "beta"}, fetch=False, timeout=2.0) - assert status == "UPDATE 2" - prepared.fetch.assert_awaited_once_with("beta", timeout=2.0) +async def test_asyncpg_connection_error_mapping() -> None: + """ConnectionDoesNotExistError and CannotConnectNowError map to DatabaseConnectionError.""" + exc1 = asyncpg.exceptions.ConnectionDoesNotExistError("connection lost") + mapped1 = create_mapped_exception(exc1) + assert isinstance(mapped1, DatabaseConnectionError) + + exc2 = asyncpg.exceptions.CannotConnectNowError("server starting up") + mapped2 = create_mapped_exception(exc2) + assert isinstance(mapped2, DatabaseConnectionError) @pytest.mark.anyio diff --git a/tests/unit/adapters/test_bigquery/test_config.py b/tests/unit/adapters/test_bigquery/test_config.py index d317260fc..006336e98 100644 --- a/tests/unit/adapters/test_bigquery/test_config.py +++ b/tests/unit/adapters/test_bigquery/test_config.py @@ -114,7 +114,7 @@ def test_bigquery_config_routes_current_client_level_settings(monkeypatch: Monke def test_bigquery_config_typed_surfaces_do_not_advertise_inert_settings() -> None: """Typed public settings should only include options that SQLSpec routes.""" assert "credentials_path" not in BigQueryConnectionParams.__annotations__ - assert "use_query_and_wait" in BigQueryDriverFeatures.__annotations__ + assert "enable_query_and_wait" in BigQueryDriverFeatures.__annotations__ assert "enable_storage_write_api" in BigQueryDriverFeatures.__annotations__ assert "enable_storage_write_api" not in BigQueryConnectionParams.__annotations__ assert "on_job_start" not in BigQueryDriverFeatures.__annotations__ diff --git a/tests/unit/adapters/test_bigquery/test_core.py b/tests/unit/adapters/test_bigquery/test_core.py index 7383e5aac..cfbf1184e 100644 --- a/tests/unit/adapters/test_bigquery/test_core.py +++ b/tests/unit/adapters/test_bigquery/test_core.py @@ -5,10 +5,11 @@ from contextlib import nullcontext from datetime import UTC, datetime from types import SimpleNamespace -from typing import Any, cast +from typing import TYPE_CHECKING, Any, cast from unittest.mock import patch import pytest +from google.api_core.exceptions import RetryError from google.cloud import bigquery from google.cloud.bigquery import QueryJobConfig from google.cloud.exceptions import Forbidden, GoogleCloudError @@ -33,7 +34,7 @@ try_bulk_insert, ) from sqlspec.adapters.bigquery.data_dictionary import BigQueryDataDictionary -from sqlspec.adapters.bigquery.driver import BigQueryDriver +from sqlspec.adapters.bigquery.driver import BigQueryDriver, BigQueryExceptionHandler from sqlspec.adapters.bigquery.events.store import BigQueryEventQueueStore from sqlspec.adapters.bigquery.litestar.store import BigQueryStore from sqlspec.core import ParameterStyle, TypedParameter @@ -48,6 +49,9 @@ ) from sqlspec.utils.serializers import to_json +if TYPE_CHECKING: + from collections.abc import Callable + class _RecordingConnection: def __init__(self) -> None: @@ -920,3 +924,19 @@ def _cloud_error(message: str, *, code: int, errors: list[dict[str, str]] | None err.code = code err._errors = errors or [] return cast("GoogleCloudError", err) + + +def test_retry_error_mapped_to_query_timeout() -> None: + """RetryError maps to QueryTimeoutError in BigQuery create_mapped_exception.""" + err = cast("Callable[[str, Exception | None], RetryError]", RetryError)("Deadline exceeded waiting for query", None) + mapped = create_mapped_exception(err) + assert isinstance(mapped, QueryTimeoutError) + + +def test_exception_handler_catches_google_api_error() -> None: + """BigQueryExceptionHandler catches GoogleAPIError and maps it.""" + handler = BigQueryExceptionHandler() + err = cast("Callable[[str, Exception | None], RetryError]", RetryError)("Deadline exceeded", None) + with handler: + raise err + assert isinstance(handler.pending_exception, QueryTimeoutError) diff --git a/tests/unit/adapters/test_bigquery/test_job_controls.py b/tests/unit/adapters/test_bigquery/test_job_controls.py index 6a3f8db83..ffb43f073 100644 --- a/tests/unit/adapters/test_bigquery/test_job_controls.py +++ b/tests/unit/adapters/test_bigquery/test_job_controls.py @@ -130,7 +130,7 @@ def test_run_query_job_passes_job_id_prefix_only_without_job_id() -> None: def test_query_and_wait_used_when_enabled() -> None: connection = _RecordingConnection() connection.row_iterator = _RecordingRowIterator(rows=[{"v": 1}], schema=_schema("v")) - driver = BigQueryDriver(cast(Any, connection), driver_features={"use_query_and_wait": True}) + driver = BigQueryDriver(cast(Any, connection), driver_features={"enable_query_and_wait": True}) result = driver.execute("SELECT 1 AS v") @@ -144,7 +144,7 @@ def test_query_and_wait_used_when_enabled() -> None: def test_query_and_wait_dml_rowcount() -> None: connection = _RecordingConnection() connection.row_iterator = _RecordingRowIterator(num_dml_affected_rows=3) - driver = BigQueryDriver(cast(Any, connection), driver_features={"use_query_and_wait": True}) + driver = BigQueryDriver(cast(Any, connection), driver_features={"enable_query_and_wait": True}) result = driver.execute("UPDATE t SET v = 1") @@ -156,7 +156,7 @@ def test_query_and_wait_explicit_off() -> None: connection = _RecordingConnection() connection.query_job = _RecordingJob(rows=[{"v": 1}], schema=_schema("v")) connection.row_iterator = _RecordingRowIterator(rows=[{"v": 1}], schema=_schema("v")) - driver = BigQueryDriver(cast(Any, connection), driver_features={"use_query_and_wait": False}) + driver = BigQueryDriver(cast(Any, connection), driver_features={"enable_query_and_wait": False}) result = driver.execute("SELECT 1 AS v") @@ -345,7 +345,11 @@ def test_select_to_storage_uses_existing_export_surface_and_job_controls(tmp_pat connection.query_job = _RecordingJob(rows=[{"id": 1}], schema=_schema("id")) driver = BigQueryDriver( cast(Any, connection), - driver_features={"job_result_timeout": 5.0, "storage_capabilities": CAPABILITIES, "use_query_and_wait": False}, + driver_features={ + "job_result_timeout": 5.0, + "storage_capabilities": CAPABILITIES, + "enable_query_and_wait": False, + }, ) job = driver.select_to_storage("SELECT 1 AS id", tmp_path / "result.parquet", format_hint="parquet") diff --git a/tests/unit/adapters/test_bigquery/test_storage_write_api.py b/tests/unit/adapters/test_bigquery/test_storage_write_api.py index 03a01c999..f658910c7 100644 --- a/tests/unit/adapters/test_bigquery/test_storage_write_api.py +++ b/tests/unit/adapters/test_bigquery/test_storage_write_api.py @@ -6,12 +6,13 @@ import pyarrow as pa import pytest +from google.cloud.bigquery import SchemaField from google.cloud.bigquery_storage_v1 import types import sqlspec.adapters.bigquery.config as bigquery_config from sqlspec.adapters.bigquery import BigQueryConfig from sqlspec.adapters.bigquery.core import build_arrow_write_stream_payload -from sqlspec.adapters.bigquery.driver import BigQueryDriver +from sqlspec.adapters.bigquery.driver import BigQueryDriver, _bigquery_arrow_reader_from_iterable from sqlspec.exceptions import ImproperConfigurationError, StorageOperationFailedError CAPABILITIES = { @@ -211,10 +212,7 @@ def test_import_failure_falls_back_to_parquet(monkeypatch: pytest.MonkeyPatch) - driver_features={"enable_storage_write_api": True, "storage_capabilities": CAPABILITIES}, ) - def _raise(_self: Any, _table: str, _arrow: Any) -> Any: - raise ImportError("bigquery_storage not available") - - monkeypatch.setattr(BigQueryDriver, "_load_arrow_via_storage_write_api", _raise) + monkeypatch.setattr("sqlspec.adapters.bigquery.driver.BigQueryStorageWriteModule", None) driver.load_from_arrow("dataset.table", pa.table({"id": [1]})) @@ -330,3 +328,14 @@ def fail_append(self: Any, **kwargs: Any) -> Any: client = _FakeWriteClient.instances[0] assert client.finalize_calls == ["projects/proj/datasets/dataset/tables/table/streams/s1"] assert client.commit_calls == [] + + +def test_arrow_reader_from_iterable_empty_with_schema() -> None: + """_bigquery_arrow_reader_from_iterable builds typed schema from field list when empty.""" + fields = [SchemaField("id", "INTEGER"), SchemaField("name", "STRING")] + reader = _bigquery_arrow_reader_from_iterable([], schema=fields) + + assert reader.schema.names == ["id", "name"] + assert reader.schema.field("id").type == pa.int64() + assert reader.schema.field("name").type == pa.string() + assert reader.read_all().num_rows == 0 diff --git a/tests/unit/adapters/test_cockroach_asyncpg/test_adk_store.py b/tests/unit/adapters/test_cockroach_asyncpg/test_adk_store.py index 55f91e177..026708cc1 100644 --- a/tests/unit/adapters/test_cockroach_asyncpg/test_adk_store.py +++ b/tests/unit/adapters/test_cockroach_asyncpg/test_adk_store.py @@ -1,11 +1,14 @@ # pyright: reportPrivateUsage=false """Unit tests for CockroachDB asyncpg ADK store extension configuration.""" +from datetime import UTC, datetime from typing import Any, NotRequired, Self, cast, get_args, get_origin -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock +from uuid import UUID import pytest +from sqlspec.adapters.cockroach_asyncpg import CockroachAsyncpgConfig from sqlspec.adapters.cockroach_asyncpg.adk import ( CockroachAsyncpgADKConfig, CockroachAsyncpgADKMemoryStore, @@ -197,3 +200,79 @@ async def test_cockroach_asyncpg_list_sessions_rejects_invalid_options_before_co await store.list_sessions("app", **options) assert conn.calls == [] + + +def test_cockroach_asyncpg_adk_hash_shard_bucket_count_validation() -> None: + """Invalid hash_shard_bucket_count must raise ValueError matching litestar store.""" + from sqlspec.adapters.cockroach_asyncpg.adk.store import _cockroach_hash_shard_clause + + bad_config = CockroachAsyncpgADKConfig(enable_hash_sharded_indexes=True, hash_shard_bucket_count=-1) + with pytest.raises(ValueError, match="hash_shard_bucket_count"): + _cockroach_hash_shard_clause(bad_config) + + +@pytest.mark.parametrize("use_fts", [False, True]) +async def test_cockroach_asyncpg_adk_memory_search_with_embedding(use_fts: bool) -> None: + """Embedding arguments retain text search against the canonical memory schema.""" + executed_sql: list[tuple[str, tuple[Any, ...]]] = [] + + class _FakeConn: + async def fetch(self, sql: str, *params: Any) -> list[Any]: + executed_sql.append((sql, params)) + return [] + + class _FakeCtx: + async def __aenter__(self) -> _FakeConn: + return _FakeConn() + + async def __aexit__(self, *_args: Any) -> None: + return None + + config = CockroachAsyncpgConfig(extension_config={"adk": {"enable_memory": True, "memory_use_fts": use_fts}}) + object.__setattr__(config, "provide_connection", MagicMock(return_value=_FakeCtx())) + store = CockroachAsyncpgADKMemoryStore(config) + + await store.search_entries(query="hello", app_name="app", user_id="u1", embedding=[0.1, 0.2, 0.3]) + assert len(executed_sql) == 1 + sql, params = executed_sql[0] + assert "plainto_tsquery" in sql if use_fts else "content_text ILIKE" in sql + assert params == ("app", "u1", "hello" if use_fts else "%hello%", 20) + + assert await store.search_entries("", "app", "u1", embedding=[0.1]) == [] + assert len(executed_sql) == 1 + + +@pytest.mark.parametrize(("owner_type", "owner_id"), [("INTEGER", 42), ("UUID", UUID(int=1))]) +async def test_cockroach_asyncpg_memory_batch_preserves_owner_type( + owner_type: str, owner_id: object, monkeypatch: pytest.MonkeyPatch +) -> None: + """A scalar owner bind retains integer and UUID ownership contracts.""" + conn = MagicMock() + conn.execute = AsyncMock(return_value="INSERT 0 1") + context = MagicMock() + context.__aenter__.return_value = conn + monkeypatch.setattr(CockroachAsyncpgConfig, "provide_connection", lambda *_a, **_k: context) + config = CockroachAsyncpgConfig(extension_config={"adk": {"owner_id_column": f"owner_id {owner_type}"}}) + now = datetime.now(UTC) + entry = { + "id": "mem-1", + "session_id": "sess-1", + "app_name": "app", + "user_id": "u1", + "scope": "user", + "event_id": "evt-1", + "author": "user", + "timestamp": now, + "content_json": {"text": "hello"}, + "content_text": "hello", + "metadata_json": None, + "inserted_at": now, + "embedding": None, + } + + assert await CockroachAsyncpgADKMemoryStore(config).insert_memory_entries([cast("Any", entry)], owner_id) == 1 + + sql, *params = conn.execute.call_args.args + assert "author, $13, timestamp" in _normalized(sql) + assert params[-1] == owner_id + assert params[0] == ["mem-1"] diff --git a/tests/unit/adapters/test_cockroach_asyncpg/test_config.py b/tests/unit/adapters/test_cockroach_asyncpg/test_config.py index 93eb123a0..67f90dc0d 100644 --- a/tests/unit/adapters/test_cockroach_asyncpg/test_config.py +++ b/tests/unit/adapters/test_cockroach_asyncpg/test_config.py @@ -549,3 +549,142 @@ async def test_cockroach_asyncpg_adk_locality_validation_and_schema_qualified_dd litestar_store = CockroachAsyncpgStore(good_config) litestar_store._table_name = "tenant_a.litestar_sessions" assert "idx_tenant_a_litestar_sessions_expires_at" in litestar_store._table_ddl() + + +async def test_cockroach_asyncpg_follower_read_only_for_readonly_statements() -> None: + """CockroachAsyncpgDriver should not begin follower-read transactions for modifying operations.""" + in_tx = False + + async def _in_tx() -> bool: + return in_tx + + async def _begin(self: Any) -> None: + nonlocal in_tx + in_tx = True + + with ( + patch.object(CockroachAsyncpgDriver, "_connection_in_transaction", _in_tx), + patch.object(CockroachAsyncpgDriver, "begin", _begin), + ): + conn = MagicMock() + conn.fetch = AsyncMock(return_value=[]) + driver = CockroachAsyncpgDriver( + connection=conn, driver_features={"enable_follower_reads": True, "default_staleness": "'-10s'"} + ) + modifying_stmt = SQL("INSERT INTO users (name) VALUES ('alice') RETURNING id") + await driver.dispatch_execute(conn, modifying_stmt) + assert in_tx is False + + +async def test_cockroach_asyncpg_begin_follower_read_failure_handling() -> None: + """Failing SET TRANSACTION in begin() must rollback and raise mapped exception.""" + conn = MagicMock() + conn.is_in_transaction.return_value = False + tx = AsyncMock() + conn.transaction.return_value = tx + conn.execute = AsyncMock(side_effect=RuntimeError("syntax error in staleness")) + + driver = CockroachAsyncpgDriver( + connection=conn, driver_features={"enable_follower_reads": True, "default_staleness": "'-10s'"} + ) + with pytest.raises(Exception): + await driver.begin() + assert driver._transaction is None + tx.rollback.assert_awaited_once() + + +async def test_cockroach_asyncpg_select_stream_follower_reads() -> None: + """select_stream on a SELECT query with follower reads emits the staleness clause and closes.""" + executed_statements: list[str] = [] + conn = MagicMock() + conn.is_in_transaction.return_value = False + tx = AsyncMock() + conn.transaction.return_value = tx + + async def fake_execute(query: str, *args: Any, **kwargs: Any) -> str: + executed_statements.append(query) + return "OK" + + conn.execute = AsyncMock(side_effect=fake_execute) + + cursor_mock = AsyncMock() + cursor_mock.fetch = AsyncMock(side_effect=[[{"id": 1}], []]) + conn.cursor = AsyncMock(return_value=cursor_mock) + + driver = CockroachAsyncpgDriver( + connection=conn, driver_features={"enable_follower_reads": True, "default_staleness": "'-10s'"} + ) + stream = driver.dispatch_select_stream(SQL("SELECT id FROM users"), chunk_size=10) + assert stream is not None + rows = [row async for row in stream] + assert rows == [{"id": 1}] + assert any("AS OF SYSTEM TIME '-10s'" in s for s in executed_statements) + tx.commit.assert_awaited() + + +async def test_cockroach_asyncpg_get_columns_and_indexes_by_table() -> None: + """CockroachAsyncpgDataDictionary executes by_table query when table is specified.""" + from sqlspec.adapters.cockroach_asyncpg.data_dictionary import CockroachAsyncpgDataDictionary + + calls: list[tuple[Any, dict[str, Any]]] = [] + driver = MagicMock() + + async def fake_select(query: Any, **kwargs: Any) -> list[Any]: + calls.append((query, kwargs)) + return [] + + driver.select = AsyncMock(side_effect=fake_select) + dd = CockroachAsyncpgDataDictionary() + + await dd.get_columns(driver, table="users", schema=None) + assert len(calls) == 1 + assert calls[0][1]["table_name"] == "users" + assert calls[0][1]["schema_name"] is None + + calls.clear() + await dd.get_indexes(driver, table="users", schema=None) + assert len(calls) == 1 + assert calls[0][1]["table_name"] == "users" + assert calls[0][1]["schema_name"] is None + + +async def test_cockroach_asyncpg_cache_hit_with_cursor() -> None: + """CockroachAsyncpgDriver executes cached queries directly on connection via _execute_cache_hit.""" + conn = MagicMock() + conn.fetch = AsyncMock(return_value=[{"val": 42}]) + conn.execute = AsyncMock(return_value="UPDATE 1") + + driver = CockroachAsyncpgDriver(cast("Any", conn)) + res1 = await driver.execute("SELECT 42") + assert res1.data == [{"val": 42}] + + res2 = await driver.execute("SELECT 42") + assert res2.data == [{"val": 42}] + assert conn.fetch.call_count == 2 + + upd1 = await driver.execute("UPDATE t SET x = 1") + assert upd1.rows_affected == 1 + + upd2 = await driver.execute("UPDATE t SET x = 1") + assert upd2.rows_affected == 1 + assert conn.execute.call_count == 2 + + +def test_cockroach_asyncpg_connection_config_timeout_annotations() -> None: + """CockroachAsyncpgConnectionConfig exposes statement_timeout and idle_in_transaction_session_timeout.""" + annotations = CockroachAsyncpgConnectionConfig.__annotations__ + assert "statement_timeout" in annotations + assert "idle_in_transaction_session_timeout" in annotations + + +async def test_cockroach_asyncpg_execute_native_storage_under_active_exception() -> None: + """_execute_native_storage inside active except block does not pollute pending exception.""" + conn = MagicMock() + conn.fetch = AsyncMock(return_value=[{"rows": 1}]) + driver = CockroachAsyncpgDriver(cast("Any", conn)) + + try: + raise ValueError("outer error") + except ValueError: + rows = await driver._execute_native_storage("SELECT 1", []) + assert rows == [{"rows": 1}] diff --git a/tests/unit/adapters/test_cockroach_psycopg/test_adk_store.py b/tests/unit/adapters/test_cockroach_psycopg/test_adk_store.py index 644270e86..6d3e31262 100644 --- a/tests/unit/adapters/test_cockroach_psycopg/test_adk_store.py +++ b/tests/unit/adapters/test_cockroach_psycopg/test_adk_store.py @@ -1,11 +1,14 @@ # pyright: reportPrivateUsage=false """Unit tests for CockroachDB psycopg ADK store extension configuration.""" +from datetime import UTC, datetime from typing import Any, NotRequired, Self, cast, get_args, get_origin -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock +from uuid import UUID import pytest +from sqlspec.adapters.cockroach_psycopg import CockroachPsycopgAsyncConfig, CockroachPsycopgSyncConfig from sqlspec.adapters.cockroach_psycopg.adk import ( CockroachPsycopgADKConfig, CockroachPsycopgAsyncADKMemoryStore, @@ -319,3 +322,86 @@ def test_cockroach_psycopg_sync_list_sessions_rejects_invalid_options(options: " store.list_sessions("app", **options) assert cursor.calls == [] + + +def _memory_store_with_cursor(async_mode: bool, adk_config: dict[str, Any], monkeypatch: pytest.MonkeyPatch) -> Any: + """Provide real stores with SDK cursor and connection boundaries substituted.""" + config_type = CockroachPsycopgAsyncConfig if async_mode else CockroachPsycopgSyncConfig + cursor = MagicMock() + cursor.rowcount = 1 + cursor.description = [] + cursor.fetchall.return_value = [] + conn = MagicMock() + context = MagicMock() + context.__enter__.return_value = conn + context.__aenter__.return_value = conn + cursor.__enter__.return_value = cursor + cursor.__aenter__.return_value = cursor + conn.cursor.return_value = cursor + if async_mode: + cursor.execute = AsyncMock() + cursor.fetchall = AsyncMock(return_value=[]) + conn.commit = AsyncMock() + monkeypatch.setattr(config_type, "provide_connection", lambda *_a, **_k: context) + if async_mode: + return CockroachPsycopgAsyncADKMemoryStore( + CockroachPsycopgAsyncConfig(extension_config={"adk": adk_config}) + ), cursor + return CockroachPsycopgSyncADKMemoryStore(CockroachPsycopgSyncConfig(extension_config={"adk": adk_config})), cursor + + +@pytest.mark.parametrize("async_mode", [False, True]) +@pytest.mark.parametrize(("owner_type", "owner_id"), [("INTEGER", 42), ("UUID", UUID(int=1))]) +async def test_cockroach_psycopg_memory_batch_preserves_owner_type( + async_mode: bool, owner_type: str, owner_id: object, monkeypatch: pytest.MonkeyPatch +) -> None: + """Sync and async batches bind one native owner scalar outside UNNEST.""" + store, cursor = _memory_store_with_cursor(async_mode, {"owner_id_column": f"owner_id {owner_type}"}, monkeypatch) + now = datetime.now(UTC) + entry = { + "id": "mem-1", + "session_id": "sess-1", + "app_name": "app", + "user_id": "u1", + "scope": "user", + "event_id": "evt-1", + "author": "user", + "timestamp": now, + "content_json": {"text": "hello"}, + "content_text": "hello", + "metadata_json": None, + "inserted_at": now, + "embedding": None, + } + inserted = store.insert_memory_entries([entry], owner_id) + if async_mode: + inserted = await inserted + assert inserted == 1 + + query, params = cursor.execute.call_args.args + assert "author, %s, timestamp" in _normalized(query.as_string()) + assert params[0] == owner_id + assert params[1] == ["mem-1"] + + +@pytest.mark.parametrize("async_mode", [False, True]) +@pytest.mark.parametrize("use_fts", [False, True]) +async def test_cockroach_psycopg_memory_embedding_keeps_text_search( + async_mode: bool, use_fts: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + """Embedding arguments keep searching text columns present in the memory schema.""" + store, cursor = _memory_store_with_cursor(async_mode, {"memory_use_fts": use_fts}, monkeypatch) + result = store.search_entries("hello", "app", "u1", embedding=[0.1]) + if async_mode: + result = await result + assert result == [] + query, params = cursor.execute.call_args.args + sql = query.decode() + assert "plainto_tsquery" in sql if use_fts else "content_text ILIKE" in sql + assert params == ("app", "u1", "hello" if use_fts else "%hello%", 20) + + empty = store.search_entries("", "app", "u1", embedding=[0.1]) + if async_mode: + empty = await empty + assert empty == [] + assert cursor.execute.call_count == 1 diff --git a/tests/unit/adapters/test_cockroach_psycopg/test_config.py b/tests/unit/adapters/test_cockroach_psycopg/test_config.py index 45d747d0c..3d7c9c825 100644 --- a/tests/unit/adapters/test_cockroach_psycopg/test_config.py +++ b/tests/unit/adapters/test_cockroach_psycopg/test_config.py @@ -8,13 +8,14 @@ - Connection config normalization """ -from typing import Any, cast -from unittest.mock import MagicMock, patch +from typing import Any, Self, cast +from unittest.mock import AsyncMock, MagicMock, patch import pytest from sqlspec.adapters.cockroach_psycopg import ( CockroachPsycopgAsyncConfig, + CockroachPsycopgAsyncDriver, CockroachPsycopgAsyncSessionContext, CockroachPsycopgDriverFeatures, CockroachPsycopgRetryConfig, @@ -810,3 +811,436 @@ async def __aexit__(self, *_args: Any) -> None: litestar_store = CockroachPsycopgSyncStore(good_cfg) litestar_store._table_name = "tenant_a.litestar_sessions" assert "idx_tenant_a_litestar_sessions_expires_at" in litestar_store._table_ddl() + + +def test_cockroach_psycopg_follower_read_only_for_readonly_statements(monkeypatch: pytest.MonkeyPatch) -> None: + """CockroachPsycopgSyncDriver should not begin follower-read transactions for modifying operations.""" + in_tx = False + + def _in_tx(_self: Any) -> bool: + return in_tx + + def _begin(_self: Any) -> None: + nonlocal in_tx + in_tx = True + + monkeypatch.setattr(CockroachPsycopgSyncDriver, "_connection_in_transaction", _in_tx) + monkeypatch.setattr(CockroachPsycopgSyncDriver, "begin", _begin) + + conn = MagicMock() + cursor = MagicMock() + col = MagicMock() + col.name = "id" + col.type_code = 23 + cursor.description = [col] + cursor.fetchall.return_value = [{"id": 1}] + driver = CockroachPsycopgSyncDriver( + connection=conn, driver_features={"enable_follower_reads": True, "default_staleness": "'-10s'"} + ) + modifying_stmt = SQL("INSERT INTO users (name) VALUES ('alice') RETURNING id") + driver.dispatch_execute(cursor, modifying_stmt) + assert in_tx is False + + +def test_cockroach_psycopg_begin_follower_read_failure_handling() -> None: + """Failing SET TRANSACTION in begin() must rollback, reset transaction active, and raise mapped exception.""" + conn = MagicMock() + conn.autocommit = True + conn.info.transaction_status = 0 + conn.execute.side_effect = RuntimeError("syntax error in staleness") + + driver = CockroachPsycopgSyncDriver( + connection=conn, driver_features={"enable_follower_reads": True, "default_staleness": "'-10s'"} + ) + with pytest.raises(Exception): + driver.begin() + assert driver._transaction_active is False + conn.rollback.assert_called_once() + + +def test_cockroach_psycopg_select_stream_follower_reads() -> None: + """select_stream on a SELECT query with follower reads emits the staleness clause and closes.""" + executed_statements: list[str] = [] + conn = MagicMock() + conn.autocommit = True + conn.info.transaction_status = 0 + + def fake_conn_execute(query: str, *args: Any, **kwargs: Any) -> Any: + executed_statements.append(query) + return MagicMock() + + conn.execute.side_effect = fake_conn_execute + + tx = MagicMock() + conn.transaction.return_value = tx + + stream_cursor = MagicMock() + col = MagicMock() + col.name = "id" + stream_cursor.description = [col] + stream_cursor.fetchmany.side_effect = [[(1,)], []] + conn.cursor.return_value = stream_cursor + + driver = CockroachPsycopgSyncDriver( + connection=conn, driver_features={"enable_follower_reads": True, "default_staleness": "'-10s'"} + ) + stream = driver.dispatch_select_stream(SQL("SELECT id FROM users"), chunk_size=10) + assert stream is not None + rows = list(stream) + assert rows == [{"id": 1}] + assert any("AS OF SYSTEM TIME '-10s'" in s for s in executed_statements) + conn.commit.assert_called() + + +def test_cockroach_psycopg_sync_user_configure_runs_after_internal_setup() -> None: + """User-supplied configure callback runs after internal connection setup in sync config.""" + call_order: list[str] = [] + fake_conn = MagicMock() + fake_conn.autocommit = False + + def user_cb(conn: Any) -> None: + assert conn.autocommit is True + call_order.append("user_cb") + + config = CockroachPsycopgSyncConfig(connection_config={"autocommit": True, "configure": user_cb}) + + with patch("sqlspec.adapters.cockroach_psycopg.config.psycopg_crdb.CrdbConnection.connect", return_value=fake_conn): + conn = config.create_connection() + assert conn is fake_conn + assert call_order == ["user_cb"] + + pool = config.create_pool() + assert pool is not None + + +async def test_cockroach_psycopg_async_user_configure_runs_after_internal_setup() -> None: + """User-supplied async configure callback runs after internal connection setup in async config.""" + call_order: list[str] = [] + fake_conn = MagicMock() + autocommit_val = False + + async def fake_set_autocommit(val: bool) -> None: + nonlocal autocommit_val + autocommit_val = val + call_order.append("internal_autocommit") + + fake_conn.set_autocommit = AsyncMock(side_effect=fake_set_autocommit) + + async def async_user_cb(conn: Any) -> None: + assert autocommit_val is True + call_order.append("user_cb") + + config = CockroachPsycopgAsyncConfig(connection_config={"autocommit": True, "configure": async_user_cb}) + + with patch( + "sqlspec.adapters.cockroach_psycopg.config.psycopg_crdb.AsyncCrdbConnection.connect", + AsyncMock(return_value=fake_conn), + ): + conn = await config.create_connection() + assert conn is fake_conn + assert call_order == ["internal_autocommit", "user_cb"] + + +@pytest.mark.parametrize("async_mode", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize( + ("override_kind", "expected"), + [ + ("inherit", ["original"]), + ("mapping", ["replacement"]), + ("mapping_none", []), + ("keyword", ["replacement"]), + ("keyword_none", []), + ("keyword_precedence", []), + ], +) +async def test_cockroach_psycopg_derived_connection_preserves_configure( + async_mode: bool, override_kind: str, expected: list[str] +) -> None: + """Derived connections inherit configure callbacks and honor explicit replacements.""" + calls: list[str] = [] + + def original_callback(_connection: Any) -> None: + calls.append("original") + + def replacement_callback(_connection: Any) -> None: + calls.append("replacement") + + async def async_original_callback(connection: Any) -> None: + original_callback(connection) + + async def async_replacement_callback(connection: Any) -> None: + replacement_callback(connection) + + original = async_original_callback if async_mode else original_callback + replacement = async_replacement_callback if async_mode else replacement_callback + source_config: dict[str, Any] = {"autocommit": True, "configure": original} + config = ( + CockroachPsycopgAsyncConfig(connection_config=source_config) + if async_mode + else CockroachPsycopgSyncConfig(connection_config=source_config) + ) + overrides: dict[str, Any] = {} + if override_kind in {"mapping", "keyword_precedence"}: + overrides["connection_config_overrides"] = {"configure": replacement} + elif override_kind == "mapping_none": + overrides["connection_config_overrides"] = {"configure": None} + if override_kind == "keyword": + overrides["configure"] = replacement + elif override_kind in {"keyword_none", "keyword_precedence"}: + overrides["configure"] = None + derived = config.derive_pool(**overrides) + fake_connection = MagicMock() + fake_connection.set_autocommit = AsyncMock() + sdk_connection = "AsyncCrdbConnection" if async_mode else "CrdbConnection" + connect = AsyncMock(return_value=fake_connection) if async_mode else MagicMock(return_value=fake_connection) + + with patch(f"sqlspec.adapters.cockroach_psycopg.config.psycopg_crdb.{sdk_connection}.connect", connect): + if isinstance(derived, CockroachPsycopgAsyncConfig): + await derived.create_connection() + else: + derived.create_connection() + assert calls == expected + calls.clear() + with ( + patch("sqlspec.adapters.cockroach_psycopg.config.ConnectionPool", _SyncPoolSpy), + patch("sqlspec.adapters.cockroach_psycopg.config.AsyncConnectionPool", _AsyncPoolSpy), + ): + if isinstance(derived, CockroachPsycopgAsyncConfig): + await derived.create_pool() + pool_kwargs = _AsyncPoolSpy.calls[-1][1] + await pool_kwargs["configure"](fake_connection) + else: + derived.create_pool() + pool_kwargs = _SyncPoolSpy.calls[-1][1] + pool_kwargs["configure"](fake_connection) + assert pool_kwargs["kwargs"] == {"autocommit": True} + assert calls == expected + calls.clear() + if isinstance(config, CockroachPsycopgAsyncConfig): + await config.create_connection() + else: + config.create_connection() + + assert calls == ["original"] + assert source_config["configure"] is original + + +def test_cockroach_psycopg_adk_hash_shard_bucket_count_validation() -> None: + """Invalid hash_shard_bucket_count must raise ValueError in cockroach_psycopg adk store.""" + from sqlspec.adapters.cockroach_psycopg.adk.store import CockroachPsycopgADKConfig, _cockroach_hash_shard_clause + + bad_config: CockroachPsycopgADKConfig = {"enable_hash_sharded_indexes": True, "hash_shard_bucket_count": -1} + with pytest.raises(ValueError, match="hash_shard_bucket_count"): + _cockroach_hash_shard_clause(bad_config) + + +async def test_cockroach_psycopg_adk_memory_search_with_embedding() -> None: + """An embedding argument preserves text search against the canonical memory schema.""" + executed_queries: list[tuple[str, Any]] = [] + + class _FakeCursor: + description = [ + ("id",), + ("session_id",), + ("app_name",), + ("user_id",), + ("scope",), + ("event_id",), + ("author",), + ("timestamp",), + ("content_json",), + ("content_text",), + ("metadata_json",), + ("inserted_at",), + ] + + async def execute(self, sql: bytes, params: Any) -> None: + executed_queries.append((sql.decode(), params)) + + async def fetchall(self) -> list[Any]: + return [] + + async def __aenter__(self) -> Self: + return self + + async def __aexit__(self, *_a: Any) -> None: + pass + + class _FakeConn: + def cursor(self) -> _FakeCursor: + return _FakeCursor() + + class _FakeCtx: + async def __aenter__(self) -> _FakeConn: + return _FakeConn() + + async def __aexit__(self, *_a: Any) -> None: + pass + + config = CockroachPsycopgAsyncConfig(extension_config={"adk": {"enable_memory": True}}) + object.__setattr__(config, "provide_connection", lambda: _FakeCtx()) + store = CockroachPsycopgAsyncADKMemoryStore(config) + + result = await store.search_entries(query="hello", app_name="app", user_id="u1", limit=7, embedding=[0.1, 0.2, 0.3]) + assert result == [] + assert len(executed_queries) == 1 + sql, params = executed_queries[0] + assert "content_text ILIKE %s" in sql + assert "ORDER BY timestamp DESC" in sql + assert params == ("app", "u1", "%hello%", 7) + + +async def test_cockroach_psycopg_adk_memory_insert_batch_unnest() -> None: + """insert_memory_entries executes a single UNNEST query rather than an N+1 loop.""" + executed_queries: list[str] = [] + + class _FakeCursor: + rowcount = 2 + + async def execute(self, sql: Any, params: Any) -> None: + executed_queries.append(str(sql)) + + async def __aenter__(self) -> Self: + return self + + async def __aexit__(self, *_a: Any) -> None: + pass + + class _FakeConn: + def cursor(self) -> _FakeCursor: + return _FakeCursor() + + async def commit(self) -> None: + pass + + class _FakeCtx: + async def __aenter__(self) -> _FakeConn: + return _FakeConn() + + async def __aexit__(self, *_a: Any) -> None: + pass + + config = CockroachPsycopgAsyncConfig(extension_config={"adk": {"enable_memory": True}}) + object.__setattr__(config, "provide_connection", lambda: _FakeCtx()) + store = CockroachPsycopgAsyncADKMemoryStore(config) + + entry1: Any = { + "id": "1", + "session_id": "s1", + "app_name": "app", + "user_id": "u1", + "scope": "user", + "event_id": "e1", + "author": "user", + "timestamp": "2026-01-01T00:00:00Z", + "content_json": {}, + "content_text": "hello", + "metadata_json": None, + "inserted_at": "2026-01-01T00:00:00Z", + } + entry2: Any = { + "id": "2", + "session_id": "s1", + "app_name": "app", + "user_id": "u1", + "scope": "user", + "event_id": "e2", + "author": "user", + "timestamp": "2026-01-01T00:00:00Z", + "content_json": {}, + "content_text": "world", + "metadata_json": None, + "inserted_at": "2026-01-01T00:00:00Z", + } + + count = await store.insert_memory_entries([entry1, entry2]) + assert count == 2 + assert len(executed_queries) == 1 + assert "unnest" in executed_queries[0].lower() + + +def test_cockroach_psycopg_sync_get_columns_and_indexes_by_table() -> None: + """CockroachPsycopgSyncDataDictionary executes by_table query when table is specified.""" + from sqlspec.adapters.cockroach_psycopg.data_dictionary import CockroachPsycopgSyncDataDictionary + + calls: list[tuple[Any, dict[str, Any]]] = [] + driver = MagicMock() + + def fake_select(query: Any, **kwargs: Any) -> list[Any]: + calls.append((query, kwargs)) + return [] + + driver.select = fake_select + dd = CockroachPsycopgSyncDataDictionary() + + dd.get_columns(driver, table="users", schema=None) + assert len(calls) == 1 + assert calls[0][1]["table_name"] == "users" + assert calls[0][1]["schema_name"] is None + + calls.clear() + dd.get_indexes(driver, table="users", schema=None) + assert len(calls) == 1 + assert calls[0][1]["table_name"] == "users" + assert calls[0][1]["schema_name"] is None + + +async def test_cockroach_psycopg_async_get_columns_and_indexes_by_table() -> None: + """CockroachPsycopgAsyncDataDictionary executes by_table query when table is specified.""" + from sqlspec.adapters.cockroach_psycopg.data_dictionary import CockroachPsycopgAsyncDataDictionary + + calls: list[tuple[Any, dict[str, Any]]] = [] + driver = MagicMock() + + async def fake_select(query: Any, **kwargs: Any) -> list[Any]: + calls.append((query, kwargs)) + return [] + + driver.select = AsyncMock(side_effect=fake_select) + dd = CockroachPsycopgAsyncDataDictionary() + + await dd.get_columns(driver, table="users", schema=None) + assert len(calls) == 1 + assert calls[0][1]["table_name"] == "users" + assert calls[0][1]["schema_name"] is None + + calls.clear() + await dd.get_indexes(driver, table="users", schema=None) + assert len(calls) == 1 + assert calls[0][1]["table_name"] == "users" + assert calls[0][1]["schema_name"] is None + + +def test_cockroach_psycopg_sync_execute_native_storage_under_active_exception() -> None: + """_execute_native_storage in sync driver inside active except block works cleanly.""" + conn = MagicMock() + conn.autocommit = True + cursor = MagicMock() + cursor.fetchall.return_value = [{"rows": 1}] + conn.cursor.return_value = cursor + + driver = CockroachPsycopgSyncDriver(connection=conn) + try: + raise ValueError("outer error") + except ValueError: + rows = driver._execute_native_storage("SELECT 1", []) + assert rows == [{"rows": 1}] + + +async def test_cockroach_psycopg_async_execute_native_storage_under_active_exception() -> None: + """_execute_native_storage in async driver inside active except block works cleanly.""" + conn = MagicMock() + conn.autocommit = True + cursor = MagicMock() + cursor.fetchall = AsyncMock(return_value=[{"rows": 1}]) + cursor.execute = AsyncMock() + cursor.__aenter__ = AsyncMock(return_value=cursor) + cursor.__aexit__ = AsyncMock() + conn.cursor.return_value = cursor + + driver = CockroachPsycopgAsyncDriver(connection=conn) + try: + raise ValueError("outer error") + except ValueError: + rows = await driver._execute_native_storage("SELECT 1", []) + assert rows == [{"rows": 1}] diff --git a/tests/unit/adapters/test_db2/_fakes.py b/tests/unit/adapters/test_db2/_fakes.py index 7607bf1f8..de6b5a3e7 100644 --- a/tests/unit/adapters/test_db2/_fakes.py +++ b/tests/unit/adapters/test_db2/_fakes.py @@ -212,6 +212,7 @@ def __init__( self.dbms_name = "DB2/LINUXX8664" self.autocommit = autocommit self.dsn: str | None = None + self.current_schema: str | None = None self.pending: list[tuple[str, object]] = [] self.committed: list[tuple[str, object]] = [] self.rollbacks = 0 @@ -241,6 +242,11 @@ def set_autocommit(self, is_on: bool) -> None: """Switch the connection's autocommit mode.""" self.autocommit = bool(is_on) + def set_current_schema(self, schema: str) -> bool: + """Update the schema used by SDK catalog operations.""" + self.current_schema = schema + return True + def commit(self) -> None: """Make pending writes durable.""" self.commits += 1 @@ -380,6 +386,10 @@ async def set_autocommit(self, is_on: bool) -> None: """Switch the wrapped connection's autocommit mode.""" self.sync_connection.set_autocommit(is_on) + async def set_current_schema(self, schema: str) -> bool: + """Update the schema used by SDK catalog operations.""" + return self.sync_connection.set_current_schema(schema) + class FakeAsyncConnectionFactory: """Stand-in for the ``ibm_db_dbi.AsyncConnection`` class and its ``connect`` factory.""" @@ -475,6 +485,7 @@ def __init__( self.connection = FakeDb2Connection(list(cursors), autocommit=True) self.extension_config: dict[str, Any] = extension_config or {} self.statement_config = default_statement_config + self.driver_features: dict[str, Any] = {} @property def executed(self) -> "list[tuple[str, Any]]": @@ -488,7 +499,9 @@ class FakeDb2SessionConfig(_FakeDb2SessionState): @contextmanager def provide_session(self, **_: Any) -> "Iterator[Db2SyncDriver]": """Yield a driver bound to the shared fake connection.""" - yield Db2SyncDriver(self.connection, statement_config=default_statement_config) + yield Db2SyncDriver( + self.connection, statement_config=default_statement_config, driver_features=self.driver_features + ) class FakeDb2AsyncSessionConfig(_FakeDb2SessionState): @@ -502,7 +515,9 @@ class FakeDb2AsyncSessionConfig(_FakeDb2SessionState): async def provide_session(self, **_: Any) -> "AsyncIterator[Any]": """Yield an async driver bound to the shared fake connection.""" yield db2_driver.Db2AsyncDriver( - FakeDb2AsyncConnection(self.connection), statement_config=default_statement_config + FakeDb2AsyncConnection(self.connection), + statement_config=default_statement_config, + driver_features=self.driver_features, ) diff --git a/tests/unit/adapters/test_db2/test_adk.py b/tests/unit/adapters/test_db2/test_adk.py index 872989e2b..ede427fc7 100644 --- a/tests/unit/adapters/test_db2/test_adk.py +++ b/tests/unit/adapters/test_db2/test_adk.py @@ -9,7 +9,7 @@ import pytest from sqlspec.adapters.db2.adk import Db2AsyncADKMemoryStore, Db2AsyncADKStore, Db2SyncADKMemoryStore, Db2SyncADKStore -from sqlspec.adapters.db2.adk.store import _datetime_value, _is_table_missing +from sqlspec.adapters.db2.adk.store import _create_index_sql, _datetime_value, _is_table_missing from sqlspec.adapters.db2.config import Db2AsyncConfig, Db2SyncConfig from sqlspec.adapters.db2.core import INDEX_EXISTS_SQL, TABLE_EXISTS_SQL, create_mapped_exception from sqlspec.exceptions import SQLParsingError, SQLSpecError @@ -421,7 +421,7 @@ async def test_insert_memory_entries_binds_naive_utc_and_skips_duplicates(db2_mo """Memory inserts bind naive-UTC times, include the owner column, and skip known event ids.""" store, config = _memory_store( db2_mode, - FakeDb2Cursor(rows=[("m0",)], description=db2_description("id")), + FakeDb2Cursor(rows=[("m0", "e0")], description=db2_description("id", "event_id")), adk={"owner_id_column": "tenant_id INTEGER"}, ) stamp = datetime(2026, 9, 24, 14, tzinfo=timezone(timedelta(hours=2))) @@ -445,8 +445,9 @@ async def test_insert_memory_entries_binds_naive_utc_and_skips_duplicates(db2_mo ] assert await db2_mode.call(store.insert_memory_entries, entries, owner_id=7) == 1 + assert len(config.executed) <= 2 - insert_sql, insert_params = config.executed[2] + insert_sql, insert_params = config.executed[1] assert f"INSERT INTO {store.memory_table} (" in insert_sql assert "inserted_at, tenant_id" in insert_sql assert insert_params[7] == datetime(2026, 9, 24, 12) @@ -642,4 +643,60 @@ async def test_schema_qualified_adk_tables_sanitize_constraint_and_index_names(d assert "CONSTRAINT pk_app_adk_memory_id" in memory_ddl assert "CONSTRAINT uq_app_adk_memory_event" in memory_ddl assert all("." not in index_name for index_name, _, _ in store._index_specs()) - assert all("." not in index_name for index_name, _, _ in mem_store._memory_index_specs()) + + +@pytest.mark.parametrize( + ("table", "schema"), [("MYSCHEMA.SESSIONS", "MYSCHEMA"), ('"App".sessions', "App"), ('"A""B".sessions', 'A"B')] +) +async def test_adk_index_catalog_probes_preserve_schema(db2_mode: DriverMode, table: str, schema: str) -> None: + """Index existence checks use the schema and index casing emitted by DDL.""" + store, config = _session_store(db2_mode) + store._session_table = table + + await db2_mode.call(store.create_tables) + + index_probes = [params for sql, params in config.executed if sql == INDEX_EXISTS_SQL] + assert index_probes[:2] == [(schema, name.upper()) for name, _, _ in store._index_specs()[:2]] + + +@pytest.mark.parametrize("lowercase_columns", [True, False]) +async def test_memory_replay_deduplicates_event_ids_with_native_casing( + db2_mode: DriverMode, lowercase_columns: bool +) -> None: + """A replayed event is skipped even when the returned memory ID differs.""" + store, config = _memory_store( + db2_mode, FakeDb2Cursor(rows=[("old-memory", "same-event")], description=db2_description("id", "event_id")) + ) + config.driver_features["enable_lowercase_column_names"] = lowercase_columns + stamp = datetime.now(UTC) + entry = StoredMemory( + id="new-memory", + session_id="s1", + app_name="app", + user_id="u1", + scope="user", + event_id="same-event", + author=None, + timestamp=stamp, + content_json={}, + content_text="text", + metadata_json=None, + inserted_at=stamp, + embedding=None, + ) + + assert await db2_mode.call(store.insert_memory_entries, [entry]) == 0 + assert len(config.executed) == 1 + + +def test_create_index_sql_qualifies_schema() -> None: + """ADK index DDL qualifies index name when table is schema-qualified.""" + assert ( + _create_index_sql('"MYSCHEMA"."SESSIONS"', "ix_sessions", "user_id") + == 'CREATE INDEX "MYSCHEMA"."IX_SESSIONS" ON "MYSCHEMA"."SESSIONS" (user_id)' + ) + assert ( + _create_index_sql("MYSCHEMA.SESSIONS", "ix_sessions", "user_id") + == 'CREATE INDEX "MYSCHEMA"."IX_SESSIONS" ON MYSCHEMA.SESSIONS (user_id)' + ) + assert _create_index_sql("sessions", "ix_sessions", "user_id") == "CREATE INDEX ix_sessions ON sessions (user_id)" diff --git a/tests/unit/adapters/test_db2/test_async_pool.py b/tests/unit/adapters/test_db2/test_async_pool.py index 4cdde9b9d..6a99e76de 100644 --- a/tests/unit/adapters/test_db2/test_async_pool.py +++ b/tests/unit/adapters/test_db2/test_async_pool.py @@ -7,11 +7,11 @@ import pytest -import sqlspec.adapters.db2.pool as pool_module from sqlspec.adapters.db2.pool import Db2AsyncConnectionPool from sqlspec.exceptions import ConnectionTimeoutError, DatabaseConnectionError from tests.unit.adapters.test_db2._fakes import ( FakeDb2AsyncConnection, + FakeDb2AsyncCursor, FakeDb2Connection, FakeDb2Cursor, FakeDb2OperationalError, @@ -26,7 +26,7 @@ class FakeClock: - """Replacement for the pool module's ``time`` with a controllable monotonic clock.""" + """Controllable external monotonic clock.""" def __init__(self) -> None: self.now = 1000.0 @@ -34,9 +34,6 @@ def __init__(self) -> None: def monotonic(self) -> float: return self.now - def time(self) -> float: - return time.time() - def advance(self, seconds: float) -> None: self.now += seconds @@ -205,7 +202,7 @@ async def test_recycle_and_health_check(fake_ibm_db: FakeModules, monkeypatch: p """Old connections are recycled, idle ones are pinged, and failed pings replace the connection.""" _, module = fake_ibm_db clock = FakeClock() - monkeypatch.setattr(pool_module, "time", clock) + monkeypatch.setattr(time, "monotonic", clock.monotonic) recycled = FakeDb2Connection() unhealthy = FakeDb2Connection([ FakeDb2Cursor(error=db2_error(-30108, "08506", "A connection failed", FakeDb2OperationalError)) @@ -327,30 +324,47 @@ async def test_async_connection_create_hook_is_awaited( await pool.close() +@pytest.mark.parametrize("close_fails", [False, True]) async def test_cancelled_health_check_closes_detached_connection( - fake_ibm_db: FakeModules, monkeypatch: pytest.MonkeyPatch + fake_ibm_db: FakeModules, close_fails: bool, monkeypatch: pytest.MonkeyPatch ) -> None: - pool = Db2AsyncConnectionPool({"database": "TESTDB"}, max_size=1) + pool = Db2AsyncConnectionPool({"database": "TESTDB"}, max_size=1, health_check_interval=-1.0) connection = _as_fake(await pool.acquire()) await pool.release(connection) checking = asyncio.Event() + finish = asyncio.Event() + execute = FakeDb2AsyncCursor.execute + close = FakeDb2AsyncCursor.close - async def check(_pool: object, _record: object) -> bool: + async def blocked_execute(cursor: FakeDb2AsyncCursor, operation: str, parameters: object = None) -> FakeDb2Cursor: checking.set() - await asyncio.Event().wait() - return True + await finish.wait() + return await execute(cursor, operation, parameters) - monkeypatch.setattr(Db2AsyncConnectionPool, "_is_reusable", check) + async def failing_close(cursor: FakeDb2AsyncCursor) -> None: + await close(cursor) + msg = "Cursor close failed during cancelled health check." + raise OSError(msg) + + monkeypatch.setattr(FakeDb2AsyncCursor, "execute", blocked_execute) + if close_fails: + monkeypatch.setattr(FakeDb2AsyncCursor, "close", failing_close) task = asyncio.create_task(pool.acquire()) - await checking.wait() - task.cancel() - with pytest.raises(asyncio.CancelledError): - await task - assert connection.sync_connection.closed is True - assert pool.size() == 0 - replacement = await pool.acquire() - await pool.release(replacement) - await pool.close() + try: + await asyncio.wait_for(checking.wait(), timeout=1.0) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout=1.0) + assert connection.sync_connection.closed is True + assert connection.sync_connection.cursors[0].closed is True + assert pool.size() == 0 + replacement = await asyncio.wait_for(pool.acquire(), timeout=1.0) + await pool.release(replacement) + finally: + finish.set() + task.cancel() + await asyncio.gather(task, return_exceptions=True) + await pool.close() async def test_pool_closed_while_opening_rejects_and_closes_connection(fake_ibm_db: FakeModules) -> None: diff --git a/tests/unit/adapters/test_db2/test_config.py b/tests/unit/adapters/test_db2/test_config.py index a0dba3b8d..c7446282e 100644 --- a/tests/unit/adapters/test_db2/test_config.py +++ b/tests/unit/adapters/test_db2/test_config.py @@ -1,5 +1,6 @@ """Tests for IBM Db2 database configuration.""" +import asyncio from datetime import datetime, timedelta, timezone from typing import Any, cast @@ -25,6 +26,7 @@ DriverMode, FakeDb2AsyncConnection, FakeDb2Connection, + FakeDb2Cursor, FakeIbmDbDbiModule, FakeIbmDbModule, ) @@ -135,6 +137,104 @@ def test_password_with_semicolon_is_brace_escaped(fake_ibm_db: FakeModules) -> N assert (user, password, host, database) == ("", "", "", "") +@pytest.mark.anyio +@pytest.mark.parametrize("schema", ["APP", "SQLSPEC;SEMICOLON", "Mixed Case"]) +async def test_schema_dsn_roundtrip_is_applied_before_creation_hook( + db2_mode: DriverMode, fake_ibm_db: FakeModules, schema: str +) -> None: + """Schema values reach SQL as one bound value instead of CLI keyword fragments.""" + _, fake_module = fake_ibm_db + cursor = FakeDb2Cursor() + native_connection = FakeDb2Connection(cursors=[cursor]) + fake_module.pending_connections.append(native_connection) + hook_observations: list[tuple[list[tuple[str, Any]], str | None]] = [] + source = db2_mode.config( + connection_config={"database": "d", "user": "u", "password": "p;x", "current_schema": schema} + ) + config = db2_mode.config( + connection_config={"dsn": source.get_connection_string()}, + driver_features={ + "on_connection_create": lambda _: hook_observations.append(( + list(cursor.executed), + native_connection.current_schema, + )) + }, + ) + expected = [("SET CURRENT SCHEMA = ?", (schema,))] + + assert config.get_connection_string() == source.get_connection_string() + connection = await config.create_connection() if db2_mode.is_async else config.create_connection() + try: + assert cursor.executed == expected + assert hook_observations == [(expected, schema)] + assert cursor.closed + assert not native_connection.closed + assert fake_module.connect_calls[0][:5] == ("DATABASE=d;UID=u;PWD={p;x};", "", "", "", "") + finally: + if db2_mode.is_async: + await connection.close() + await config.close_pool() + else: + connection.close() + config.close_pool() + + +@pytest.mark.anyio +@pytest.mark.parametrize("failure", ["execute", "close", "setter", "execute-and-close", "cancel-and-close"]) +async def test_schema_initialization_failure_closes_connection_and_preserves_error( + db2_mode: DriverMode, fake_ibm_db: FakeModules, monkeypatch: pytest.MonkeyPatch, failure: str +) -> None: + """Schema initialization never hands an uninitialized or cancelled connection to the hook.""" + _, fake_module = fake_ibm_db + primary: BaseException = ( + asyncio.CancelledError("schema cancelled") + if failure == "cancel-and-close" + else RuntimeError("schema initialization failed") + ) + cursor = FakeDb2Cursor(error=primary if failure in {"execute", "execute-and-close", "cancel-and-close"} else None) + native_connection = FakeDb2Connection(cursors=[cursor]) + fake_module.pending_connections.append(native_connection) + hook_calls: list[object] = [] + + def close_cursor() -> None: + cursor.closed = True + raise primary if failure == "close" else RuntimeError("schema cursor close failed") + + def set_schema(_: str) -> bool: + raise primary + + original_close = native_connection.close + + def close_connection() -> None: + original_close() + raise RuntimeError("connection close failed") + + if failure in {"close", "execute-and-close", "cancel-and-close"}: + monkeypatch.setattr(cursor, "close", close_cursor) + if failure == "setter": + monkeypatch.setattr(native_connection, "set_current_schema", set_schema) + monkeypatch.setattr(native_connection, "close", close_connection) + config = db2_mode.config( + connection_config={"database": "d", "current_schema": "SQLSPEC;SEMICOLON"}, + driver_features={"on_connection_create": hook_calls.append}, + ) + try: + with pytest.raises(type(primary)) as caught: + if db2_mode.is_async: + await config.create_connection() + else: + config.create_connection() + assert caught.value is primary + assert cursor.closed + assert native_connection.closed + assert hook_calls == [] + finally: + if db2_mode.is_async: + await config.close_pool() + else: + config.close_pool() + + @pytest.mark.parametrize( ("value", "rendered"), [("a{b", "{a{b}"), (" padded", "{ padded}"), (True, "1"), (False, "0"), ("plain", "plain")], diff --git a/tests/unit/adapters/test_db2/test_data_dictionary.py b/tests/unit/adapters/test_db2/test_data_dictionary.py index 8c003f9d5..39a190a3c 100644 --- a/tests/unit/adapters/test_db2/test_data_dictionary.py +++ b/tests/unit/adapters/test_db2/test_data_dictionary.py @@ -493,3 +493,11 @@ async def test_get_columns_reports_unique_separately_from_primary(case: Dictiona assert columns[0]["is_primary"] is False assert columns[0]["is_unique"] is True + + +def test_all_authored_db2_dialect_queries_are_wired() -> None: + """Every authored query in Db2 SQL template files is loaded and callable by the data dictionary.""" + dictionary = Db2SyncDataDictionary() + assert dictionary.get_query("tables", "by_schema") is not None + assert dictionary.get_query("tables", "all_by_schema") is not None + assert dictionary.get_query("schemas", "by_schema") is not None diff --git a/tests/unit/adapters/test_db2/test_driver.py b/tests/unit/adapters/test_db2/test_driver.py index e5e3d5bf1..c9dc0e6f3 100644 --- a/tests/unit/adapters/test_db2/test_driver.py +++ b/tests/unit/adapters/test_db2/test_driver.py @@ -3,10 +3,12 @@ Every test runs in both driver modes through the ``db2_mode`` fixture. """ +import importlib from typing import Any import pytest +import sqlspec.adapters.db2._typing as db2_typing from sqlspec.adapters.db2.core import create_mapped_exception, default_statement_config, resolve_column_names from sqlspec.adapters.db2.data_dictionary import Db2AsyncDataDictionary, Db2SyncDataDictionary from sqlspec.core import SQL @@ -311,3 +313,43 @@ async def test_driver_exposes_matching_data_dictionary(db2_mode: DriverMode) -> assert isinstance(dictionary, Db2AsyncDataDictionary if db2_mode.is_async else Db2SyncDataDictionary) assert driver.data_dictionary is dictionary + + +async def test_select_stream_inside_active_exception_block(db2_mode: DriverMode) -> None: + """Consuming select_stream inside an active except block works cleanly without masking exceptions.""" + cursor = FakeDb2Cursor(rows=[(1, "Ada"), (2, "Grace")], description=[("id",), ("name",)]) + driver = db2_mode.driver(FakeDb2Connection(lambda: cursor)) + + try: + raise ValueError("ambient error") + except ValueError: + rows = await db2_mode.collect(driver.select_stream("SELECT id, name FROM users", chunk_size=1)) + assert len(rows) == 2 + + +def test_session_driver_type_resolution_is_cached(monkeypatch: pytest.MonkeyPatch) -> None: + """Db2SyncSessionContext caches driver class resolution and does not re-import module.""" + import_calls: list[str] = [] + original_import = importlib.import_module + + def _spy_import(name: str, *args: Any, **kwargs: Any) -> Any: + import_calls.append(name) + return original_import(name, *args, **kwargs) + + monkeypatch.setattr(importlib, "import_module", _spy_import) + db2_typing._SYNC_DRIVER_CLS = None + + context = db2_typing.Db2SyncSessionContext( + acquire_connection=FakeDb2Connection, + release_connection=lambda *args, **kwargs: None, + statement_config=default_statement_config, + driver_features={}, + prepare_driver=lambda d: d, + ) + with context: + pass + with context: + pass + + db2_driver_imports = [c for c in import_calls if c == "sqlspec.adapters.db2.driver"] + assert len(db2_driver_imports) == 1 diff --git a/tests/unit/adapters/test_db2/test_exceptions.py b/tests/unit/adapters/test_db2/test_exceptions.py index a31b350fe..ea5a10205 100644 --- a/tests/unit/adapters/test_db2/test_exceptions.py +++ b/tests/unit/adapters/test_db2/test_exceptions.py @@ -52,6 +52,9 @@ ("22003", DataError), ("22007", DataError), ("22012", DataError), + ("42703", SQLParsingError), + ("42710", SQLParsingError), + ("42884", SQLParsingError), ], ) def test_sqlstate_attribute_mapping(sqlstate: str, expected_type: type[SQLSpecError]) -> None: @@ -94,6 +97,11 @@ def test_sqlstate_embedded_in_message(sqlstate_str: str, expected_type: type[SQL ("SQL30081N", DatabaseConnectionError), ("SQL0104N", SQLParsingError), ("SQL0204N", SQLParsingError), + ("SQL0206N", SQLParsingError), + ("SQL0440N", SQLParsingError), + ("SQL0601N", SQLParsingError), + ("SQL30082N", PermissionDeniedError), + ("30082", PermissionDeniedError), ("-803", UniqueViolationError), ], ) diff --git a/tests/unit/adapters/test_db2/test_litestar.py b/tests/unit/adapters/test_db2/test_litestar.py index 3e05a8e84..7f4de1a16 100644 --- a/tests/unit/adapters/test_db2/test_litestar.py +++ b/tests/unit/adapters/test_db2/test_litestar.py @@ -15,6 +15,7 @@ from sqlspec.adapters.db2.config import Db2AsyncConfig, Db2SyncConfig from sqlspec.adapters.db2.core import TABLE_EXISTS_SQL, to_db_timestamp, utc_now from sqlspec.adapters.db2.litestar import Db2AsyncStore, Db2SyncStore +from sqlspec.adapters.db2.litestar.store import session_index_ddl from tests.unit.adapters.test_db2._fakes import DriverMode, FakeDb2Cursor, db2_description pytestmark = pytest.mark.anyio @@ -268,3 +269,16 @@ def test_store_resolves_by_config_name(config_class: Any, store_class: Any) -> N context = SimpleNamespace(config=config_class(connection_config={"database": "d"})) assert migration._get_store_class(context) is store_class + + +def test_session_index_ddl_qualifies_schema() -> None: + """Session index DDL qualifies index name when table is schema-qualified.""" + assert ( + session_index_ddl('"MYSCHEMA"."SESSIONS"') + == 'CREATE INDEX "MYSCHEMA"."IX_SESSIONS_EXP" ON "MYSCHEMA"."SESSIONS"(expires_at)' + ) + assert ( + session_index_ddl("MYSCHEMA.SESSIONS") + == 'CREATE INDEX "MYSCHEMA"."IX_SESSIONS_EXP" ON MYSCHEMA.SESSIONS(expires_at)' + ) + assert session_index_ddl("sessions") == "CREATE INDEX IX_SESSIONS_EXP ON sessions(expires_at)" diff --git a/tests/unit/adapters/test_db2/test_migrations.py b/tests/unit/adapters/test_db2/test_migrations.py index 1d8f17907..67a158cd3 100644 --- a/tests/unit/adapters/test_db2/test_migrations.py +++ b/tests/unit/adapters/test_db2/test_migrations.py @@ -11,9 +11,9 @@ from sqlspec.adapters.db2.config import Db2AsyncConfig, Db2SyncConfig from sqlspec.adapters.db2.core import TABLE_EXISTS_SQL, build_set_schema_sql, split_db2_table_name from sqlspec.adapters.db2.migrations import Db2AsyncMigrationTracker, Db2SyncMigrationTracker -from sqlspec.exceptions import ImproperConfigurationError +from sqlspec.exceptions import ImproperConfigurationError, PermissionDeniedError from sqlspec.observability import resolve_db_system -from tests.unit.adapters.test_db2._fakes import DriverMode, FakeDb2Connection, FakeDb2Cursor +from tests.unit.adapters.test_db2._fakes import DriverMode, FakeDb2Connection, FakeDb2Cursor, db2_error pytestmark = pytest.mark.anyio @@ -278,6 +278,19 @@ async def test_has_schema(db2_mode: DriverMode, rows: "list[Any]", expected: boo assert connection.cursors[0].executed == [("SELECT 1 FROM SYSCAT.SCHEMATA WHERE SCHEMANAME = ?", ("APP",))] +async def test_migration_schema_operations_translate_exceptions(db2_mode: DriverMode) -> None: + """Migration schema operations map native driver errors to SQLSpec exceptions.""" + error = db2_error(sqlcode=-551, sqlstate="42501", text="permission denied") + connection = FakeDb2Connection(cursors=lambda: FakeDb2Cursor(error=error)) + driver = db2_mode.driver(connection) + + with pytest.raises(PermissionDeniedError): + await db2_mode.call(driver.has_schema, "app") + + with pytest.raises(PermissionDeniedError): + await db2_mode.call(driver.set_migration_session_schema, "app") + + def test_configs_use_db2_trackers() -> None: """Each Db2 config tracks migrations with the Db2 tracker of its mode.""" assert Db2SyncConfig.migration_tracker_type is Db2SyncMigrationTracker diff --git a/tests/unit/adapters/test_db2/test_pool.py b/tests/unit/adapters/test_db2/test_pool.py index 573155bb7..e3c4d048c 100644 --- a/tests/unit/adapters/test_db2/test_pool.py +++ b/tests/unit/adapters/test_db2/test_pool.py @@ -6,9 +6,7 @@ import pytest -import sqlspec.adapters.db2.pool as pool_module from sqlspec.adapters.db2.pool import Db2SyncConnectionPool -from sqlspec.exceptions import MissingDependencyError from tests.unit.adapters.test_db2._fakes import FakeDb2Connection, FakeDb2Cursor, FakeIbmDbDbiModule, FakeIbmDbModule FakeModules = tuple[FakeIbmDbModule, FakeIbmDbDbiModule] @@ -38,7 +36,10 @@ def test_pool_connects_with_config_and_runs_hook(fake_ibm_db: FakeModules) -> No assert "PORT=50000" in dsn_called assert seen == [connection] assert pool.size() == 1 - assert pool.checked_out() == 0 + assert pool.checked_out() == 1 + + with pool.get_connection(): + assert pool.checked_out() == 1 def test_pool_renders_dsn_and_passes_empty_connect_arguments(fake_ibm_db: FakeModules) -> None: @@ -64,32 +65,19 @@ def test_pool_renders_dsn_and_passes_empty_connect_arguments(fake_ibm_db: FakeMo ) -def test_pool_missing_dependency_raises(monkeypatch: pytest.MonkeyPatch) -> None: - """MissingDependencyError is raised when ibm_db_dbi is not available.""" - monkeypatch.setattr(pool_module, "_IBM_DB_DBI", None) - monkeypatch.setattr(pool_module, "import_optional", lambda _name: None) - - pool = Db2SyncConnectionPool({"database": "TESTDB"}) - - with pytest.raises(MissingDependencyError) as exc_info: - pool.acquire() - - assert "ibm_db" in str(exc_info.value) - - -def test_pool_recycles_failed_health_check(fake_ibm_db: FakeModules, monkeypatch: pytest.MonkeyPatch) -> None: +def test_pool_recycles_failed_health_check(fake_ibm_db: FakeModules) -> None: """Failed health check should close and replace thread connection.""" _, fake_module = fake_ibm_db - first = FakeDb2Connection() + first = FakeDb2Connection([FakeDb2Cursor(error=RuntimeError("connection lost"))]) second = FakeDb2Connection() fake_module.pending_connections.extend([first, second]) - monkeypatch.setattr(Db2SyncConnectionPool, "_is_connection_alive", lambda *_: False) pool = Db2SyncConnectionPool({"database": "TESTDB"}, health_check_interval=-1.0) assert cast(object, pool.acquire()) is first assert cast(object, pool.acquire()) is second assert first.closed is True + assert first.cursors[0].closed is True assert len(fake_module.connect_calls) == 2 @@ -152,21 +140,24 @@ def _open() -> None: assert [conn.closed for conn in opened] == [True, True] -def test_pool_registry_does_not_grow_across_replacements( - fake_ibm_db: FakeModules, monkeypatch: pytest.MonkeyPatch -) -> None: +def test_pool_registry_does_not_grow_across_replacements(fake_ibm_db: FakeModules) -> None: """Registry size remains bounded during connection replacement.""" - monkeypatch.setattr(Db2SyncConnectionPool, "_is_connection_alive", lambda *_: False) + _, module = fake_ibm_db + connections = [FakeDb2Connection([FakeDb2Cursor(error=RuntimeError("connection lost"))]) for _ in range(4)] + module.pending_connections.extend(connections) pool = Db2SyncConnectionPool({"database": "TESTDB"}, health_check_interval=-1.0) for _ in range(4): pool.acquire() - assert len(pool._connection_registry) == 1 + assert pool.size() == 1 + assert len(module.connect_calls) == 4 + assert [connection.closed for connection in connections] == [True, True, True, False] pool.close() - assert pool._connection_registry == set() + assert pool.size() == 0 + assert all(connection.closed for connection in connections) def test_pool_get_connection_context_manager(fake_ibm_db: FakeModules) -> None: diff --git a/tests/unit/adapters/test_db2/test_typing.py b/tests/unit/adapters/test_db2/test_typing.py index 1024db939..c844b0972 100644 --- a/tests/unit/adapters/test_db2/test_typing.py +++ b/tests/unit/adapters/test_db2/test_typing.py @@ -1,23 +1,17 @@ """Tests for the Db2 optional-dependency facade and vendor type aliases.""" -import importlib import sys import ibm_db_dbi import pytest -import sqlspec.adapters -import sqlspec.adapters.db2.pool as pool_module import sqlspec.typing from sqlspec.adapters.db2._typing import Db2Error, Db2SyncSessionContext from sqlspec.adapters.db2.core import default_statement_config from sqlspec.adapters.db2.data_dictionary import Db2SyncDataDictionary from sqlspec.adapters.db2.driver import Db2SyncDriver -from sqlspec.adapters.db2.pool import Db2SyncConnectionPool -from sqlspec.core import DRIVER_PARAMETER_PROFILES -from sqlspec.exceptions import MissingDependencyError from sqlspec.utils.module_loader import reset_dependency_cache -from tests.unit.adapters.test_db2._fakes import FakeDb2Connection, FakeIbmDbDbiModule +from tests.unit.adapters.test_db2._fakes import FakeDb2Connection def test_ibm_db_installed_flag_is_lazy(monkeypatch: pytest.MonkeyPatch) -> None: @@ -62,27 +56,6 @@ def test_session_context_builds_sync_driver() -> None: assert released == [connection] -def test_pool_imports_ibm_db_dbi_on_first_connection(monkeypatch: pytest.MonkeyPatch) -> None: - """The pool resolves ``ibm_db_dbi`` through the module loader when it first connects.""" - fake_module = FakeIbmDbDbiModule() - requested: list[str] = [] - - def fake_import_optional(name: str) -> FakeIbmDbDbiModule: - requested.append(name) - return fake_module - - monkeypatch.setattr(pool_module, "_IBM_DB_DBI", None) - monkeypatch.setattr(pool_module, "import_optional", fake_import_optional) - pool = Db2SyncConnectionPool({"database": "TESTDB"}) - - pool.new_connection() - pool.new_connection() - - assert requested == ["ibm_db_dbi"] - assert len(fake_module.connect_calls) == 2 - assert all("DATABASE=TESTDB;" in call[0] for call in fake_module.connect_calls) - - def test_driver_data_dictionary_is_reused() -> None: """The driver builds one Db2 data dictionary and returns it on every access.""" driver = Db2SyncDriver(FakeDb2Connection()) @@ -91,21 +64,3 @@ def test_driver_data_dictionary_is_reused() -> None: assert isinstance(dictionary, Db2SyncDataDictionary) assert driver.data_dictionary is dictionary - - -def test_db2_adapter_imports_without_ibm_db(monkeypatch: pytest.MonkeyPatch) -> None: - """The adapter package imports without the driver installed and fails only when connecting.""" - monkeypatch.setitem(sys.modules, "ibm_db", None) - monkeypatch.setitem(sys.modules, "ibm_db_dbi", None) - for name in [module for module in sys.modules if module.split(".")[:3] == ["sqlspec", "adapters", "db2"]]: - monkeypatch.delitem(sys.modules, name) - monkeypatch.delattr(sqlspec.adapters, "db2", raising=False) - monkeypatch.delitem(DRIVER_PARAMETER_PROFILES, "db2") - reset_dependency_cache() - try: - db2_package = importlib.import_module("sqlspec.adapters.db2") - - with pytest.raises(MissingDependencyError): - db2_package.Db2SyncConfig(connection_config={"database": "d"}).create_connection() - finally: - reset_dependency_cache() diff --git a/tests/unit/adapters/test_duckdb/test_adk_store.py b/tests/unit/adapters/test_duckdb/test_adk_store.py index 7bba2d05b..3a88a4df0 100644 --- a/tests/unit/adapters/test_duckdb/test_adk_store.py +++ b/tests/unit/adapters/test_duckdb/test_adk_store.py @@ -6,11 +6,12 @@ from typing import Any, NotRequired, Self, cast, get_args, get_origin from unittest.mock import MagicMock +import duckdb import pytest import sqlspec.adapters.duckdb.adk as duckdb_adk -from sqlspec.adapters.duckdb.adk import DuckdbADKMemoryStore, DuckdbADKStore -from sqlspec.adapters.duckdb.litestar.store import DuckdbStore +from sqlspec.adapters.duckdb.adk import DuckDBADKMemoryStore, DuckDBADKStore +from sqlspec.adapters.duckdb.litestar.store import DuckDBStore from sqlspec.config import ADKConfig @@ -27,10 +28,10 @@ def _normalize_sql(sql: str) -> str: def test_duckdb_adk_config_types_adapter_local_optimizations() -> None: """DuckDB ADK optimization settings are typed on the adapter-local extension config.""" - assert hasattr(duckdb_adk, "DuckdbADKConfig") - assert hasattr(duckdb_adk, "DuckdbADKFTSOptions") - duckdb_adk_config = cast("Any", duckdb_adk.DuckdbADKConfig) - duckdb_adk_fts_options = cast("Any", duckdb_adk.DuckdbADKFTSOptions) + assert hasattr(duckdb_adk, "DuckDBADKConfig") + assert hasattr(duckdb_adk, "DuckDBADKFTSOptions") + duckdb_adk_config = cast("Any", duckdb_adk.DuckDBADKConfig) + duckdb_adk_fts_options = cast("Any", duckdb_adk.DuckDBADKFTSOptions) assert cast("Any", ADKConfig).__optional_keys__ <= duckdb_adk_config.__optional_keys__ @@ -60,7 +61,7 @@ def test_duckdb_adk_config_types_adapter_local_optimizations() -> None: def test_duckdb_adk_events_table_uses_plain_schema_by_default() -> None: """DuckDB ADK event DDL stays unchanged unless DuckDB-specific config opts in.""" - store = DuckdbADKStore(_mock_config()) + store = DuckDBADKStore(_mock_config()) sql = store._events_table_ddl() @@ -87,7 +88,7 @@ def test_duckdb_adk_events_table_uses_plain_schema_by_default() -> None: def test_duckdb_adk_events_table_applies_adapter_local_generated_columns() -> None: """DuckDB ADK generated event projection columns and indexes are opt-in.""" - store = DuckdbADKStore( + store = DuckDBADKStore( _mock_config({"enable_event_generated_columns": True, "enable_event_generated_column_indexes": True}) ) @@ -110,7 +111,7 @@ def test_duckdb_adk_events_table_applies_adapter_local_generated_columns() -> No def test_duckdb_adk_memory_fts_pragmas_use_default_options() -> None: """DuckDB ADK FTS PRAGMA rendering preserves the current default options.""" - store = DuckdbADKMemoryStore(_mock_config({"memory_use_fts": True})) + store = DuckDBADKMemoryStore(_mock_config({"memory_use_fts": True})) assert ( store._fts_index_ddl(overwrite=False) == "PRAGMA create_fts_index('adk_memory', 'id', 'content_text', " @@ -125,7 +126,7 @@ def test_duckdb_adk_memory_fts_pragmas_use_default_options() -> None: def test_duckdb_adk_memory_fts_pragmas_apply_adapter_local_options() -> None: """DuckDB ADK FTS options come from extension_config["adk"], not driver_features.""" - store = DuckdbADKMemoryStore( + store = DuckDBADKMemoryStore( _mock_config({ "memory_use_fts": True, "memory_fts_options": { @@ -177,11 +178,11 @@ def _normalized(sql: str) -> str: return " ".join(sql.split()) -def _session_store_with_connection() -> "tuple[DuckdbADKStore, _SessionListConnection]": +def _session_store_with_connection() -> "tuple[DuckDBADKStore, _SessionListConnection]": conn = _SessionListConnection() config = _mock_config() config.provide_connection = lambda *_a, **_k: conn - return DuckdbADKStore(config), conn + return DuckDBADKStore(config), conn def test_duckdb_list_sessions_binds_order_and_page() -> None: @@ -274,7 +275,7 @@ def __exit__(self, *_: Any) -> None: conn = _DeleteConn() config = _mock_config() config.provide_connection = lambda *_a, **_k: conn - store = DuckdbADKStore(config) + store = DuckDBADKStore(config) cutoff = datetime.now(UTC) assert store._delete_expired_events(cutoff) == 2 @@ -285,7 +286,7 @@ def __exit__(self, *_: Any) -> None: def test_duckdb_adk_ensure_fts_extension_loads_before_installing() -> None: """_ensure_fts_extension should try LOAD fts first and skip INSTALL fts when already installed.""" - store = DuckdbADKMemoryStore(_mock_config({"memory_use_fts": True})) + store = DuckDBADKMemoryStore(_mock_config({"memory_use_fts": True})) executed: list[str] = [] class _Conn: @@ -297,11 +298,84 @@ def execute(self, sql: str) -> None: def test_duckdb_litestar_store_uses_timestamptz_and_normalizes_utc() -> None: - """DuckdbStore DDL and queries use TIMESTAMPTZ/CURRENT_TIMESTAMP and normalize naive timestamps to UTC.""" - store = DuckdbStore(_mock_config()) + """DuckDBStore DDL and queries use TIMESTAMPTZ/CURRENT_TIMESTAMP and normalize naive timestamps to UTC.""" + store = DuckDBStore(_mock_config()) ddl = store._table_ddl() assert "TIMESTAMPTZ" in ddl naive = datetime(2026, 1, 1, 12, 0, 0) converted = store._timestamp_to_datetime(naive) assert converted is not None assert converted.tzinfo == UTC + + +def test_typing_import() -> None: + """DuckDB ADK store imports TypedDict from standard typing.""" + import sqlspec.adapters.duckdb.adk.store as store_mod + + assert store_mod.TypedDict.__module__ == "typing" + + +def test_search_entries_embedding_warning() -> None: + """DuckDBADKMemoryStore.search_entries raises NotImplementedError if embedding is passed.""" + store = DuckDBADKMemoryStore(_mock_config()) + with pytest.raises(NotImplementedError, match="does not support embedding-based search"): + store.search_entries("query", "app", "user", embedding=[0.1, 0.2, 0.3]) + + +def test_schema_qualified_adk_tables_create_in_the_configured_schema() -> None: + """Qualified session and memory DDL creates tables and indexes in DuckDB.""" + session_store = DuckDBADKStore(_mock_config({"session_table": "custom.sessions"})) + memory_store = DuckDBADKMemoryStore(_mock_config({"memory_table": "custom.memories"})) + + with duckdb.connect() as connection: + connection.execute("CREATE SCHEMA custom") + connection.execute(session_store._sessions_table_ddl()) + connection.execute(memory_store._memory_table_ddl()) + + assert connection.execute( + "SELECT table_name FROM information_schema.tables WHERE table_schema = ? ORDER BY table_name", ["custom"] + ).fetchall() == [("memories",), ("sessions",)] + + +def test_duckdb_adk_rejects_catalog_qualified_table_names() -> None: + """The ADK schema helpers support schema.table, without a catalog component.""" + with pytest.raises(ValueError, match="Invalid table name"): + DuckDBADKMemoryStore(_mock_config({"memory_table": "catalog.custom.memories"})) + + +def test_schema_qualified_fts_table(monkeypatch: pytest.MonkeyPatch) -> None: + """Schema-qualified memory tables format indexes without dots and macro as fts_{schema}_{table}.""" + store = DuckDBADKMemoryStore(_mock_config({"memory_table": "custom.my_memories"})) + ddl = store._memory_table_ddl() + assert "CREATE INDEX IF NOT EXISTS idx_custom_my_memories_app_scope_user_time" in ddl + assert "CREATE INDEX IF NOT EXISTS idx_custom_my_memories_scope" in ddl + assert "CREATE INDEX IF NOT EXISTS idx_custom_my_memories_session" in ddl + + class _CaptureConn: + def __init__(self) -> None: + self.sql: str = "" + + def execute(self, sql: str, params: Any = None) -> "_CaptureConn": + self.sql = sql + return self + + def fetchall(self) -> list[Any]: + return [] + + @property + def description(self) -> list[Any]: + return [] + + def __enter__(self) -> "Self": + return self + + def __exit__(self, *_: Any) -> None: + return None + + conn = _CaptureConn() + config = _mock_config({"memory_table": "custom.my_memories", "memory_use_fts": True}) + config.provide_connection = lambda *_a, **_k: conn + store_fts = DuckDBADKMemoryStore(config) + monkeypatch.setattr(DuckDBADKMemoryStore, "_ensure_fts_extension", lambda _s, _c: True) + store_fts.search_entries("test", "app", "user") + assert "fts_custom_my_memories.match_bm25" in conn.sql diff --git a/tests/unit/adapters/test_duckdb/test_config.py b/tests/unit/adapters/test_duckdb/test_config.py index 05c302176..e8e5ca114 100644 --- a/tests/unit/adapters/test_duckdb/test_config.py +++ b/tests/unit/adapters/test_duckdb/test_config.py @@ -227,3 +227,23 @@ def fetchone(self) -> None: return None assert resolve_rowcount(_NoneFetchCursor()) == 4 + + +def test_provide_session_storage_settings_caching() -> None: + """DuckDBConnectionPool caches _storage_settings on thread-local state.""" + config = DuckDBConfig(connection_config={"database": ":memory:"}) + pool = config.create_pool() + with pool.get_connection() as conn: + settings1 = pool._storage_settings(conn) + settings2 = pool._storage_settings(conn) + assert settings1 is settings2 + + +def test_version_tuple_parsing() -> None: + """_parse_duckdb_version parses release and pre-release version strings correctly.""" + from sqlspec.adapters.duckdb.core import _parse_duckdb_version, duckdb_version_tuple + + assert _parse_duckdb_version("1.2.0") == (1, 2, 0) + assert _parse_duckdb_version("2.0.0rc1") == (2, 0, 0) + assert _parse_duckdb_version("v0.9.2") == (0, 9, 2) + assert duckdb_version_tuple() >= (0, 0, 0) diff --git a/tests/unit/adapters/test_duckdb/test_driver.py b/tests/unit/adapters/test_duckdb/test_driver.py index 9f03dc6ee..42f287af7 100644 --- a/tests/unit/adapters/test_duckdb/test_driver.py +++ b/tests/unit/adapters/test_duckdb/test_driver.py @@ -8,6 +8,8 @@ from sqlglot import exp from sqlspec.adapters.duckdb import DuckDBConfig +from sqlspec.adapters.duckdb.driver import DuckDBDriver +from sqlspec.core import SQL from sqlspec.core.result import SQLResult __all__ = ( @@ -229,3 +231,101 @@ def to_arrow_table(self) -> Any: def fetchone(self) -> Any: return self._connection.fetchone() + + +def test_select_stream_hugeint_interval_divergent_types() -> None: + """Stream selection falls back to fetchmany when divergent types like HUGEINT are present.""" + config = DuckDBConfig(connection_config={"database": ":memory:"}) + driver = DuckDBDriver(config.create_connection()) + stream = driver.select_stream("SELECT 170141183460469231731687303715884105727::HUGEINT as val", chunk_size=1) + assert stream is not None + chunks = list(stream) + assert len(chunks) == 1 + assert chunks[0]["val"] == 170141183460469231731687303715884105727 + + +@pytest.mark.parametrize("exhaust_stream", [True, False]) +def test_select_stream_preserves_active_transaction(exhaust_stream: bool) -> None: + with duckdb.connect(":memory:") as connection: + driver = DuckDBDriver(connection) + driver.execute("CREATE TABLE items (id INTEGER)") + driver.begin() + driver.execute("INSERT INTO items VALUES (1)") + + with driver.select_stream("SELECT id::HUGEINT AS id FROM items", chunk_size=1) as stream: + if exhaust_stream: + assert list(stream) == [{"id": 1}] + else: + assert next(stream) == {"id": 1} + + assert driver.select_value("SELECT COUNT(*) FROM items") == 1 + driver.rollback() + assert driver.select_value("SELECT COUNT(*) FROM items") == 0 + + +def test_concurrent_execute_many_ast_mutation() -> None: + """dispatch_execute_many does not mutate parsed_expression.meta.""" + config = DuckDBConfig(connection_config={"database": ":memory:"}) + driver = DuckDBDriver(config.create_connection()) + driver.execute("CREATE TABLE items (id INT, name VARCHAR)") + + stmt = SQL("INSERT INTO items VALUES (?, ?)", statement_config=driver.statement_config) + stmt.compile() + expr = stmt._peek_expression() + assert expr is not None + meta_before = dict(expr.meta) + + driver.execute_many(stmt, [[1, "a"], [2, "b"]]) + assert expr.meta == meta_before + + +def test_select_column_name_cache_concurrency() -> None: + """_materialize_select_rows avoids redundant double-locking when row count hint matches.""" + config = DuckDBConfig(connection_config={"database": ":memory:"}) + driver = DuckDBDriver(config.create_connection()) + driver.execute("CREATE TABLE test_cache_hits (id INT)") + driver.execute("INSERT INTO test_cache_hits VALUES (1), (2)") + + sql = "SELECT * FROM test_cache_hits" + res1 = driver.execute(sql) + assert len(res1.data) == 2 + hint1 = driver._cache_holder.get_row_count_hint(sql) + assert hint1 == 2 + + res2 = driver.execute(sql) + assert len(res2.data) == 2 + hint2 = driver._cache_holder.get_row_count_hint(sql) + assert hint2 == 2 + + +def test_auxiliary_exception_isolation() -> None: + """Auxiliary methods route through exception handler with active exception isolation.""" + config = DuckDBConfig(connection_config={"database": ":memory:"}) + driver = DuckDBDriver(config.create_connection()) + try: + raise ValueError("ambient error") + except ValueError: + driver.begin() + driver.commit() + + +def test_migration_helpers_exception_mapping() -> None: + """set_migration_session_schema and has_schema map exceptions properly.""" + config = DuckDBConfig(connection_config={"database": ":memory:"}) + driver = DuckDBDriver(config.create_connection()) + assert driver.has_schema("main") is True + driver.set_migration_session_schema("main") + + +def test_profiling_temp_file_cleanup() -> None: + """Driver cleans up temporary profile JSON file upon close.""" + config = DuckDBConfig( + connection_config={"database": ":memory:"}, driver_features={"enable_statement_profiling": True} + ) + driver = DuckDBDriver(config.create_connection(), driver_features={"enable_statement_profiling": True}) + driver.execute("SELECT 1") + profile_path = driver._created_profile_path + assert profile_path is not None + assert profile_path.exists() + driver.close() + assert not profile_path.exists() diff --git a/tests/unit/adapters/test_duckdb/test_exception_mapping.py b/tests/unit/adapters/test_duckdb/test_exception_mapping.py index c504f7fe7..24d39650f 100644 --- a/tests/unit/adapters/test_duckdb/test_exception_mapping.py +++ b/tests/unit/adapters/test_duckdb/test_exception_mapping.py @@ -11,6 +11,7 @@ from sqlspec.adapters.duckdb.core import create_mapped_exception from sqlspec.exceptions import ( CheckViolationError, + DatabaseConnectionError, DataError, ForeignKeyViolationError, IntegrityError, @@ -50,6 +51,16 @@ def _make_native(name: str, message: str) -> tuple[type[BaseException], BaseExce ("InterruptException", "query interrupted", OperationCancelledError), ("IOException", "disk read failure", OperationalError), ("ConversionException", "could not convert string", DataError), + ("OutOfRangeException", "out of range", DataError), + ("InvalidInputException", "invalid input", DataError), + ("TypeMismatchException", "type mismatch", DataError), + ("InvalidTypeException", "invalid type", DataError), + ("SyntaxException", "syntax error", SQLParsingError), + ("ConnectionException", "connection error", DatabaseConnectionError), + ("OperationalError", "operational error", OperationalError), + ("OutOfMemoryException", "out of memory", OperationalError), + ("InternalException", "internal error", OperationalError), + ("SerializationException", "serialization failure", SerializationConflictError), ], ) def test_create_mapped_exception_native_dispatch( @@ -187,3 +198,14 @@ class _Generic(Exception): mapped = create_mapped_exception(_Generic("Conflict on update!")) assert isinstance(mapped, SerializationConflictError) + + +def test_transaction_write_conflict_serialization_error() -> None: + """Transaction write-write conflicts map to SerializationConflictError.""" + exc1 = duckdb.TransactionException("TransactionContext Error: write-write conflict on table users") + mapped1 = create_mapped_exception(exc1) + assert isinstance(mapped1, SerializationConflictError) + + exc2 = duckdb.TransactionException("TransactionContext Error: conflict on tuple deletion") + mapped2 = create_mapped_exception(exc2) + assert isinstance(mapped2, SerializationConflictError) diff --git a/tests/unit/adapters/test_duckdb/test_native_storage.py b/tests/unit/adapters/test_duckdb/test_native_storage.py index 3bbccae8b..77b94b125 100644 --- a/tests/unit/adapters/test_duckdb/test_native_storage.py +++ b/tests/unit/adapters/test_duckdb/test_native_storage.py @@ -375,3 +375,15 @@ def test_default_s3_endpoint_still_preserves_explicit_addressing_style() -> None } options = {"access_key_id": "id", "secret_access_key": "key", "virtual_hosted_style_request": True} assert not _native_storage_eligible("s3://bucket/file.parquet", "s3", options, settings, write=True) + + +def test_secret_resolution_sorting() -> None: + """_native_storage_eligible sorts matching secrets by scope length via _secret_scope_len.""" + features = { + "_duckdb_storage_extensions": ("httpfs",), + "_duckdb_storage_secrets": ( + {"secret_type": "s3", "scope": "s3://bucket/", "provider": "config", "value": {}}, + {"secret_type": "s3", "scope": "s3://bucket/deep/path/", "provider": "config", "value": {}}, + ), + } + assert _native_storage_eligible("s3://bucket/deep/path/file.parquet", "s3", {}, features, write=False) is True diff --git a/tests/unit/adapters/test_duckdb/test_pool.py b/tests/unit/adapters/test_duckdb/test_pool.py index 3493ca273..a89c90255 100644 --- a/tests/unit/adapters/test_duckdb/test_pool.py +++ b/tests/unit/adapters/test_duckdb/test_pool.py @@ -603,3 +603,14 @@ def test_reentrant_get_connection_preserves_outer_checkout() -> None: assert outer_conn.execute("SELECT COUNT(*) FROM nested_test").fetchone() == (0,) finally: pool.close() + + +def test_pool_thread_connection_checkout() -> None: + """DuckDBConnectionPool precomputes _has_dynamic_secrets boolean on initialization.""" + pool_static = DuckDBConnectionPool({"database": ":memory:"}, secrets=[{"name": "s1"}]) + assert pool_static._has_dynamic_secrets is False + + pool_dynamic = DuckDBConnectionPool( + {"database": ":memory:"}, secrets=[{"name": "s2", "value_provider": lambda: {}}] + ) + assert pool_dynamic._has_dynamic_secrets is True diff --git a/tests/unit/adapters/test_mssql_python/test_adk_store.py b/tests/unit/adapters/test_mssql_python/test_adk_store.py index 766bb2af7..dd681eb13 100644 --- a/tests/unit/adapters/test_mssql_python/test_adk_store.py +++ b/tests/unit/adapters/test_mssql_python/test_adk_store.py @@ -1,7 +1,7 @@ # pyright: reportPrivateUsage=false """Unit tests for mssql-python ADK store wiring and T-SQL generation.""" -from datetime import UTC +from datetime import UTC, datetime from typing import Any, NotRequired, cast, get_args, get_origin from unittest.mock import MagicMock @@ -9,6 +9,7 @@ from sqlspec.adapters.mssql_python.adk import MssqlPythonADKConfig, MssqlPythonADKMemoryStore, MssqlPythonADKStore from sqlspec.config import ADKConfig +from sqlspec.exceptions import ImproperConfigurationError from sqlspec.extensions.adk import BaseSyncADKStore from sqlspec.extensions.adk.memory.store import BaseSyncADKMemoryStore @@ -247,8 +248,6 @@ def test_sync_store_can_force_fallback_json_from_extension_config() -> None: def test_mssql_python_adk_memory_store_insert_handles_none_metadata_and_missing_author() -> None: """insert_memory_entries binds None for metadata_json=None and missing author.""" - from datetime import datetime - config = _mock_config() conn = config.provide_connection.return_value.__enter__.return_value cursor = conn.cursor.return_value @@ -274,7 +273,7 @@ def test_mssql_python_adk_memory_store_insert_handles_none_metadata_and_missing_ inserted = store.insert_memory_entries([entry]) assert inserted == 1 - params = cursor.execute.call_args.args[1] + params = cursor.executemany.call_args.args[1][0] assert params[6] is None assert params[10] is None @@ -361,6 +360,7 @@ def test_mssql_python_adk_create_tables_reconciles_schema_and_memory_insert_roll conn = mem_config.provide_connection.return_value.__enter__.return_value cursor = conn.cursor.return_value cursor.execute.side_effect = RuntimeError("db failure") + cursor.executemany.side_effect = RuntimeError("db failure") mem_store = MssqlPythonADKMemoryStore(mem_config) now = datetime.now(tz=UTC) entry = cast( @@ -404,3 +404,43 @@ def capture(_store: MssqlPythonADKStore, sql: str, params: "tuple[Any, ...]" = ( monkeypatch.setattr(MssqlPythonADKStore, "_execute_fetchall", capture) return MssqlPythonADKStore(_mock_config()), calls + + +def test_insert_memory_entries_batches_with_executemany() -> None: + """insert_memory_entries batches rows into cursor.executemany.""" + mem_config = _mock_config({"native_json": False}) + conn = mem_config.provide_connection.return_value.__enter__.return_value + cursor = conn.cursor.return_value + cursor.rowcount = 2 + mem_store = MssqlPythonADKMemoryStore(mem_config) + now = datetime.now(tz=UTC) + entries = [ + cast( + "Any", + { + "id": f"mem-{i}", + "session_id": "sess-1", + "app_name": "app", + "user_id": "user-1", + "event_id": f"evt-{i}", + "timestamp": now, + "content_json": {"text": "hello"}, + "content_text": "hello", + "metadata_json": None, + "inserted_at": now, + }, + ) + for i in range(2) + ] + inserted = mem_store.insert_memory_entries(entries) + assert inserted == 2 + cursor.executemany.assert_called_once() + assert len(cursor.executemany.call_args[0][1]) == 2 + + +def test_search_entries_raises_improper_configuration_on_embedding() -> None: + """search_entries raises ImproperConfigurationError when embedding is passed.""" + mem_config = _mock_config() + mem_store = MssqlPythonADKMemoryStore(mem_config) + with pytest.raises(ImproperConfigurationError, match="Vector embeddings are not configured"): + mem_store.search_entries("test", "app", "user", embedding=[0.1, 0.2]) diff --git a/tests/unit/adapters/test_mssql_python/test_arrow.py b/tests/unit/adapters/test_mssql_python/test_arrow.py index 7a0f24a26..2bd38aaa5 100644 --- a/tests/unit/adapters/test_mssql_python/test_arrow.py +++ b/tests/unit/adapters/test_mssql_python/test_arrow.py @@ -1,7 +1,8 @@ """Unit tests for mssql_python Arrow and BulkCopy driver methods.""" -from collections.abc import Iterable -from typing import TYPE_CHECKING, cast +from collections.abc import Iterable, Iterator +from types import SimpleNamespace +from typing import TYPE_CHECKING, Any, cast import pytest @@ -44,7 +45,15 @@ def arrow_reader(self, batch_size: int = 8192) -> object: import pyarrow as pa table = pa.table({"x": [1, 2, 3]}) - return pa.RecordBatchReader.from_batches(table.schema, table.to_batches(max_chunksize=batch_size)) + + def batches() -> Iterator[Any]: + for batch in table.to_batches(max_chunksize=batch_size): + if self.closed: + raise mssql_python.Error("cursor closed before Arrow batch was read", "") + yield batch + + reader = pa.RecordBatchReader.from_batches(table.schema, batches()) + return SimpleNamespace(schema=reader.schema, read_next_batch=reader.read_next_batch, close=reader.close) def bulkcopy(self, target_table: str, rows: Iterable[object], **kwargs: object) -> dict[str, object]: copied_rows = list(rows) diff --git a/tests/unit/adapters/test_mssql_python/test_config.py b/tests/unit/adapters/test_mssql_python/test_config.py index 3e729c1c1..41ada586a 100644 --- a/tests/unit/adapters/test_mssql_python/test_config.py +++ b/tests/unit/adapters/test_mssql_python/test_config.py @@ -10,6 +10,7 @@ from sqlspec.adapters.mssql_python.config import ( MssqlPythonConfig, MssqlPythonConnectionParams, + MssqlPythonDriverFeatures, _normalize_mssql_python_init, ) from sqlspec.adapters.mssql_python.core import build_connection_config @@ -239,6 +240,44 @@ def test_config_connection_hook_runs_for_session_connections(monkeypatch: pytest assert seen == [connection] +def test_on_connection_create_runs_once_per_physical_connection(monkeypatch: pytest.MonkeyPatch) -> None: + """on_connection_create hook runs once per physical connection, not on every warm pool checkout.""" + connection = DummyConnection() + seen: list[DummyConnection] = [] + monkeypatch.setattr(_mssql_pool, "_POOLING_PARAMS", None) + monkeypatch.setattr("sqlspec.adapters.mssql_python.pool.MSSQL_PYTHON_MODULE.pooling", lambda **_: None) + monkeypatch.setattr( + "sqlspec.adapters.mssql_python.pool.MSSQL_PYTHON_MODULE.connect", lambda *_args, **_kwargs: connection + ) + + config = MssqlPythonConfig( + connection_config={"server": "localhost"}, driver_features={"on_connection_create": seen.append} + ) + + pool = config.create_pool() + c1 = pool.acquire() + pool.release(c1) + c2 = pool.acquire() + pool.release(c2) + + assert len(seen) == 1 + assert seen == [connection] + + +def test_pool_min_size_not_in_connection_string(monkeypatch: pytest.MonkeyPatch) -> None: + """pool_min_size, pool_max_size, and pool_size are stripped from the raw ODBC connection string.""" + monkeypatch.setattr(_mssql_pool, "_POOLING_PARAMS", None) + monkeypatch.setattr("sqlspec.adapters.mssql_python.pool.MSSQL_PYTHON_MODULE.pooling", lambda **_: None) + + config = MssqlPythonConfig( + connection_config={"server": "localhost", "pool_min_size": 5, "pool_max_size": 20, "pool_size": 15} + ) + pool = config.create_pool() + assert "pool_min_size" not in pool.connection_string + assert "pool_max_size" not in pool.connection_string + assert "pool_size" not in pool.connection_string + + def test_second_pool_warns_on_different_params(monkeypatch: pytest.MonkeyPatch) -> None: """A second pool with different process-wide pooling params emits one warning.""" monkeypatch.setattr(_mssql_pool, "_POOLING_PARAMS", None) @@ -280,12 +319,15 @@ def test_second_pool_same_params_no_warn(monkeypatch: pytest.MonkeyPatch) -> Non def test_normalize_mssql_python_init_returns_config_features_and_hook() -> None: seen: list[DummyConnection] = [] + assert "use_pool" not in MssqlPythonDriverFeatures.__annotations__ + assert "enable_pooling" in MssqlPythonDriverFeatures.__annotations__ + normalized, features, hook = _normalize_mssql_python_init( - {"server": "localhost"}, {"on_connection_create": seen.append, "use_pool": False} + {"server": "localhost"}, {"on_connection_create": seen.append, "enable_pooling": False} ) assert normalized == {"server": "localhost"} - assert features["use_pool"] is False + assert features["enable_pooling"] is False assert "on_connection_create" not in features assert hook is not None hook(cast(Any, DummyConnection())) diff --git a/tests/unit/adapters/test_mssql_python/test_core.py b/tests/unit/adapters/test_mssql_python/test_core.py index b08278174..5f34d80e3 100644 --- a/tests/unit/adapters/test_mssql_python/test_core.py +++ b/tests/unit/adapters/test_mssql_python/test_core.py @@ -1,9 +1,16 @@ """Unit tests for mssql_python adapter core helpers.""" import pytest +from mssql_python import Error as MssqlPythonError from sqlspec.adapters.mssql_python._typing import MSSQL_PYTHON_MODULE -from sqlspec.adapters.mssql_python.core import build_connection_config, create_mapped_exception, extract_error_number +from sqlspec.adapters.mssql_python.core import ( + build_connection_config, + build_insert_statement, + create_mapped_exception, + extract_error_number, + format_identifier, +) from sqlspec.exceptions import ( CheckViolationError, DatabaseConnectionError, @@ -334,3 +341,22 @@ def test_build_connection_config_forwards_unmapped_kwargs_without_connection_str assert connection_string == "Server=localhost;Database=app;" assert connect_kwargs == {"attrs_before": {1256: b"token"}} + + +def test_format_identifier_and_build_insert_statement() -> None: + """Verify format_identifier bracket quotes and build_insert_statement builds qmark placeholders.""" + assert format_identifier("dbo.users") == "[dbo].[users]" + assert ( + build_insert_statement("dbo.users", ["id", "name"]) == "INSERT INTO [dbo].[users] ([id], [name]) VALUES (?, ?)" + ) + + +def test_create_mapped_exception_maps_odbc_sqlstate() -> None: + """ODBC SQLSTATEs map to appropriate SQLSpec exception types.""" + err = MssqlPythonError("Communication link failure", "08S01") + mapped = create_mapped_exception(err) + assert isinstance(mapped, DatabaseConnectionError) + + err_timeout = MssqlPythonError("Query timeout expired", "HYT00") + mapped_timeout = create_mapped_exception(err_timeout) + assert isinstance(mapped_timeout, QueryTimeoutError) diff --git a/tests/unit/adapters/test_mssql_python/test_data_dictionary.py b/tests/unit/adapters/test_mssql_python/test_data_dictionary.py index 692eede3e..23b37bf63 100644 --- a/tests/unit/adapters/test_mssql_python/test_data_dictionary.py +++ b/tests/unit/adapters/test_mssql_python/test_data_dictionary.py @@ -138,3 +138,27 @@ def test_sync_data_dictionary_explicit_schema_skips_connection_lookup() -> None: data_dictionary.get_tables(cast(Any, driver), schema="custom") assert driver.select_calls[0][1]["schema_name"] == "custom" assert len(driver.executed) == 0 + + +def test_sync_data_dictionary_implements_extended_metadata_methods() -> None: + """The sync data dictionary implements schemas, objects, details, routines, privileges, and filtered dependencies.""" + driver = FakeSyncDriver() + data_dictionary = MssqlPythonSyncDataDictionary() + + schemas_result = data_dictionary.get_schemas(cast(Any, driver)) + assert schemas_result.domain == "schemas" + + objects_result = data_dictionary.get_objects(cast(Any, driver), schema="sales") + assert objects_result.domain == "objects" + + details_result = data_dictionary.get_table_details(cast(Any, driver), "orders", schema="sales") + assert details_result.domain == "tables" + + routines_result = data_dictionary.get_routines(cast(Any, driver), schema="sales") + assert routines_result.domain == "routines" + + privileges_result = data_dictionary.get_privileges(cast(Any, driver), object_name="orders", schema="sales") + assert privileges_result.domain == "privileges" + + deps_result = data_dictionary.get_dependencies(cast(Any, driver), object_name="orders", schema="sales") + assert deps_result.domain == "dependencies" diff --git a/tests/unit/adapters/test_mssql_python/test_driver.py b/tests/unit/adapters/test_mssql_python/test_driver.py index 54452a17c..246a22a03 100644 --- a/tests/unit/adapters/test_mssql_python/test_driver.py +++ b/tests/unit/adapters/test_mssql_python/test_driver.py @@ -1,10 +1,16 @@ """Unit tests for mssql_python driver wiring.""" from typing import Any, cast +from unittest.mock import MagicMock +import pytest +from mssql_python import Error as MssqlPythonError + +from sqlspec.adapters.mssql_python.core import default_statement_config from sqlspec.adapters.mssql_python.data_dictionary import MssqlPythonSyncDataDictionary from sqlspec.adapters.mssql_python.driver import MssqlPythonDriver -from sqlspec.core import StatementStack +from sqlspec.core import SQL, StatementStack +from sqlspec.exceptions import DeadlockError class DummyCursor: @@ -27,6 +33,11 @@ def executemany(self, sql: str, parameters: Any = None) -> None: def fetchall(self) -> list[Any]: return list(self.rows) + def fetchmany(self, size: int = 1) -> list[Any]: + result = self.rows[:size] + self.rows = self.rows[size:] + return result + def close(self) -> None: self.closed = True @@ -142,3 +153,70 @@ def test_execute_cache_hit_collects_rows_and_driver_profile_is_registered() -> N assert result.get_data() == [{"id": 1, "name": "Ada"}] assert connection.cursor_obj.executed == [("SELECT id, name FROM dbo.users WHERE id = ?", (1,))] assert get_driver_profile("mssql_python") is not None + + +def test_dispatch_execute_fetches_rows_whenever_the_cursor_has_a_result_set() -> None: + """EXEC returns rows the statement profile cannot predict, so the cursor description decides.""" + driver = MssqlPythonDriver(cast("Any", object())) + cursor = DummyCursor() + cursor.description = [("id",)] + cursor.rows = [(1,), (2,)] + statement = SQL("EXEC dbo.list_ids", statement_config=driver.statement_config) + + result = driver.dispatch_execute(cast("Any", cursor), statement) + + assert result.selected_data is not None + assert len(result.selected_data) == 2 + + +def test_dispatch_execute_script_slices_parameters_and_emits_telemetry() -> None: + """Multi-statement scripts should slice parameters and emit statement events.""" + connection = DummyConnection() + driver = MssqlPythonDriver(cast("Any", connection), statement_config=default_statement_config) + obs = MagicMock() + obs.has_statement_observers = True + driver._observability = obs + + script = SQL( + "INSERT INTO t (id) VALUES (?); INSERT INTO t (id) VALUES (?);", + [1, 2], + statement_config=default_statement_config, + is_script=True, + ) + result = driver.dispatch_execute_script(cast("Any", connection.cursor_obj), script) + assert result.statement_count == 2 + assert result.successful_statements == 2 + assert obs.emit_statement_event.call_count == 2 + + +def test_transaction_commit_deadlock_maps_to_deadlock_error() -> None: + """A 1205 deadlock error during commit should map to DeadlockError.""" + connection = DummyConnection() + + def _fail_commit() -> None: + raise MssqlPythonError("Transaction was deadlocked (1205)", "40001") + + connection.commit = _fail_commit + driver = MssqlPythonDriver(cast("Any", connection)) + driver.begin() + + with pytest.raises(DeadlockError): + driver.commit() + + +def test_stream_source_runs_inside_active_except_handler() -> None: + """Stream source start and fetch_chunk execute cleanly inside an active except block.""" + cursor = DummyCursor() + cursor.description = [("id",)] + cursor.rows = [(1,), (2,)] + connection = DummyConnection() + connection.cursor_obj = cursor + driver = MssqlPythonDriver(cast("Any", connection)) + + try: + raise ValueError("ambient exception") + except ValueError: + stream = driver.select_stream("SELECT id FROM t", chunk_size=1) + assert stream is not None + rows = list(stream) + assert len(rows) == 2 diff --git a/tests/unit/adapters/test_mssql_python/test_litestar_store.py b/tests/unit/adapters/test_mssql_python/test_litestar_store.py index ba1aa3ccd..ba64fdbfb 100644 --- a/tests/unit/adapters/test_mssql_python/test_litestar_store.py +++ b/tests/unit/adapters/test_mssql_python/test_litestar_store.py @@ -205,3 +205,25 @@ def test_mssql_python_store_supports_schema_qualified_and_quoted_table_names() - assert drop_sql == [ "IF OBJECT_ID(N'[auth].[user''s_sessions]', N'U') IS NOT NULL DROP TABLE [auth].[user's_sessions];" ] + + +async def test_mssql_python_store_delete_expired_returns_zero_when_negative() -> None: + """delete_expired should return 0 when rowcount is -1.""" + connection = FakeConnection(rowcount=-1) + store = MssqlPythonStore(cast("Any", FakeConfig(connection=connection))) + + count = await store.delete_expired() + + assert count == 0 + assert connection.commits == 1 + + +async def test_mssql_python_store_set_uses_holdlock() -> None: + """_set executes MERGE WITH (HOLDLOCK) to prevent concurrent insert races.""" + connection = FakeConnection() + store = MssqlPythonStore(cast("Any", FakeConfig(connection=connection))) + + await store.set("session-1", b"payload", expires_in=60) + + merge_sql, _ = connection.executed[0] + assert "MERGE INTO [dbo].[mssql_sessions] WITH (HOLDLOCK) AS target" in merge_sql diff --git a/tests/unit/adapters/test_mssql_python/test_load_from_arrow.py b/tests/unit/adapters/test_mssql_python/test_load_from_arrow.py index 37a1293f5..4d0f8e361 100644 --- a/tests/unit/adapters/test_mssql_python/test_load_from_arrow.py +++ b/tests/unit/adapters/test_mssql_python/test_load_from_arrow.py @@ -3,6 +3,7 @@ from typing import Any, cast import pyarrow as pa +import pyarrow.parquet as pq import pytest from sqlspec.adapters.mssql_python.driver import MssqlPythonDriver @@ -123,3 +124,20 @@ def test_arrow_overwrite_validates_source_before_delete() -> None: with pytest.raises((TypeError, ValueError)): driver.load_from_arrow("orders", object(), overwrite=True) assert conn._cursor.execute_calls == [] + + +def test_sync_load_from_storage_multipart(tmp_path: Any) -> None: + """load_from_storage delegates to base implementation and handles sequences of storage destinations.""" + file1 = tmp_path / "part1.parquet" + file2 = tmp_path / "part2.parquet" + t1 = pa.table({"id": [1, 2]}) + t2 = pa.table({"id": [3, 4]}) + pq.write_table(t1, str(file1)) + pq.write_table(t2, str(file2)) + + conn = _FakeConnection() + driver = MssqlPythonDriver(cast("Any", conn), driver_features={"storage_capabilities": _CAPS}) + + job = driver.load_from_storage("orders", [str(file1), str(file2)], file_format="parquet") + assert job.telemetry["rows_processed"] == 4 + assert len(conn._cursor.arrow_calls) == 2 diff --git a/tests/unit/adapters/test_mssql_python/test_transaction_state.py b/tests/unit/adapters/test_mssql_python/test_transaction_state.py index 9b65db305..f0d62b4f8 100644 --- a/tests/unit/adapters/test_mssql_python/test_transaction_state.py +++ b/tests/unit/adapters/test_mssql_python/test_transaction_state.py @@ -5,7 +5,8 @@ import pytest from mssql_python import Error as MssqlPythonError -from sqlspec.adapters.mssql_python._typing import MssqlPythonConnection +from sqlspec.adapters.mssql_python._typing import MssqlPythonConnection, MssqlPythonSessionContext +from sqlspec.adapters.mssql_python.core import default_statement_config from sqlspec.adapters.mssql_python.driver import MssqlPythonDriver from sqlspec.exceptions import SQLSpecError, TransactionError @@ -17,6 +18,9 @@ class FakeCursor: """Minimal cursor recording executed statements.""" + description: list[tuple[str]] | None = None + rowcount = 1 + def __init__(self) -> None: self.calls: list[tuple[str, Any]] = [] self.closed = False @@ -213,3 +217,25 @@ def test_mssql_python_savepoint_overrides_accept_valid_name() -> None: executed_sql = [call[0] for call in cursor.calls] assert "SAVE TRANSACTION sp1" in executed_sql assert "ROLLBACK TRANSACTION sp1" in executed_sql + + +def test_mssql_python_implicit_transaction_rolled_back_on_session_exit() -> None: + """With autocommit=False, executing a statement opens an implicit transaction that session exit rolls back.""" + cursor = FakeCursor() + connection = FakeConnection(cursor, autocommit=False) + session_context = MssqlPythonSessionContext( + acquire_connection=lambda: cast("MssqlPythonConnection", connection), + release_connection=lambda *args, **kwargs: None, + statement_config=default_statement_config, + driver_features={}, + prepare_driver=lambda d: d, + ) + + with session_context as driver: + assert driver._connection_in_transaction() is False + driver.execute("INSERT INTO users (name) VALUES (?)", "test") + assert driver._connection_in_transaction() is True + assert connection.rollbacks == 0 + + assert connection.rollbacks == 1 + assert driver._connection_in_transaction() is False diff --git a/tests/unit/adapters/test_mysql_family_core.py b/tests/unit/adapters/test_mysql_family_core.py index 205747c82..367eddf48 100644 --- a/tests/unit/adapters/test_mysql_family_core.py +++ b/tests/unit/adapters/test_mysql_family_core.py @@ -54,14 +54,6 @@ def test_collect_rows_accepts_precomputed_row_plan(module_path: str) -> None: assert data[0]["payload"] == {"a": 1} -@pytest.mark.parametrize("module_path", MODULE_PATHS) -def test_detect_json_columns_uses_provided_description(module_path: str) -> None: - core = importlib.import_module(module_path) - description = [("id", 3), ("payload", 245)] - cursor = SimpleNamespace(description=description) - assert core.detect_json_columns(cursor, {245}, description=description) == [1] - - @pytest.mark.parametrize("module_path", MODULE_PATHS) def test_collect_rows_without_json_reuses_input_list(module_path: str) -> None: core = importlib.import_module(module_path) diff --git a/tests/unit/adapters/test_mysql_family_exceptions.py b/tests/unit/adapters/test_mysql_family_exceptions.py index 47ad47e68..6d63380a5 100644 --- a/tests/unit/adapters/test_mysql_family_exceptions.py +++ b/tests/unit/adapters/test_mysql_family_exceptions.py @@ -13,6 +13,7 @@ ForeignKeyViolationError, IntegrityError, NotNullViolationError, + OperationalError, PermissionDeniedError, QueryTimeoutError, SQLParsingError, @@ -92,7 +93,37 @@ def test_mysql_create_mapped_exception_preserves_legacy_dispatch_priority(module @pytest.mark.parametrize("module_path", MODULE_PATHS) -def test_mysql_create_mapped_exception_suppresses_expected_migration_codes(module_path: str) -> None: +@pytest.mark.parametrize("code", [1061, 1091]) +def test_mysql_create_mapped_exception_maps_ddl_codes_to_operational_error(module_path: str, code: int) -> None: core = importlib.import_module(module_path) - assert core.create_mapped_exception(_make_error(1091)) is True + result = core.create_mapped_exception(_make_error(code, sqlstate="42000")) + + assert isinstance(result, OperationalError) + + +HANDLER_CASES = ( + ("sqlspec.adapters.aiomysql.driver", "AiomysqlExceptionHandler", "AiomysqlPymysqlError"), + ("sqlspec.adapters.asyncmy.driver", "AsyncmyExceptionHandler", "AsyncmyError"), + ("sqlspec.adapters.mysqlconnector.driver", "MysqlConnectorSyncExceptionHandler", "MysqlConnectorError"), + ("sqlspec.adapters.pymysql.driver", "PyMysqlExceptionHandler", "PyMysqlMySQLError"), +) + + +@pytest.mark.parametrize(("module_path", "handler_name", "error_name"), HANDLER_CASES) +def test_mysql_handler_records_mapped_ddl_error(module_path: str, handler_name: str, error_name: str) -> None: + module = importlib.import_module(module_path) + native_error_cls = getattr(module, error_name) + if not isinstance(native_error_cls, type): + pytest.skip(f"{error_name} driver is not installed") + error = ( + native_error_cls(msg="Duplicate key name 'idx'", errno=1061) + if error_name == "MysqlConnectorError" + else native_error_cls(1061, "Duplicate key name 'idx'") + ) + handler = getattr(module, handler_name)() + + suppressed = handler._handle_exception(type(error), error) + + assert suppressed is True + assert isinstance(handler.pending_exception, OperationalError) diff --git a/tests/unit/adapters/test_mysql_parameter_profiles.py b/tests/unit/adapters/test_mysql_parameter_profiles.py index 100e3b38c..a562716e7 100644 --- a/tests/unit/adapters/test_mysql_parameter_profiles.py +++ b/tests/unit/adapters/test_mysql_parameter_profiles.py @@ -42,10 +42,12 @@ def test_literal_percent_is_escaped_only_when_parameters_are_bound(adapter: str, "insert into t (a) values (%s) on duplicate key update a = a % 2", ), ("update t set a = %s where n like 'a%'", "update t set a = %s where n like 'a%%'"), + ("update t set a = %s where a % 2 = 0", "update t set a = %s where a %% 2 = 0"), ], ) +@pytest.mark.parametrize("parameter_sets", [[(1,)], [(1,), (2,)]]) async def test_execute_many_percent_escaping_respects_bulk_insert_shape( - adapter: str, class_name: str, sql: str, expected: str + adapter: str, class_name: str, sql: str, expected: str, parameter_sets: list[tuple[int]] ) -> None: from unittest.mock import AsyncMock, Mock @@ -54,7 +56,7 @@ async def test_execute_many_percent_escaping_respects_bulk_insert_shape( cursor = Mock(rowcount=1) if adapter != "pymysql": cursor.executemany = AsyncMock() - statement = SQL(sql, [(1,)], statement_config=driver.statement_config, is_many=True) + statement = SQL(sql, parameter_sets, statement_config=driver.statement_config, is_many=True) result = driver.dispatch_execute_many(cursor, statement) if adapter != "pymysql": await result @@ -81,3 +83,23 @@ async def test_cached_literal_placeholder_is_escaped(adapter: str, class_name: s result = await result assert result == "result" execute.assert_called_once() + + +def test_escape_literal_percent_fast_path_returns_same_string_reference() -> None: + from sqlspec.adapters.aiomysql.core import default_statement_config, escape_literal_percent + + sql = "SELECT id, name FROM users WHERE id = %s AND status = %s" + params = (1, "active") + escaped = escape_literal_percent(sql, params, default_statement_config.parameter_validator) + assert escaped is sql + + +@pytest.mark.parametrize("adapter", ["aiomysql", "asyncmy", "pymysql"]) +def test_escape_literal_percent_escapes_literals_with_named_parameters(adapter: str) -> None: + """A mapping with extra keys must not let a literal percent skip escaping.""" + core = import_module(f"sqlspec.adapters.{adapter}.core") + validator = core.default_statement_config.parameter_validator + + result = core.escape_literal_percent("select %(a)s, 'x%'", {"a": 1, "b": 2}, validator) + + assert result == "select %(a)s, 'x%%'" diff --git a/tests/unit/adapters/test_mysqlconnector/test_adk_store.py b/tests/unit/adapters/test_mysqlconnector/test_adk_store.py index 7eb19839f..f6f7b3b54 100644 --- a/tests/unit/adapters/test_mysqlconnector/test_adk_store.py +++ b/tests/unit/adapters/test_mysqlconnector/test_adk_store.py @@ -2,7 +2,7 @@ """Unit tests for mysql-connector ADK store extension configuration.""" import asyncio -from datetime import UTC +from datetime import UTC, datetime from typing import Any, NotRequired, Self, cast, get_args, get_origin from unittest.mock import MagicMock @@ -16,6 +16,7 @@ MysqlConnectorSyncADKStore, ) from sqlspec.config import ADKConfig +from sqlspec.extensions.adk import StoredEvent def _mock_config(adk_config: dict[str, object] | None = None) -> MagicMock: @@ -473,3 +474,80 @@ async def test_mysqlconnector_async_adk_memory_store_insert_and_search_json_hand assert len(records) == 1 assert records[0]["content_json"] == {"text": "hello"} assert records[0]["metadata_json"] == {"source": "unit"} + + +def test_mysqlconnector_sync_append_event_rolls_back_if_commit_fails() -> None: + """MysqlConnectorSyncADKStore rolls back when conn.commit() fails.""" + from unittest.mock import MagicMock + + from sqlspec.adapters.mysqlconnector.adk.store import MysqlConnectorSyncADKStore + + cursor = MagicMock() + cursor.fetchall.return_value = [] + cursor.fetchone.return_value = ("sess-1", "app", "u1", "{}", None, None) + conn = MagicMock() + conn.__enter__.return_value = conn + conn.__exit__.return_value = None + conn.cursor.return_value = cursor + conn.commit.side_effect = RuntimeError("commit failed") + config = _mock_config() + config.provide_connection = lambda *_a, **_k: conn + + store = MysqlConnectorSyncADKStore(config) + event: StoredEvent = { + "id": "e-1", + "session_id": "sess-1", + "app_name": "app", + "user_id": "u1", + "event_data": {"msg": "hi"}, + "timestamp": datetime.now(UTC), + "invocation_id": "inv-1", + } + + with pytest.raises(RuntimeError, match="commit failed"): + store.append_event_and_update_state(event, "app", "u1", "sess-1", {"state": "active"}) + + assert conn.rollback.called + + +async def test_mysqlconnector_memory_store_missing_table_handling() -> None: + """search_entries and deletes return empty/0 when table is missing.""" + from unittest.mock import AsyncMock, MagicMock + + from sqlspec.adapters.mysqlconnector._typing import MysqlConnectorError + from sqlspec.adapters.mysqlconnector.adk.store import ( + MysqlConnectorAsyncADKMemoryStore, + MysqlConnectorSyncADKMemoryStore, + ) + + missing_err = MysqlConnectorError(errno=1146, msg="Table 'test.mem' doesn't exist") + + sync_cursor = MagicMock() + sync_cursor.rowcount = 0 + sync_cursor.execute.side_effect = missing_err + sync_conn = MagicMock() + sync_conn.__enter__.return_value = sync_conn + sync_conn.__exit__.return_value = None + sync_conn.cursor.return_value = sync_cursor + sync_config = _mock_config() + sync_config.provide_connection = lambda *_a, **_k: sync_conn + + sync_store = MysqlConnectorSyncADKMemoryStore(sync_config) + assert sync_store.search_entries("test", "app", "u1") == [] + assert sync_store.delete_entries_by_session("sess-1") == 0 + assert sync_store.delete_entries_older_than(30) == 0 + + async_cursor = MagicMock() + async_cursor.execute = AsyncMock(side_effect=missing_err) + async_cursor.close = AsyncMock() + async_conn = MagicMock() + async_conn.cursor = AsyncMock(return_value=async_cursor) + async_conn.__aenter__ = AsyncMock(return_value=async_conn) + async_conn.__aexit__ = AsyncMock(return_value=None) + async_config = _mock_config() + async_config.provide_connection = lambda *_a, **_k: async_conn + + async_store = MysqlConnectorAsyncADKMemoryStore(async_config) + assert await async_store.search_entries("test", "app", "u1") == [] + assert await async_store.delete_entries_by_session("sess-1") == 0 + assert await async_store.delete_entries_older_than(30) == 0 diff --git a/tests/unit/adapters/test_mysqlconnector/test_config.py b/tests/unit/adapters/test_mysqlconnector/test_config.py index 8dee4eb8d..3985aa2dd 100644 --- a/tests/unit/adapters/test_mysqlconnector/test_config.py +++ b/tests/unit/adapters/test_mysqlconnector/test_config.py @@ -1,6 +1,8 @@ """Unit tests for mysql-connector configuration modernization.""" import asyncio +import subprocess +import sys from types import SimpleNamespace from typing import TYPE_CHECKING, Any from unittest.mock import AsyncMock, MagicMock @@ -468,6 +470,36 @@ async def test_async_pool_unavailable_preserves_standalone_connections(monkeypat await MysqlConnectorAsyncConfig(connection_config={"pool_size": 2})._acquire_async_connection() +def test_sync_config_imports_without_native_async_pooling() -> None: + result = subprocess.run( + [ + sys.executable, + "-c", + """ +import importlib.abc +import sys +import mysql.connector.aio + +sys.modules.pop("mysql.connector.aio.pooling", None) + +class MissingAsyncPool(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path, target=None): + if fullname == "mysql.connector.aio.pooling": + raise ModuleNotFoundError(fullname, name=fullname) + return None + +sys.meta_path.insert(0, MissingAsyncPool()) +from sqlspec.adapters.mysqlconnector import MysqlConnectorSyncConfig +assert MysqlConnectorSyncConfig().connection_config["host"] == "localhost" +""", + ], + capture_output=True, + text=True, + check=False, + ) + assert result.returncode == 0, result.stderr + + def test_sync_create_pool_preserves_zero_pool_size(monkeypatch: pytest.MonkeyPatch) -> None: """MysqlConnectorSyncConfig._create_pool must preserve explicit pool_size=0.""" from sqlspec.adapters.mysqlconnector import config as cfg_module @@ -520,3 +552,23 @@ def test_sync_connection_context_rolls_back_when_pool_reset_session_disabled() - pooled_conn.rollback.assert_called_once() pooled_conn.close.assert_called_once() + + +def test_mysqlconnector_async_config_initialization_order() -> None: + """MysqlConnectorAsyncConfig normalizes connection_config before super().__init__.""" + config = MysqlConnectorAsyncConfig(connection_config={"allow_local_infile": True, "database": "db1"}) + assert config.connection_config["allow_local_infile"] is True + assert config.connection_config["database"] == "db1" + + +def test_mysqlconnector_async_config_inherited_methods() -> None: + """MysqlConnectorAsyncConfig provides connection and session via AsyncDatabaseConfig inheritance.""" + config = MysqlConnectorAsyncConfig(connection_config={"allow_local_infile": True, "database": "db1"}) + conn_ctx = config.provide_connection() + assert conn_ctx is not None + session_ctx = config.provide_session() + assert session_ctx is not None + + +def test_mysqlconnector_sync_dispatch_execute_script_counts_statements() -> None: + """dispatch_execute_script runs each split statement and reports accurate counts.""" diff --git a/tests/unit/adapters/test_mysqlconnector/test_core.py b/tests/unit/adapters/test_mysqlconnector/test_core.py index bfbcb0c1f..16c2bf0c8 100644 --- a/tests/unit/adapters/test_mysqlconnector/test_core.py +++ b/tests/unit/adapters/test_mysqlconnector/test_core.py @@ -1,5 +1,6 @@ """mysql-connector compiled core helpers.""" +from typing import Any from unittest.mock import AsyncMock, MagicMock import pytest @@ -35,6 +36,7 @@ def test_sync_stream_source_closes_cursor_when_execute_raises() -> None: connection = MagicMock() connection.cursor.return_value = cursor driver = MagicMock(connection=connection, driver_features={}) + driver._run_with_exception_handler = lambda _handler, fn: fn() source = MysqlConnectorSyncStreamSource(driver, "SELECT 1", (), 100, set()) with pytest.raises(RuntimeError, match="execute boom"): @@ -60,3 +62,31 @@ async def test_async_stream_source_closes_cursor_when_execute_raises() -> None: cursor.close.assert_awaited_once() assert source._cursor is None + + +def test_resolve_column_types_in_core() -> None: + """_resolve_column_types lives in compiled core and maps FIELD_TYPE codes to Arrow tokens.""" + from sqlspec.adapters.mysqlconnector.core import _resolve_column_types + + description = [("id", 3), ("amount", 0), ("created", 7), ("unknown", 999)] + types = _resolve_column_types(description) + assert types == {"id": "int64", "amount": "decimal", "created": "timestamp"} + assert _resolve_column_types(None) is None + + +def test_stream_source_caches_deserializer() -> None: + """Stream sources cache the deserializer on initialization without per-chunk getattr.""" + from sqlspec.adapters.mysqlconnector.core import MysqlConnectorAsyncStreamSource, MysqlConnectorSyncStreamSource + + def custom_deserializer(v: Any) -> Any: + return v + + driver = MagicMock( + _json_deserializer=custom_deserializer, driver_features={"json_deserializer": custom_deserializer} + ) + + sync_source = MysqlConnectorSyncStreamSource(driver, "SELECT 1", (), 10, set()) + assert sync_source._deserializer is custom_deserializer + + async_source = MysqlConnectorAsyncStreamSource(driver, "SELECT 1", (), 10, set()) + assert async_source._deserializer is custom_deserializer diff --git a/tests/unit/adapters/test_mysqlconnector/test_driver.py b/tests/unit/adapters/test_mysqlconnector/test_driver.py index d3a5a7444..ff9bfc062 100644 --- a/tests/unit/adapters/test_mysqlconnector/test_driver.py +++ b/tests/unit/adapters/test_mysqlconnector/test_driver.py @@ -41,3 +41,93 @@ def _fail_commit() -> None: driver = MysqlConnectorSyncDriver(cast(Any, connection)) with pytest.raises(DeadlockError): driver.commit() + + +def test_mysqlconnector_pool_error_and_connection_error_mapping() -> None: + """PoolError maps to ConnectionTimeoutError and 2055 maps to DatabaseConnectionError.""" + from sqlspec.adapters.mysqlconnector.core import create_mapped_exception + from sqlspec.exceptions import ConnectionTimeoutError, DatabaseConnectionError + + class PoolError(Exception): + pass + + class InterfaceError(Exception): + errno = -1 + + assert isinstance(create_mapped_exception(PoolError("connection pool exhausted")), ConnectionTimeoutError) + assert isinstance(create_mapped_exception(InterfaceError("Connection disconnected")), DatabaseConnectionError) + assert isinstance( + create_mapped_exception(MysqlConnectorError(errno=2055, msg="Lost connection")), DatabaseConnectionError + ) + + +def test_mysqlconnector_sync_stream_source_in_active_except_block() -> None: + """MysqlConnectorSyncStreamSource safely handles execution inside an active except block.""" + from unittest.mock import MagicMock + + from sqlspec.adapters.mysqlconnector.core import MysqlConnectorSyncStreamSource + + cursor = MagicMock() + cursor.fetchmany.return_value = [{"id": 1}] + cursor.description = [("id", 3, None, None, None, None, None)] + connection = MagicMock() + connection.cursor.return_value = cursor + driver = MysqlConnectorSyncDriver(connection) + + try: + raise ValueError("simulated active exception") + except ValueError: + source = MysqlConnectorSyncStreamSource(driver, "SELECT 1", (), 10, set()) + source.start() + chunk = source.fetch_chunk() + assert len(chunk) == 1 + source.close() + + +def test_mysqlconnector_data_dictionary_cached_property() -> None: + """MysqlConnectorSyncDriver and MysqlConnectorAsyncDriver cache the data dictionary instance.""" + from unittest.mock import MagicMock + + sync_driver = MysqlConnectorSyncDriver(connection=MagicMock()) + assert sync_driver.data_dictionary is sync_driver.data_dictionary + + async_driver = MysqlConnectorAsyncDriver(connection=MagicMock()) + assert async_driver.data_dictionary is async_driver.data_dictionary + + +def test_with_cursor_avoids_dictionary_allocation() -> None: + """with_cursor passes None options when cursor_options is not configured.""" + from unittest.mock import MagicMock + + conn = MagicMock() + driver = MysqlConnectorSyncDriver(connection=conn) + cursor_cm = driver.with_cursor(conn) + assert cursor_cm.cursor_options is None + + +def test_mysqlconnector_table_and_column_not_found_map_to_operational_error() -> None: + """Errno 1146 and 1054 map to OperationalError rather than SQLParsingError.""" + from sqlspec.adapters.mysqlconnector.core import create_mapped_exception + from sqlspec.exceptions import OperationalError + + class _Err(Exception): + pass + + assert isinstance(create_mapped_exception(_Err(1146, "Table doesn't exist")), OperationalError) + assert isinstance(create_mapped_exception(_Err(1054, "Unknown column")), OperationalError) + + +def test_mysqlconnector_dispatch_execute_script_delegates_to_super() -> None: + """dispatch_execute_script delegates to base driver and emits telemetry.""" + from unittest.mock import MagicMock + + from sqlspec import SQL + + cursor = MagicMock(with_rows=False) + conn = MagicMock() + conn.cursor.return_value = cursor + driver = MysqlConnectorSyncDriver(connection=conn) + statement = SQL("SELECT 1; SELECT 2").as_script() + result = driver.dispatch_execute_script(cursor, statement) + assert result.statement_count == 2 + assert result.successful_statements == 2 diff --git a/tests/unit/adapters/test_mysqlconnector/test_litestar_store.py b/tests/unit/adapters/test_mysqlconnector/test_litestar_store.py index 45e49365a..10a0a1f66 100644 --- a/tests/unit/adapters/test_mysqlconnector/test_litestar_store.py +++ b/tests/unit/adapters/test_mysqlconnector/test_litestar_store.py @@ -48,3 +48,54 @@ def test_mysqlconnector_sync_store_honors_table_options() -> None: ddl = store._table_ddl() assert f"{DEFAULT_ENGINE_CLAUSE} COMMENT='litestar-session'" in ddl + + +def test_mysqlconnector_drop_table_sql_omits_drop_index() -> None: + """_drop_table_sql must emit only DROP TABLE IF EXISTS without preceding DROP INDEX.""" + async_store = MysqlConnectorAsyncStore(_mock_config()) + sync_store = MysqlConnectorSyncStore(_mock_config()) + + assert async_store._drop_table_sql() == ["DROP TABLE IF EXISTS litestar_session"] + assert sync_store._drop_table_sql() == ["DROP TABLE IF EXISTS litestar_session"] + + +async def test_mysqlconnector_set_uses_mariadb_compatible_upsert() -> None: + """set must use VALUES(col) in ON DUPLICATE KEY UPDATE for MySQL and MariaDB compatibility.""" + from unittest.mock import AsyncMock + + cursor = MagicMock() + cursor.execute = AsyncMock() + cursor.close = AsyncMock() + conn = MagicMock() + conn.cursor = AsyncMock(return_value=cursor) + conn.commit = AsyncMock() + conn.__aenter__ = AsyncMock(return_value=conn) + conn.__aexit__ = AsyncMock(return_value=None) + config = _mock_config() + config.provide_connection = lambda *_a, **_k: conn + + store = MysqlConnectorAsyncStore(config) + await store.set("session-1", b"payload", 300) + + assert cursor.execute.call_count == 1 + sql = cursor.execute.call_args_list[0].args[0] + assert "ON DUPLICATE KEY UPDATE" in sql + assert "data = VALUES(data)" in sql + assert "expires_at = VALUES(expires_at)" in sql + assert "AS new" not in sql + + +def test_mysqlconnector_sync_create_table_does_not_commit() -> None: + """Sync _create_table does not call driver.commit() after execute_script.""" + driver = MagicMock() + session = MagicMock() + session.__enter__ = MagicMock(return_value=driver) + session.__exit__ = MagicMock(return_value=None) + config = _mock_config() + config.provide_session = lambda *_a, **_k: session + + store = MysqlConnectorSyncStore(config) + store._create_table() + + assert driver.execute_script.called + assert not driver.commit.called diff --git a/tests/unit/adapters/test_oracledb/test_arrow_batches.py b/tests/unit/adapters/test_oracledb/test_arrow_batches.py index 73545af34..7386dc8b1 100644 --- a/tests/unit/adapters/test_oracledb/test_arrow_batches.py +++ b/tests/unit/adapters/test_oracledb/test_arrow_batches.py @@ -6,12 +6,13 @@ import pyarrow as pa import pytest -from sqlspec.adapters.oracledb.core import ( - build_arrow_fetch_kwargs, +from sqlspec.adapters.oracledb.core import build_arrow_fetch_kwargs +from sqlspec.adapters.oracledb.driver import OracleAsyncDriver, OracleSyncDriver +from sqlspec.utils.arrow_helpers import ( coerce_oracle_df_to_arrow_batches, coerce_oracle_df_to_arrow_table, + oracle_df_to_arrow_table, ) -from sqlspec.adapters.oracledb.driver import OracleAsyncDriver, OracleSyncDriver if TYPE_CHECKING: from sqlspec.adapters.oracledb._typing import OracleAsyncConnection, OracleSyncConnection @@ -188,4 +189,115 @@ async def fetch_df_all(self, **_kwargs: object) -> _StreamOnlyOracleDataFrame: ) res_table = await driver.select_to_arrow("SELECT id, value FROM example", return_format="table") assert res_table.get_data().to_pydict() == {"id": [1, 2], "value": [10, 20]} - assert stream_df.stream_calls == 1 + + +def test_oracle_df_to_arrow_table_lazy_pyarrow_resolution() -> None: + """oracle_df_to_arrow_table lazily resolves pyarrow when pa_module is None.""" + raw_table = pa.table({"ID": [1, 2], "VALUE": [10, 20]}) + stream_df = _StreamOnlyOracleDataFrame(raw_table) + table = oracle_df_to_arrow_table(stream_df, pa_module=None) + assert table.num_rows == 2 + + +def test_select_to_arrow_maps_native_database_error() -> None: + """select_to_arrow routes oracledb.DatabaseError through exception handler.""" + import oracledb + + from sqlspec.exceptions import SQLSpecError + + class _FailingOracleConn: + def fetch_df_all(self, **_kwargs: object) -> None: + raise oracledb.DatabaseError("ORA-00942: table or view does not exist") + + driver = OracleSyncDriver(cast("OracleSyncConnection", _FailingOracleConn())) + with pytest.raises(SQLSpecError): + driver.select_to_arrow("SELECT * FROM non_existent", return_format="table") + + +@pytest.mark.anyio +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.parametrize("return_format", ["batches", "reader"]) +async def test_arrow_batch_failure_preserves_database_error(is_async: bool, return_format: Any) -> None: + import oracledb + + from sqlspec.exceptions import SQLParsingError + + class _FailingConnection: + def fetch_df_all(self, **_kwargs: Any) -> None: + raise AssertionError("batch export should use fetch_df_batches") + + def fetch_df_batches(self, **_kwargs: Any) -> Iterator[pa.Table]: + raise oracledb.DatabaseError("ORA-00942: missing table") + yield pa.table({"ID": [1]}) + + class _FailingAsyncConnection: + async def fetch_df_all(self, **_kwargs: Any) -> None: + raise AssertionError("batch export should use fetch_df_batches") + + async def fetch_df_batches(self, **_kwargs: Any) -> AsyncIterator[pa.Table]: + raise oracledb.DatabaseError("ORA-00942: missing table") + yield pa.table({"ID": [1]}) + + with pytest.raises(SQLParsingError, match="ORA-00942") as exc_info: + if is_async: + await OracleAsyncDriver(cast("Any", _FailingAsyncConnection())).select_to_arrow( + "SELECT * FROM missing", return_format=return_format + ) + else: + OracleSyncDriver(cast("Any", _FailingConnection())).select_to_arrow( + "SELECT * FROM missing", return_format=return_format + ) + + assert isinstance(exc_info.value.__cause__, oracledb.DatabaseError) + + +def test_lazy_arrow_reader_maps_errors_after_the_first_batch() -> None: + import oracledb + + from sqlspec.exceptions import SQLParsingError + + class _FailingConnection: + def fetch_df_batches(self, **_kwargs: Any) -> Iterator[pa.Table]: + yield pa.table({"ID": [1]}) + raise oracledb.DatabaseError("ORA-00942: missing table") + + def fetch_df_all(self, **_kwargs: Any) -> None: + raise AssertionError("reader export should use fetch_df_batches") + + result = OracleSyncDriver(cast("Any", _FailingConnection())).select_to_arrow( + "SELECT * FROM example", return_format="reader" + ) + reader = result.get_data() + assert reader.read_next_batch().to_pydict() == {"ID": [1]} + with pytest.raises(SQLParsingError, match="ORA-00942") as exc_info: + reader.read_next_batch() + assert isinstance(exc_info.value.__cause__, oracledb.DatabaseError) + + +def test_sync_select_to_arrow_reader_yields_lazily() -> None: + """return_format=reader yields record batches lazily from the underlying stream.""" + yielded: list[int] = [] + + def _generate_tables() -> Iterator[pa.Table]: + yielded.append(1) + yield pa.table({"ID": [1, 2], "VALUE": [10, 20]}) + yielded.append(2) + yield pa.table({"ID": [3], "VALUE": [30]}) + + class _LazyBatchConnection: + def fetch_df_batches(self, **_kwargs: object) -> Iterator[pa.Table]: + return _generate_tables() + + driver = OracleSyncDriver( + cast("OracleSyncConnection", _LazyBatchConnection()), driver_features={"enable_lowercase_column_names": True} + ) + + result = driver.select_to_arrow("SELECT id, value FROM example", return_format="reader", batch_size=2) + reader = result.get_data() + assert yielded == [1] + first_batch = reader.read_next_batch() + assert first_batch.num_rows == 2 + assert yielded == [1] + second_batch = reader.read_next_batch() + assert second_batch.num_rows == 1 + assert yielded == [1, 2] diff --git a/tests/unit/adapters/test_oracledb/test_config.py b/tests/unit/adapters/test_oracledb/test_config.py index 18d25eb13..28af3fac4 100644 --- a/tests/unit/adapters/test_oracledb/test_config.py +++ b/tests/unit/adapters/test_oracledb/test_config.py @@ -1,5 +1,6 @@ """OracleDB configuration tests covering driver kwargs and typed options.""" +import asyncio import io import os import ssl @@ -9,6 +10,7 @@ from inspect import isawaitable from pathlib import Path from ssl import TLSVersion +from types import SimpleNamespace from typing import Any, NotRequired, cast, get_args, get_origin, get_type_hints from unittest.mock import AsyncMock, MagicMock, Mock, call @@ -67,6 +69,11 @@ def _stub_sync_connection_setup(monkeypatch: pytest.MonkeyPatch, calls: list[str monkeypatch.setattr(oracle_config_module, "register_uuid_handlers", lambda _connection: calls.append("uuid")) +async def _create_standalone_oracle_connection(config: OracleSyncConfig | OracleAsyncConfig) -> object: + result = config.create_connection() + return await result if isawaitable(result) else result + + def test_oracle_connection_params_expose_current_driver_options() -> None: """Connection params should mirror current python-oracledb connection knobs.""" annotations = _oracle_config_hints(OracleConnectionParams) @@ -347,17 +354,187 @@ async def test_async_pool_rejects_thick_mode(options: dict[str, Any], monkeypatc create_pool.assert_not_called() -async def test_async_soda_option_does_not_initialize_thick_mode(monkeypatch: pytest.MonkeyPatch) -> None: +async def test_async_soda_option_rejects_thick_mode_option(monkeypatch: pytest.MonkeyPatch) -> None: initialize = Mock() create_pool = Mock() monkeypatch.setattr(oracle_config_module.oracledb, "init_oracle_client", initialize) monkeypatch.setattr(oracle_config_module.oracledb, "create_pool_async", create_pool) config = OracleAsyncConfig(connection_config={"soda_metadata_cache": True}) - assert await config._create_pool() is create_pool.return_value + with pytest.raises(ImproperConfigurationError, match="only supports Thin mode"): + await config._create_pool() initialize.assert_not_called() - assert create_pool.call_args.kwargs["soda_metadata_cache"] is True + create_pool.assert_not_called() + + +def test_sync_create_connection_unpooled(monkeypatch: pytest.MonkeyPatch) -> None: + """create_connection creates an unpooled standalone connection without instantiating a pool.""" + connect_mock = Mock() + create_pool_mock = Mock() + monkeypatch.setattr(oracle_config_module.oracledb, "connect", connect_mock) + monkeypatch.setattr(oracle_config_module.oracledb, "create_pool", create_pool_mock) + config = OracleSyncConfig(connection_config={"user": "scott", "password": "tiger", "dsn": "localhost/xe"}) + conn = config.create_connection() + assert conn is connect_mock.return_value + connect_mock.return_value.close.assert_not_called() + create_pool_mock.assert_not_called() + assert config.connection_instance is None + + +@pytest.mark.anyio +async def test_async_create_connection_unpooled(monkeypatch: pytest.MonkeyPatch) -> None: + """create_connection creates an unpooled standalone async connection without instantiating a pool.""" + connect_mock = AsyncMock() + create_pool_mock = AsyncMock() + monkeypatch.setattr(oracle_config_module.oracledb, "connect_async", connect_mock) + monkeypatch.setattr(oracle_config_module.oracledb, "create_pool_async", create_pool_mock) + config = OracleAsyncConfig(connection_config={"user": "scott", "password": "tiger", "dsn": "localhost/xe"}) + conn = await config.create_connection() + assert conn is connect_mock.return_value + connect_mock.return_value.close.assert_not_called() + create_pool_mock.assert_not_called() + assert config.connection_instance is None + + +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.parametrize("hook_kind", ["on_connection_create", "session_callback"]) +@pytest.mark.parametrize("cancelled", [False, True]) +@pytest.mark.parametrize("cleanup_failure", [None, OSError("close failed"), asyncio.CancelledError("close cancelled")]) +async def test_standalone_connection_setup_failure_closes_owned_connection( + is_async: bool, + hook_kind: str, + cancelled: bool, + cleanup_failure: BaseException | None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Configured setup failures close new SDK connections without replacing the primary error.""" + failure = asyncio.CancelledError("setup cancelled") if cancelled else RuntimeError("setup failed") + close = AsyncMock(side_effect=cleanup_failure) if is_async else Mock(side_effect=cleanup_failure) + connection = SimpleNamespace(version="23.5.0.0.0", inputtypehandler=None, outputtypehandler=None, close=close) + connect = AsyncMock(return_value=connection) if is_async else Mock(return_value=connection) + monkeypatch.setattr(oracle_config_module.oracledb, "connect_async" if is_async else "connect", connect) + + async def fail_async_setup() -> None: + await asyncio.sleep(0) + raise failure + + def setup_hook(acquired: object, tag: str) -> Awaitable[None]: + assert acquired is connection + assert tag == "" + assert callable(connection.inputtypehandler) + assert callable(connection.outputtypehandler) + if is_async: + return fail_async_setup() + raise failure + + connection_options: dict[str, Any] = {"dsn": "localhost/xe"} + driver_features: dict[str, Any] = {} + if hook_kind == "session_callback": + connection_options[hook_kind] = setup_hook + else: + driver_features[hook_kind] = setup_hook + config_type = OracleAsyncConfig if is_async else OracleSyncConfig + config = config_type(connection_config=connection_options, driver_features=driver_features) + + with pytest.raises(type(failure)) as caught: + await _create_standalone_oracle_connection(config) + + assert caught.value is failure + if is_async: + cast("AsyncMock", close).assert_awaited_once_with() + else: + close.assert_called_once_with() + assert config.connection_instance is None + + +@pytest.mark.anyio +@pytest.mark.parametrize("is_async", [False, True]) +async def test_standalone_connections_filter_pool_options(is_async: bool, monkeypatch: pytest.MonkeyPatch) -> None: + """A pool-configured config can also create a standalone connection.""" + connection_type = ( + oracle_config_module.oracledb.AsyncConnection if is_async else oracle_config_module.oracledb.Connection + ) + connect = AsyncMock() if is_async else Mock() + monkeypatch.setattr(oracle_config_module.oracledb, "connect_async" if is_async else "connect", connect) + config_type = OracleAsyncConfig if is_async else OracleSyncConfig + config = config_type( + connection_config={ + "user": "scott", + "password": "tiger", + "dsn": "localhost/xe", + "min": 1, + "max": 4, + "increment": 1, + "getmode": PoolGetMode.WAIT, + "ping_interval": 60, + "ping_timeout": 5000, + "max_sessions_per_shard": 2, + "pool_class": oracle_config_module.oracledb.ConnectionPool, + "connectiontype": connection_type, + } + ) + + result = config.create_connection() + if isawaitable(result): + result = await result + + assert result is connect.return_value + connect.assert_called_once_with(user="scott", password="tiger", dsn="localhost/xe", conn_class=connection_type) + + +@pytest.mark.parametrize("options", [{"thick_mode": True}, {"lib_dir": "/oracle/lib"}, {"soda_metadata_cache": True}]) +def test_standalone_sync_connection_initializes_thick_mode( + options: dict[str, Any], monkeypatch: pytest.MonkeyPatch +) -> None: + calls: list[str] = [] + + def fake_connect(**_kwargs: object) -> _StubConnection: + calls.append("connect") + return _StubConnection() + + initialize = Mock(side_effect=lambda **_kwargs: calls.append("initialize")) + connect = Mock(side_effect=fake_connect) + monkeypatch.setattr(oracle_config_module.oracledb, "is_thin_mode", lambda: True) + monkeypatch.setattr(oracle_config_module.oracledb, "init_oracle_client", initialize) + monkeypatch.setattr(oracle_config_module.oracledb, "connect", connect) + config = OracleSyncConfig(connection_config=options) + + config.create_connection() + + assert calls == ["initialize", "connect"] + initialize.assert_called_once_with(**({"lib_dir": options["lib_dir"]} if "lib_dir" in options else {})) + + +@pytest.mark.anyio +@pytest.mark.parametrize("options", [{"thick_mode": True}, {"lib_dir": "/oracle/lib"}, {"soda_metadata_cache": True}]) +async def test_standalone_async_connection_rejects_thick_mode( + options: dict[str, Any], monkeypatch: pytest.MonkeyPatch +) -> None: + connect = AsyncMock() + monkeypatch.setattr(oracle_config_module.oracledb, "connect_async", connect) + config = OracleAsyncConfig(connection_config=options) + + with pytest.raises(ImproperConfigurationError, match="only supports Thin mode"): + await config.create_connection() + + connect.assert_not_called() + + +def test_derive_pool_produces_identical_configurations_across_sync_and_async() -> None: + """derive_pool produces identical connection configurations across sync and async configs.""" + common_cfg = {"user": "scott", "password": "tiger", "dsn": "localhost/xe", "min": 2, "max": 10} + common_features = {"enable_lowercase_column_names": True, "fetch_lobs": False} + sync_config = OracleSyncConfig(connection_config=common_cfg, driver_features=common_features) + async_config = OracleAsyncConfig(connection_config=common_cfg, driver_features=common_features) + + derived_sync = sync_config.derive_pool(connection_config_overrides={"max": 20}, fetch_lobs=True) + derived_async = async_config.derive_pool(connection_config_overrides={"max": 20}, fetch_lobs=True) + + assert derived_sync.connection_config == derived_async.connection_config + sync_features = {k: v for k, v in derived_sync.driver_features.items() if k != "_statement_cache_holder"} + async_features = {k: v for k, v in derived_async.driver_features.items() if k != "_statement_cache_holder"} + assert sync_features == async_features def test_oracle_config_wallet_zip_extracted_and_cleaned_up_on_close_pool() -> None: diff --git a/tests/unit/adapters/test_oracledb/test_core_driver_features.py b/tests/unit/adapters/test_oracledb/test_core_driver_features.py index 58cee2653..8adb5f874 100644 --- a/tests/unit/adapters/test_oracledb/test_core_driver_features.py +++ b/tests/unit/adapters/test_oracledb/test_core_driver_features.py @@ -1,4 +1,3 @@ -# pyright: reportArgumentType=false """Unit tests for OracleDB ``apply_driver_features`` defaults.""" import inspect @@ -506,3 +505,55 @@ def test_apply_wallet_ssl_context_noop_when_no_wallet() -> None: conn_config: dict[str, Any] = {"user": "app", "password": "secret"} apply_wallet_ssl_context(conn_config) assert "ssl_context" not in conn_config + + +def test_ora_08177_maps_to_serialization_conflict_error() -> None: + """ORA-08177 maps to SerializationConflictError.""" + import oracledb + + from sqlspec.adapters.oracledb.core import create_mapped_exception + from sqlspec.exceptions import SerializationConflictError + + error = oracledb.DatabaseError("ORA-08177: can't serialize access for this transaction") + mapped = create_mapped_exception(error) + assert isinstance(mapped, SerializationConflictError) + + +def test_ora_42699_maps_to_serialization_conflict_error() -> None: + """ORA-42699 maps to SerializationConflictError.""" + import oracledb + + from sqlspec.adapters.oracledb.core import create_mapped_exception + from sqlspec.exceptions import SerializationConflictError + + error = oracledb.DatabaseError("ORA-42699: transaction isolation violation") + mapped = create_mapped_exception(error) + assert isinstance(mapped, SerializationConflictError) + + +def test_dpy_4024_maps_to_query_timeout_error() -> None: + """DPY-4024 maps to QueryTimeoutError.""" + import oracledb + + from sqlspec.adapters.oracledb.core import create_mapped_exception + from sqlspec.exceptions import QueryTimeoutError + + error = oracledb.OperationalError("DPY-4024: call timeout expired") + mapped = create_mapped_exception(error) + assert isinstance(mapped, QueryTimeoutError) + + +def test_ora_data_and_timeout_error_mappings() -> None: + """Verify ORA-40441, ORA-51803, ORA-01476 map to DataError, ORA-03136 maps to ConnectionTimeoutError.""" + import oracledb + + from sqlspec.adapters.oracledb.core import create_mapped_exception + from sqlspec.exceptions import ConnectionTimeoutError, DataError + + assert isinstance(create_mapped_exception(oracledb.DatabaseError("ORA-40441: JSON syntax error")), DataError) + assert isinstance(create_mapped_exception(oracledb.DatabaseError("ORA-51803: vector operation error")), DataError) + assert isinstance(create_mapped_exception(oracledb.DatabaseError("ORA-01476: divisor is equal to zero")), DataError) + assert isinstance( + create_mapped_exception(oracledb.DatabaseError("ORA-03136: inbound connection timed out")), + ConnectionTimeoutError, + ) diff --git a/tests/unit/adapters/test_oracledb/test_core_row_materialization.py b/tests/unit/adapters/test_oracledb/test_core_row_materialization.py index f0059597a..164675889 100644 --- a/tests/unit/adapters/test_oracledb/test_core_row_materialization.py +++ b/tests/unit/adapters/test_oracledb/test_core_row_materialization.py @@ -131,12 +131,18 @@ def test_resolve_row_metadata_single_pass_unknown_type_code_shape_is_conservativ def test_resolve_row_metadata_cache_contract_still_holds_after_single_pass_rewrite() -> None: description = [("ID", _TypeCode("DB_TYPE_NUMBER"))] - cache: dict[int, tuple[object, list[str], bool]] = {} + cache: dict[object, tuple[object, list[str], bool]] = {} (first_names, first_requires_lob) = resolve_row_metadata(description, {}, cache) (second_names, second_requires_lob) = resolve_row_metadata(description, {}, cache) assert second_names is first_names assert first_requires_lob is second_requires_lob - assert cache[id(description)][0] is description + cache_key = (("ID", "DB_TYPE_NUMBER"),) + assert cache_key in cache + assert cache[cache_key][0] is description + description2 = [("ID", _TypeCode("DB_TYPE_NUMBER"))] + (third_names, third_requires_lob) = resolve_row_metadata(description2, {}, cache) + assert third_names is first_names + assert third_requires_lob is first_requires_lob @pytest.mark.parametrize("asynchronous", [False, True]) diff --git a/tests/unit/adapters/test_oracledb/test_data_dictionary_oracledb.py b/tests/unit/adapters/test_oracledb/test_data_dictionary_oracledb.py index ed73f3fd3..0396e2cf6 100644 --- a/tests/unit/adapters/test_oracledb/test_data_dictionary_oracledb.py +++ b/tests/unit/adapters/test_oracledb/test_data_dictionary_oracledb.py @@ -5,6 +5,7 @@ import pytest from sqlspec.adapters.oracledb.data_dictionary import OracledbAsyncDataDictionary, OracledbSyncDataDictionary +from sqlspec.adapters.oracledb.driver import OracleAsyncDriver, OracleSyncDriver from sqlspec.data_dictionary import ( DDLResult, MetadataFidelity, @@ -32,7 +33,7 @@ ) if TYPE_CHECKING: - from sqlspec.adapters.oracledb.driver import OracleAsyncDriver, OracleSyncDriver + from sqlspec.adapters.oracledb._typing import OracleAsyncConnection, OracleSyncConnection ORACLE_AI_COMPONENT_ROW = { "product": "Oracle AI Database 26ai Free", @@ -220,6 +221,8 @@ def test_oracle_dialect_helpers_resolve_shared_version_rules() -> None: config = get_dialect_config("oracle") assert extract_oracle_version_value({"VERSION": "23.26.0.0.0"}) == "23.26.0.0.0" + assert extract_oracle_version_value({"VERSION": "23.0.0.0.0", "VERSION_FULL": "23.26.0.0.0"}) == "23.26.0.0.0" + assert extract_oracle_version_value({"version": "19.0.0.0.0", "version_full": None}) == "19.0.0.0.0" assert extract_oracle_version_value(("19.20.0.0.0",)) == "19.20.0.0.0" assert parse_oracle_version_components("23.26.0.0.0") == (23, 26, 0) assert parse_oracle_compatible_major("23.20.0.0.0") == 23 @@ -298,7 +301,7 @@ def test_oracle_domain_query_packs_cover_c5_sources() -> None: "tables": ("by_schema", "ALL_TABLES"), "columns": ("by_schema", "ALL_TAB_COLS"), "constraints": ("by_schema", "ALL_CONSTRAINTS"), - "indexes": ("by_schema", "ALL_IND_EXPRESSIONS"), + "indexes": ("by_schema", "ALL_INDEXES"), "views": ("by_schema", "ALL_VIEWS"), "materialized_views": ("by_schema", "ALL_MVIEWS"), "sequences": ("by_schema", "ALL_SEQUENCES"), @@ -526,3 +529,281 @@ async def test_oracle_26ai_async_data_dictionary_introspection_methods() -> None dviews = await data_dictionary.get_json_duality_views(async_driver, schema="app") assert isinstance(dviews, list) + + +def test_oracle_indexes_sql_does_not_listagg_column_expression() -> None: + """Oracle indexes SQL text must not aggregate column_expression via LISTAGG.""" + data_dictionary = OracledbSyncDataDictionary() + by_schema_sql = str(data_dictionary.get_query("indexes", "by_schema")).upper() + by_table_sql = str(data_dictionary.get_query("indexes", "by_table")).upper() + + assert "LISTAGG(IE.COLUMN_EXPRESSION" not in by_schema_sql + assert "COLUMN_EXPRESSION" not in by_schema_sql + assert "LISTAGG(IE.COLUMN_EXPRESSION" not in by_table_sql + + +def test_oracle_sync_get_indexes_merges_expressions() -> None: + """get_indexes must query expressions and merge them into index records.""" + data_dictionary = OracledbSyncDataDictionary() + + class _StubSyncDriver: + def __init__(self) -> None: + self.queries: list[str] = [] + + def select(self, query: Any, **kwargs: Any) -> list[dict[str, Any]]: + sql_text = str(query) + self.queries.append(sql_text) + if "all_ind_expressions" in sql_text: + return [ + { + "schema_name": "HR", + "table_name": "EMPLOYEES", + "index_name": "IDX_EMP_UPPER_EMAIL", + "column_position": 1, + "column_expression": "UPPER(EMAIL)", + } + ] + return [ + { + "schema_name": "HR", + "table_name": "EMPLOYEES", + "index_name": "IDX_EMP_UPPER_EMAIL", + "index_type": "NORMAL", + "uniqueness": "NONUNIQUE", + "columns": None, + } + ] + + driver = cast("OracleSyncDriver", _StubSyncDriver()) + indexes = data_dictionary.get_indexes(driver, schema="HR") + assert len(indexes) == 1 + assert indexes[0]["index_name"] == "IDX_EMP_UPPER_EMAIL" + assert indexes[0]["expressions"] == "UPPER(EMAIL)" + assert indexes[0]["expression"] == "UPPER(EMAIL)" + + +@pytest.mark.anyio +async def test_oracle_async_get_indexes_merges_expressions() -> None: + """Async get_indexes must query expressions and merge them into index records.""" + data_dictionary = OracledbAsyncDataDictionary() + + class _StubAsyncDriver: + def __init__(self) -> None: + self.queries: list[str] = [] + + async def select(self, query: Any, **kwargs: Any) -> list[dict[str, Any]]: + sql_text = str(query) + self.queries.append(sql_text) + if "all_ind_expressions" in sql_text: + return [ + { + "schema_name": "HR", + "table_name": "EMPLOYEES", + "index_name": "IDX_EMP_UPPER_EMAIL", + "column_position": 1, + "column_expression": "UPPER(EMAIL)", + } + ] + return [ + { + "schema_name": "HR", + "table_name": "EMPLOYEES", + "index_name": "IDX_EMP_UPPER_EMAIL", + "index_type": "NORMAL", + "uniqueness": "NONUNIQUE", + "columns": None, + } + ] + + driver = cast("OracleAsyncDriver", _StubAsyncDriver()) + indexes = await data_dictionary.get_indexes(driver, schema="HR") + assert len(indexes) == 1 + assert indexes[0]["index_name"] == "IDX_EMP_UPPER_EMAIL" + assert indexes[0]["expressions"] == "UPPER(EMAIL)" + assert indexes[0]["expression"] == "UPPER(EMAIL)" + + +@pytest.mark.anyio +async def test_oracle_async_get_native_object_ddl_awaits_async_lob() -> None: + """Async _get_native_object_ddl awaits an AsyncLOB read() instead of stringifying the coroutine.""" + data_dictionary = OracledbAsyncDataDictionary() + + class _AsyncLOB: + async def read(self) -> str: + return "CREATE TABLE FOO (ID NUMBER)" + + class _AsyncDriverStub: + async def select_value(self, query: Any, **kwargs: Any) -> Any: + return _AsyncLOB() + + driver = cast("OracleAsyncDriver", _AsyncDriverStub()) + result = await data_dictionary._get_native_object_ddl( + driver, "FOO", schema="HR", object_type="TABLE", include_dependencies=False, redact=False + ) + assert result is not None + assert result.ddl == "CREATE TABLE FOO (ID NUMBER)" + assert " None: + """get_tables issues a single catalog query (tables.by_schema) without a redundant names query.""" + data_dictionary = OracledbSyncDataDictionary() + queries: list[str] = [] + + class _StubSyncDriver: + def select(self, query: Any, **kwargs: Any) -> list[dict[str, Any]]: + queries.append(str(query)) + return [{"table_name": "T1"}] + + driver = cast("OracleSyncDriver", _StubSyncDriver()) + tables = data_dictionary.get_tables(driver, schema="HR") + assert len(tables) == 1 + assert len(queries) == 1 + assert "names_by_schema" not in queries[0] + + +class _CatalogCursor: + def __init__(self, connection: "_CatalogConnection") -> None: + self.connection = connection + self.description: list[tuple[str, object]] = [] + self.rows: list[tuple[object, ...]] = [] + self.rowcount = 0 + + def execute(self, statement: str, parameters: object = None, **kwargs: object) -> None: + import oracledb + + self.connection.statements.append(statement) + if "product_component_version" in statement: + columns = ["PRODUCT", "VERSION", "STATUS"] + self.rows = [("Oracle Database", self.connection.catalog_version, "Production")] + if int(self.connection.catalog_version.split(".", 1)[0]) >= 18: + columns.insert(2, "VERSION_FULL") + self.rows = [ + ("Oracle Database", self.connection.catalog_version, self.connection.full_version, "Production") + ] + elif "v$parameter" in statement: + columns = ["COMPATIBLE"] + self.rows = [("23.4.0.0.0",)] + elif "sys_context" in statement: + columns = ["SERVICE"] + self.rows = [(None,)] + elif "all_ind_expressions" in statement: + columns = ["SCHEMA_NAME", "TABLE_NAME", "INDEX_NAME", "COLUMN_EXPRESSION"] + self.rows = [("HR", "EMPLOYEES", "IDX_EMAIL", 'UPPER("EMAIL")')] + elif "all_indexes" in statement: + columns = ["INDEX_NAME", "TABLE_NAME", "IS_UNIQUE", "IS_PRIMARY", "COLUMNS"] + self.rows = [("IDX_EMAIL", "EMPLOYEES", 0, 0, "SYS_NC00003$")] + if "i.owner AS schema_name" in statement: + columns.insert(0, "SCHEMA_NAME") + self.rows = [("HR", *row) for row in self.rows] + else: + raise ValueError(statement) + self.description = [(name, oracledb.DB_TYPE_VARCHAR) for name in columns] + self.rowcount = len(self.rows) + + def fetchall(self) -> list[tuple[object, ...]]: + rows, self.rows = self.rows, [] + return rows + + def fetchone(self) -> tuple[object, ...] | None: + return self.rows.pop(0) if self.rows else None + + def close(self) -> None: + return None + + +class _AsyncCatalogCursor: + def __init__(self, connection: "_CatalogConnection") -> None: + self.cursor = _CatalogCursor(connection) + + @property + def description(self) -> list[tuple[str, object]]: + return self.cursor.description + + @property + def rowcount(self) -> int: + return self.cursor.rowcount + + async def execute(self, statement: str, parameters: object = None, **kwargs: object) -> None: + self.cursor.execute(statement, parameters, **kwargs) + + async def fetchall(self) -> list[tuple[object, ...]]: + return self.cursor.fetchall() + + async def fetchone(self) -> tuple[object, ...] | None: + return self.cursor.fetchone() + + def close(self) -> None: + self.cursor.close() + + +class _CatalogConnection: + def __init__( + self, + *, + is_async: bool, + version: str = "23.0.0.0.0", + full_version: str = "23.26.0.0.0", + sdk_version_present: bool = True, + ) -> None: + self.is_async = is_async + self.version = version if sdk_version_present else None + self.catalog_version = version + self.full_version = full_version + self.statements: list[str] = [] + + def cursor(self) -> _CatalogCursor | _AsyncCatalogCursor: + return _AsyncCatalogCursor(self) if self.is_async else _CatalogCursor(self) + + +@pytest.mark.anyio +@pytest.mark.parametrize("is_async", [False, True]) +async def test_oracle_default_schema_index_expression_uses_native_catalog_owner(is_async: bool) -> None: + connection = _CatalogConnection(is_async=is_async) + if is_async: + driver = OracleAsyncDriver( + cast("OracleAsyncConnection", connection), driver_features={"enable_lowercase_column_names": True} + ) + indexes = await driver.data_dictionary.get_indexes(driver, table="employees") + else: + sync_driver = OracleSyncDriver( + cast("OracleSyncConnection", connection), driver_features={"enable_lowercase_column_names": True} + ) + indexes = sync_driver.data_dictionary.get_indexes(sync_driver, table="employees") + assert len(indexes) == 1 + assert indexes[0]["index_name"] == "IDX_EMAIL" + assert indexes[0]["expression"] == 'UPPER("EMAIL")' + + +@pytest.mark.anyio +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.parametrize( + ("major", "sdk_version_present", "lowercase_columns"), + [(12, True, True), (18, True, False), (23, True, True), (23, False, False)], +) +async def test_oracle_version_preserves_release_update_and_legacy_query( + is_async: bool, major: int, sdk_version_present: bool, lowercase_columns: bool +) -> None: + base_version = f"{major}.0.0.0.0" + full_version = f"{major}.26.0.0.0" + connection = _CatalogConnection( + is_async=is_async, version=base_version, full_version=full_version, sdk_version_present=sdk_version_present + ) + if is_async: + driver = OracleAsyncDriver( + cast("OracleAsyncConnection", connection), + driver_features={"enable_lowercase_column_names": lowercase_columns}, + ) + version = await driver.data_dictionary.get_version(driver) + cached_version = await driver.data_dictionary.get_version(driver) + else: + sync_driver = OracleSyncDriver( + cast("OracleSyncConnection", connection), + driver_features={"enable_lowercase_column_names": lowercase_columns}, + ) + version = sync_driver.data_dictionary.get_version(sync_driver) + cached_version = sync_driver.data_dictionary.get_version(sync_driver) + assert version == VersionInfo(major, 26 if major >= 18 else 0, 0) + assert cached_version is version + version_statements = [query for query in connection.statements if "product_component_version" in query] + assert len(version_statements) == 1 diff --git a/tests/unit/adapters/test_oracledb/test_direct_path_load.py b/tests/unit/adapters/test_oracledb/test_direct_path_load.py index 16fe1c2a5..82e951d84 100644 --- a/tests/unit/adapters/test_oracledb/test_direct_path_load.py +++ b/tests/unit/adapters/test_oracledb/test_direct_path_load.py @@ -329,3 +329,24 @@ def test_convertible_scalar_columns_keep_the_native_path(arrow_type: "pa.DataTyp assert len(conn.dpl_calls) == 1 assert conn._cursor.executemany_calls == [] + + +def test_direct_path_load_failure_preserves_cause() -> None: + """direct_path_load failure preserves the underlying cause on the raised exception.""" + import oracledb + + from sqlspec.exceptions import SQLSpecError + + native_err = oracledb.DatabaseError("ORA-00001: unique constraint violated") + + class _FailingDPLConn(_DPLConnection): + def direct_path_load(self, **kwargs: Any) -> None: + raise native_err + + conn = _FailingDPLConn(thin=True, username="SCOTT") + driver = OracleSyncDriver(cast("Any", conn), driver_features={"storage_capabilities": _CAPS}) + + with pytest.raises(SQLSpecError) as exc_info: + driver.load_from_arrow("MYTAB", _arrow()) + + assert exc_info.value.__cause__ is native_err diff --git a/tests/unit/adapters/test_oracledb/test_driver.py b/tests/unit/adapters/test_oracledb/test_driver.py index 24c928eb3..bc31c4e8f 100644 --- a/tests/unit/adapters/test_oracledb/test_driver.py +++ b/tests/unit/adapters/test_oracledb/test_driver.py @@ -8,11 +8,13 @@ import pytest from sqlspec.adapters.oracledb import build_dbms_cloud_copy_data_sql as exported_build_dbms_cloud_copy_data_sql +from sqlspec.adapters.oracledb._typing import OracleSyncConnection from sqlspec.adapters.oracledb.core import ( ORACLE_ACTION_MAX_BYTES, ORACLE_CLIENT_IDENTIFIER_MAX_BYTES, ORACLE_ECID_MAX_BYTES, ORACLE_MODULE_MAX_BYTES, + OracleSyncStreamSource, apply_driver_features, apply_oracle_session_trace_attributes, build_dbms_cloud_copy_data_sql, @@ -537,3 +539,133 @@ def close(self) -> None: assert len(async_executed) == 1 assert "DBMS_CLOUD.COPY_DATA" not in async_executed[0][0] assert "INSERT INTO" in async_executed[0][0] + + +def test_oracle_sync_stream_source_runs_with_exception_handler() -> None: + """Sync stream source start and fetch_chunk route execution through _run_with_exception_handler.""" + cursor = MagicMock() + cursor.description = [("id", int, None, None, None, None, None)] + cursor.fetchmany.return_value = [(1,)] + conn = MagicMock() + conn.cursor.return_value = cursor + driver = OracleSyncDriver(cast("OracleSyncConnection", conn)) + + source = OracleSyncStreamSource(driver, "SELECT 1 FROM dual", None, 100) + source.start() + rows = source.fetch_chunk() + assert len(rows) == 1 + assert rows[0] == {"id": 1} + + +def test_oracle_statement_cache_hit_executes_with_fetch_kwargs_and_coercion() -> None: + """Fast-path statement cache hit passes fetch kwargs and coerces parameters.""" + cursor = MagicMock() + cursor.description = [("content", "CLOB", None, None, None, None, None)] + cursor.fetchall.return_value = [("clob text",)] + conn = MagicMock() + conn.cursor.return_value = cursor + driver = OracleSyncDriver(cast("OracleSyncConnection", conn)) + + res1 = driver.execute("SELECT content FROM articles WHERE id = :id", {"id": "1"}) + assert res1.get_data() == [{"content": "clob text"}] + + res2 = driver.execute("SELECT content FROM articles WHERE id = :id", {"id": "1"}) + assert res2.get_data() == [{"content": "clob text"}] + + +def test_oracle_dispatch_execute_script_ignores_colons_in_literals_and_slices_positional() -> None: + """dispatch_execute_script ignores colons in literals and slices positional parameters.""" + from sqlspec import SQL + + executed_statements: list[tuple[str, Any]] = [] + + class _MockCursor: + def execute(self, stmt: str, params: Any = None) -> None: + executed_statements.append((stmt, params)) + + conn = MagicMock() + cursor = _MockCursor() + conn.cursor.return_value = cursor + driver = OracleSyncDriver(cast("OracleSyncConnection", conn), default_statement_config) + + script_sql = "INSERT INTO t (url) VALUES ('http://example.com:8080'); INSERT INTO t (id, name) VALUES (:1, :2)" + stmt = SQL(script_sql, [101, "Alice"], statement_config=default_statement_config) + res = driver.dispatch_execute_script(cursor, stmt) + assert res.statement_count == 2 + assert len(executed_statements) == 2 + assert executed_statements[0][1] == [] + assert executed_statements[1][1] == [101, "Alice"] + + +def test_parse_message_string_or_bytes_payload() -> None: + """_parse_message parses JSON string or bytes into a structured EventMessage.""" + import json + + from sqlspec.adapters.oracledb.events.backend import _parse_message + + raw_dict = { + "event_id": "ev-123", + "channel": "orders", + "payload": {"order_id": 42}, + "metadata": {"source": "unit"}, + "published_at": "2026-05-10T12:00:00Z", + } + json_str = json.dumps(raw_dict) + msg_from_str = _parse_message("orders", json_str) + assert msg_from_str.event_id == "ev-123" + assert msg_from_str.channel == "orders" + assert msg_from_str.payload == {"order_id": 42} + assert msg_from_str.metadata == {"source": "unit"} + + json_bytes = json_str.encode("utf-8") + msg_from_bytes = _parse_message("orders", json_bytes) + assert msg_from_bytes.event_id == "ev-123" + assert msg_from_bytes.payload == {"order_id": 42} + + +def test_oracle_aq_commit_at_dequeue_and_noop_ack_nack() -> None: + """Oracle AQ hub commits at dequeue for at-most-once delivery; ack and nack are no-ops.""" + import json + + from sqlspec.adapters.oracledb.events._hub import OracleSyncAQHub + from sqlspec.adapters.oracledb.events.backend import OracleSyncAQEventBackend + + class _MockOracleSyncConfig: + __module__ = "sqlspec.adapters.oracledb.config" + is_async = False + driver_features: dict[str, object] = {} + + def get_observability_runtime(self) -> object: + return MagicMock() + + driver = MagicMock() + session_cm = MagicMock() + session_cm.__enter__.return_value = driver + + config = _MockOracleSyncConfig() + config.provide_session = lambda: session_cm # type: ignore[attr-defined] + + msg = MagicMock() + msg.payload = json.dumps({"event_id": "ev-pinned", "payload": {"foo": "bar"}}) + queue = MagicMock() + queue.deqone.return_value = msg + + hub = OracleSyncAQHub( + cast("Any", config), + queue_name_template="TEST_QUEUE", + visibility=None, + default_visibility=None, + wait_ceiling=5, + backend_name="aq", + ) + hub._session_driver = driver + hub._queues["test_chan"] = queue + + payload = hub.dequeue("test_chan", poll_interval=0.1) + assert payload is not None + driver.commit.assert_called_once() + + backend = OracleSyncAQEventBackend(cast("Any", config)) + backend._hub = hub + backend.ack("ev-pinned") + backend.nack("ev-pinned") diff --git a/tests/unit/adapters/test_oracledb/test_fetch_tuning.py b/tests/unit/adapters/test_oracledb/test_fetch_tuning.py index 324654c84..38f7344a4 100644 --- a/tests/unit/adapters/test_oracledb/test_fetch_tuning.py +++ b/tests/unit/adapters/test_oracledb/test_fetch_tuning.py @@ -167,3 +167,34 @@ async def test_oracle_async_stream_fetch_lobs_call_option_overrides_default() -> rows: list[dict[str, Any]] = [row async for row in stream] assert rows == [] assert raw_cursor.execute_kwargs["fetch_lobs"] is True + + +def test_resolve_rowcount_negative_rowcount_returns_zero() -> None: + """resolve_rowcount returns 0 when cursor.rowcount is negative.""" + from unittest.mock import Mock + + from sqlspec.adapters.oracledb.core import resolve_rowcount + + cursor = Mock() + cursor.rowcount = -1 + assert resolve_rowcount(cursor) == 0 + + +def test_build_fetch_kwargs_returns_cached_instance_when_features_unchanged() -> None: + """build_fetch_kwargs returns the identical cached dict when features are unchanged.""" + features = {"fetch_lobs": False, "fetch_decimals": True} + res1 = build_fetch_kwargs(features) + res2 = build_fetch_kwargs(features) + assert res1 is res2 + + +def test_fetch_kwargs_precomputed_on_driver_instance() -> None: + """_fetch_kwargs is precomputed and cached on the driver instance.""" + raw_cursor = _RawCursor() + driver = OracleSyncDriver( + cast("OracleSyncConnection", _SyncConnection(raw_cursor)), + driver_features={"fetch_lobs": False, "fetch_decimals": True}, + ) + assert hasattr(driver, "_fetch_kwargs") + assert driver._fetch_kwargs == {"fetch_lobs": False, "fetch_decimals": True} + assert driver._fetch_kwargs is build_fetch_kwargs({"fetch_lobs": False, "fetch_decimals": True}) diff --git a/tests/unit/adapters/test_oracledb/test_json_handlers.py b/tests/unit/adapters/test_oracledb/test_json_handlers.py index e8d81c651..e91ccf8db 100644 --- a/tests/unit/adapters/test_oracledb/test_json_handlers.py +++ b/tests/unit/adapters/test_oracledb/test_json_handlers.py @@ -154,16 +154,18 @@ def test_input_handler_routes_dict_to_db_type_blob_on_12c() -> None: cursor.var.assert_called_once_with(oracledb.DB_TYPE_BLOB, arraysize=2, inconverter=json_converter_in_blob) -def test_input_handler_routes_dict_to_db_type_blob_pre_12c() -> None: - """Dict bound on Oracle 11g (BLOB_PLAIN storage) should use DB_TYPE_BLOB.""" +def test_input_handler_routes_dict_to_db_type_clob_on_11g() -> None: + """Dict bound on Oracle 11g (<12c) uses DB_TYPE_CLOB with json_converter_in_clob.""" + import oracledb + cursor = _mock_cursor_with_major(11) cursor_var = Mock() cursor.var = Mock(return_value=cursor_var) - result = json_input_type_handler(cursor, {"foo": "bar"}, 2) + result = json_input_type_handler(cursor, {"foo": "bar"}, 4) assert result is cursor_var - cursor.var.assert_called_once_with(oracledb.DB_TYPE_BLOB, arraysize=2, inconverter=json_converter_in_blob) + cursor.var.assert_called_once_with(oracledb.DB_TYPE_CLOB, arraysize=4, inconverter=json_converter_in_clob) def test_input_handler_respects_compatible_major_on_21c() -> None: diff --git a/tests/unit/adapters/test_oracledb/test_litestar_store.py b/tests/unit/adapters/test_oracledb/test_litestar_store.py index df71233c7..722e5554a 100644 --- a/tests/unit/adapters/test_oracledb/test_litestar_store.py +++ b/tests/unit/adapters/test_oracledb/test_litestar_store.py @@ -2,10 +2,13 @@ from datetime import UTC, datetime, timedelta from typing import Any, Self, cast +from unittest.mock import AsyncMock, MagicMock + +import pytest from sqlspec.adapters.oracledb.core import DB_TYPE_BLOB from sqlspec.adapters.oracledb.data_dictionary import OracleVersionCache -from sqlspec.adapters.oracledb.litestar import OracleSyncStore +from sqlspec.adapters.oracledb.litestar import OracleAsyncStore, OracleSyncStore class _FakeCursor: @@ -13,7 +16,7 @@ class _FakeCursor: def __init__(self, rows: "list[tuple[Any, ...]] | None" = None) -> None: self.rows = list(rows or []) - self.executed: list[tuple[str, dict[str, Any] | None]] = [] + self.executed: list[tuple[str, dict[str, Any] | tuple[object, ...] | None]] = [] self.rowcount = 0 self.closed = False @@ -26,7 +29,7 @@ def __exit__(self, *_: object) -> None: def close(self) -> None: self.closed = True - def execute(self, sql: str, parameters: "dict[str, Any] | None" = None) -> None: + def execute(self, sql: str, parameters: "dict[str, Any] | tuple[object, ...] | None" = None) -> None: self.executed.append((sql, parameters)) def fetchone(self) -> "tuple[Any, ...] | None": @@ -113,7 +116,7 @@ def test_oracle_sync_store_expires_in_uses_database_clock() -> None: sql, parameters = cursor.executed[0] assert "SELECT expires_at, SYSTIMESTAMP" in sql - assert parameters == {"session_id": "session-1"} + assert parameters == ("session-1",) def test_oracle_litestar_store_interval_partition_on_expires_at() -> None: @@ -146,6 +149,7 @@ def test_oracle_sync_store_set_large_blob_single_statement() -> None: sql, parameters = cursor.executed[0] assert "MERGE INTO oracle_sessions t" in sql assert parameters is not None + parameters = cast("dict[str, object]", parameters) assert parameters["session_id"] == "session-large" assert connection.created_lobs == [(DB_TYPE_BLOB, large_payload)] assert parameters["data"] == (DB_TYPE_BLOB, large_payload) @@ -163,3 +167,51 @@ def test_oracle_litestar_store_table_ddl_sanitizes_and_truncates_index_name() -> assert "CREATE INDEX idx_my_very_long_schema_oracle" in ddl assert "idx_my_very_long_schema.oracle" not in ddl + + +def test_oracle_store_set_rolls_back_on_execute_failure() -> None: + """Store mutations roll back connection on execution failure.""" + from unittest.mock import MagicMock + + cursor = MagicMock() + cursor.execute.side_effect = RuntimeError("execute failed") + connection = MagicMock() + connection.cursor.return_value.__enter__.return_value = cursor + context = MagicMock() + context.__enter__.return_value = connection + config = MagicMock() + config.provide_connection.return_value = context + config.extension_config = {} + + store = OracleSyncStore(config) + with pytest.raises(RuntimeError, match="execute failed"): + store._set("fail-key", b"data") + + connection.rollback.assert_called_once() + + +@pytest.mark.anyio +@pytest.mark.parametrize("is_async", [False, True]) +async def test_session_renewal_binds_each_sql_placeholder(is_async: bool) -> None: + config = MagicMock() + config.extension_config = {} + cursor = MagicMock() + cursor.fetchone.return_value = (b"payload", datetime(2026, 5, 10, tzinfo=UTC)) + connection = MagicMock() + connection.cursor.return_value.__enter__.return_value = cursor + if is_async: + cursor.execute = AsyncMock() + cursor.fetchone = AsyncMock(return_value=cursor.fetchone.return_value) + connection.commit = AsyncMock() + config.provide_connection.return_value.__aenter__.return_value = connection + store = OracleAsyncStore(config) + else: + config.provide_connection.return_value.__enter__.return_value = connection + store = OracleSyncStore(config) + + result = await store.get("session-1", renew_for=60) + + assert result == b"payload" + assert cursor.execute.call_count == 2 + assert cursor.execute.call_args.args[1] == (60, 60, "session-1") + connection.commit.assert_called_once() diff --git a/tests/unit/adapters/test_oracledb/test_lob_coercion.py b/tests/unit/adapters/test_oracledb/test_lob_coercion.py index 09aeda895..63826b455 100644 --- a/tests/unit/adapters/test_oracledb/test_lob_coercion.py +++ b/tests/unit/adapters/test_oracledb/test_lob_coercion.py @@ -1127,3 +1127,17 @@ async def test_oracle_async_stream_source_closes_cursor_when_execute_raises() -> cursor.close.assert_called_once() assert source._cursor is None + + +def test_collect_sync_rows_preserves_lob_locators_when_fetch_lobs_true() -> None: + """collect_sync_rows preserves raw LOB locators when fetch_lobs=True.""" + from sqlspec.adapters.oracledb.core import collect_sync_rows + + raw_lob = MagicMock() + raw_lob.read = MagicMock(side_effect=AssertionError("read should not be called")) + description = [("payload", "CLOB", None, None, None, None, None)] + fetched = [(raw_lob,)] + rows, _cols = collect_sync_rows(fetched, description, {"fetch_lobs": True}) + assert len(rows) == 1 + assert rows[0][0] is raw_lob + raw_lob.read.assert_not_called() diff --git a/tests/unit/adapters/test_oracledb/test_oracle_adk_store.py b/tests/unit/adapters/test_oracledb/test_oracle_adk_store.py index 0f0b89af7..87e8adf41 100644 --- a/tests/unit/adapters/test_oracledb/test_oracle_adk_store.py +++ b/tests/unit/adapters/test_oracledb/test_oracle_adk_store.py @@ -3,8 +3,10 @@ from datetime import UTC, datetime from decimal import Decimal -from typing import Any, NotRequired, Self, cast, get_args, get_origin -from unittest.mock import MagicMock +from inspect import isawaitable +from types import SimpleNamespace +from typing import TYPE_CHECKING, Any, NotRequired, Self, cast, get_args, get_origin +from unittest.mock import AsyncMock, MagicMock from uuid import uuid4 import pytest @@ -22,7 +24,11 @@ from sqlspec.adapters.oracledb.adk.store import _event_data_column_ddl from sqlspec.adapters.oracledb.data_dictionary import OracleVersionCache, OracleVersionInfo from sqlspec.config import ADKConfig -from sqlspec.exceptions import SQLSpecError +from sqlspec.exceptions import ImproperConfigurationError, SQLSpecError +from sqlspec.extensions.adk import StoredMemory + +if TYPE_CHECKING: + from collections.abc import Awaitable def _mock_config(adk_config: dict[str, object]) -> MagicMock: @@ -47,12 +53,184 @@ def _config_for_storage(storage_type: JSONStorageType) -> MagicMock: return config +def _config_for_vectors() -> MagicMock: + config = _config_for_storage(JSONStorageType.JSON_NATIVE) + config._oracle_version_cache.version = OracleVersionInfo(23, 4, 0, compatible="23.4.0") + return config + + +def _memory_entry() -> StoredMemory: + timestamp = datetime(2026, 5, 10, 12, tzinfo=UTC) + return { + "id": "memory-1", + "session_id": "session-1", + "app_name": "app", + "user_id": "user", + "scope": "user", + "event_id": "event-1", + "author": "user", + "timestamp": timestamp, + "content_json": {"text": "hello"}, + "content_text": "hello", + "metadata_json": None, + "inserted_at": timestamp, + "embedding": None, + } + + +@pytest.mark.anyio +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.parametrize("error_code", [1, 1400, 12899, 51803]) +async def test_memory_batch_errors_commit_only_duplicate_skips(is_async: bool, error_code: int) -> None: + import oracledb + + config = _config_for_storage(JSONStorageType.BLOB_JSON) + config.extension_config = {"adk": {}} + cursor = MagicMock() + batch_error = SimpleNamespace(code=error_code, message=f"ORA-{error_code:05d}: rejected memory") + cursor.getbatcherrors.return_value = [batch_error] + connection = MagicMock() + connection.cursor.return_value.__enter__.return_value = cursor + if is_async: + cursor.executemany = AsyncMock() + connection.commit = AsyncMock() + connection.rollback = AsyncMock() + config.provide_connection.return_value.__aenter__.return_value = connection + store = OracleAsyncADKMemoryStore(config) + else: + config.provide_connection.return_value.__enter__.return_value = connection + store = OracleSyncADKMemoryStore(config) + + entries: list[StoredMemory] = [_memory_entry(), {**_memory_entry(), "id": "memory-2"}] + if error_code == 1: + result: int | Awaitable[int] = store.insert_memory_entries(entries) + if isawaitable(result): + result = await result + assert result == 1 + connection.commit.assert_called_once() + connection.rollback.assert_not_called() + else: + with pytest.raises(oracledb.DatabaseError) as exc_info: + result = store.insert_memory_entries(entries) + if isawaitable(result): + await result + assert exc_info.value.args[0] is batch_error + connection.commit.assert_not_called() + connection.rollback.assert_called_once() + + +@pytest.mark.anyio +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.parametrize("use_fts", [False, True]) +@pytest.mark.parametrize( + "version", + [ + OracleVersionInfo(19, 0, 0, compatible="19.0.0"), + OracleVersionInfo(21, 3, 0, compatible="21.0.0"), + OracleVersionInfo(23, 4, 0, compatible="23.0.0"), + ], +) +async def test_legacy_memory_retains_text_storage_and_rejects_vectors( + is_async: bool, use_fts: bool, version: OracleVersionInfo +) -> None: + config = _config_for_storage(JSONStorageType.BLOB_JSON) + config._oracle_version_cache.version = version + config.extension_config = {"adk": {"memory_use_fts": use_fts}} + cursor = MagicMock() + cursor.getbatcherrors.return_value = [] + timestamp = _memory_entry()["timestamp"] + cursor.fetchall.return_value = [ + ( + "memory-1", + "session-1", + "app", + "user", + "user", + "event-1", + "user", + timestamp, + b'{"text":"hello"}', + "hello", + None, + timestamp, + None, + ) + ] + connection = MagicMock() + connection.cursor.return_value.__enter__.return_value = cursor + executed: list[str] = [] + + def execute(sql: str, _params: Any) -> None: + assert "NULL AS embedding" in sql + executed.append(sql) + + if is_async: + cursor.execute = AsyncMock(side_effect=execute) + cursor.executemany = AsyncMock() + cursor.fetchall = AsyncMock(return_value=cursor.fetchall.return_value) + connection.commit = AsyncMock() + config.provide_connection.return_value.__aenter__.return_value = connection + store = OracleAsyncADKMemoryStore(config) + else: + cursor.execute.side_effect = execute + config.provide_connection.return_value.__enter__.return_value = connection + store = OracleSyncADKMemoryStore(config) + + ddl = store._memory_table_ddl() + if isawaitable(ddl): + ddl = await ddl + from sqlglot import parse_one + from sqlglot.expressions import Create + + create_sql = ddl.split("EXECUTE IMMEDIATE '", 1)[1].split("';", 1)[0].replace("''", "'") + table = parse_one(create_sql, dialect="oracle") + assert isinstance(table, Create) + assert {column.name for column in table.this.expressions} == { + "id", + "session_id", + "app_name", + "user_id", + "scope", + "event_id", + "author", + "timestamp", + "content_json", + "metadata_json", + "content_text", + "inserted_at", + } + + inserted: int | Awaitable[int] = store.insert_memory_entries([_memory_entry()]) + if isawaitable(inserted): + inserted = await inserted + assert inserted == 1 + insert_sql, bind_rows = cursor.executemany.call_args.args + assert "timestamp, content_json" in insert_sql + assert set(bind_rows[0]) == set(_memory_entry()) - {"embedding"} + + records = store.search_entries("hello", "app", "user") + if isawaitable(records): + records = await records + assert len(executed) == 1 + assert [(record["id"], record["embedding"]) for record in records] == [("memory-1", None)] + + with pytest.raises(ImproperConfigurationError, match=r"23\.4"): + result = store.insert_memory_entries([{**_memory_entry(), "embedding": [0.1]}]) + if isawaitable(result): + await result + with pytest.raises(ImproperConfigurationError, match=r"23\.4"): + result = store.search_entries("hello", "app", "user", embedding=[0.1]) + if isawaitable(result): + await result + + def test_oracle_adk_config_types_adapter_local_optimizations() -> None: """Oracle ADK optimization settings are typed on the adapter-local extension config.""" assert cast("Any", ADKConfig).__optional_keys__ <= cast("Any", OracleADKConfig).__optional_keys__ expected_types: dict[str, object] = { + "vector_dimensions": int, "in_memory": bool, "compression": OracleADKCompressionConfig, "partitioning": OracleADKPartitionConfig, @@ -301,6 +479,160 @@ def test_oracle_sync_adk_memory_rows_to_records_deserializes_json_fields() -> No assert records[0]["embedding"] is None +@pytest.mark.anyio +async def test_oracle_adk_memory_ddl_declares_embedding_vector() -> None: + """Memory table DDL declares an embedding VECTOR column with configured dimensions.""" + config = _config_for_vectors() + config.extension_config = {"adk": {}} + store = OracleSyncADKMemoryStore(config) + sql = store._memory_table_ddl() + assert "embedding VECTOR(768, FLOAT32)" in sql + + custom_config = _config_for_vectors() + custom_config.extension_config = {"adk": {"vector_dimensions": 1536}} + custom_store = OracleAsyncADKMemoryStore(custom_config) + custom_sql = await custom_store._memory_table_ddl() + assert "embedding VECTOR(1536, FLOAT32)" in custom_sql + + +def test_oracle_adk_sync_insert_binds_embedding_array() -> None: + """Sync memory store binds embeddings as float arrays.""" + import array + + config = _config_for_vectors() + config.extension_config = {"adk": {}} + store = OracleSyncADKMemoryStore(config) + conn = MagicMock() + cursor = MagicMock() + cursor.rowcount = 1 + conn.cursor.return_value.__enter__.return_value = cursor + config.provide_connection.return_value.__enter__.return_value = conn + + timestamp = datetime(2026, 5, 10, 12, 0, tzinfo=UTC) + entries: list[StoredMemory] = [ + { + "id": "mem-1", + "session_id": "sess-1", + "app_name": "app", + "user_id": "user", + "scope": "user", + "event_id": "ev-1", + "author": "user", + "timestamp": timestamp, + "content_json": {"k": "v"}, + "content_text": "text", + "metadata_json": None, + "inserted_at": timestamp, + "embedding": [0.1, 0.2, 0.3], + } + ] + inserted = store.insert_memory_entries(entries) + assert inserted == 1 + assert cursor.executemany.called + _call_sql, call_params = cursor.executemany.call_args[0] + first_param = call_params[0] + assert "embedding" in first_param + assert isinstance(first_param["embedding"], array.array) + assert first_param["embedding"].typecode == "f" + assert list(first_param["embedding"]) == pytest.approx([0.1, 0.2, 0.3]) + + +def test_oracle_adk_memory_rows_to_records_decodes_embedding() -> None: + """Memory rows to records round-trips vector embeddings as float lists.""" + import array + + store = OracleSyncADKMemoryStore.__new__(OracleSyncADKMemoryStore) # type: ignore[call-arg] + store._config = _config_for_storage(JSONStorageType.BLOB_JSON) + timestamp = datetime(2026, 5, 10, 12, 0, tzinfo=UTC) + row = ( + "memory-vec", + "session-vec", + "app", + "user", + "user", + "event-vec", + "user", + timestamp, + b'{"text":"vec"}', + "vec", + b'{"source":"unit"}', + timestamp, + array.array("f", [0.1, 0.2, 0.3]), + ) + records = store._rows_to_records([row]) + assert len(records) == 1 + assert records[0]["embedding"] == pytest.approx([0.1, 0.2, 0.3]) + + +def test_adk_store_mutation_rolls_back_on_error() -> None: + """ADK store mutation rolls back connection on error.""" + cursor = MagicMock() + cursor.execute.side_effect = RuntimeError("mutation error") + conn = MagicMock() + conn.cursor.return_value.__enter__.return_value = cursor + config = _config_for_storage(JSONStorageType.BLOB_JSON) + config.extension_config = {"adk": {}} + config.provide_connection.return_value.__enter__.return_value = conn + + store = OracleSyncADKStore(config) + with pytest.raises(RuntimeError, match="mutation error"): + store.delete_session("app", "user", "sess-1") + + conn.rollback.assert_called_once() + + +def test_oracle_adk_memory_store_batch_insert() -> None: + """ADK memory store executes batch insert via executemany.""" + cursor = MagicMock() + cursor.rowcount = 2 + cursor.getbatcherrors.return_value = [] + conn = MagicMock() + conn.cursor.return_value.__enter__.return_value = cursor + config = _config_for_storage(JSONStorageType.JSON_NATIVE) + config.extension_config = {"adk": {}} + config.provide_connection.return_value.__enter__.return_value = conn + + store = OracleSyncADKMemoryStore(config) + now = datetime(2026, 5, 10, 12, 0, tzinfo=UTC) + entries: list[StoredMemory] = [ + { + "id": "mem-1", + "session_id": "sess-1", + "app_name": "app", + "user_id": "user", + "scope": "user", + "event_id": "ev-1", + "author": "user", + "timestamp": now, + "content_json": {"k": 1}, + "content_text": "text 1", + "metadata_json": None, + "inserted_at": now, + "embedding": None, + }, + { + "id": "mem-2", + "session_id": "sess-1", + "app_name": "app", + "user_id": "user", + "scope": "user", + "event_id": "ev-2", + "author": "user", + "timestamp": now, + "content_json": {"k": 2}, + "content_text": "text 2", + "metadata_json": None, + "inserted_at": now, + "embedding": None, + }, + ] + inserted = store.insert_memory_entries(entries) + assert inserted == 2 + cursor.executemany.assert_called_once() + assert len(cursor.executemany.call_args[0][1]) == 2 + conn.commit.assert_called_once() + + def _sync_store_with_driver() -> "tuple[Any, MagicMock, MagicMock]": config = _config_for_storage(JSONStorageType.BLOB_JSON) config.extension_config = {"adk": {}} @@ -364,6 +696,13 @@ def close(self) -> None: def execute(self, sql: str, params: "dict[str, Any] | None" = None) -> None: self.calls.append((sql, dict(params or {}))) + def executemany(self, sql: str, params: Any = None, **_kwargs: Any) -> None: + for p in params or []: + self.calls.append((sql, dict(p or {}))) + + def getbatcherrors(self) -> list[Any]: + return [] + def fetchall(self) -> "list[Any]": return [] @@ -629,7 +968,7 @@ def test_oracle_adk_owner_id_column_ddl_escapes_single_quotes() -> None: def test_oracle_sync_adk_memory_fts_sanitizes_query_and_hybrid_rrf_fuses_candidates() -> None: cursor = _RecordingCursor() conn = _RecordingConnection(cursor) - config = _config_for_storage(JSONStorageType.JSON_NATIVE) + config = _config_for_vectors() config.extension_config = {"adk": {"memory_use_fts": True}} config.provide_connection = lambda *_a, **_k: conn store = OracleSyncADKMemoryStore(config) @@ -657,7 +996,7 @@ def test_oracle_sync_adk_memory_fts_sanitizes_query_and_hybrid_rrf_fuses_candida async def test_oracle_async_adk_memory_fts_sanitizes_query_and_hybrid_rrf_fuses_candidates() -> None: cursor = _AsyncRecordingCursor() conn = _RecordingConnection(cursor) - config = _config_for_storage(JSONStorageType.JSON_NATIVE) + config = _config_for_vectors() config.extension_config = {"adk": {"memory_use_fts": True}} config.provide_connection = lambda *_a, **_k: conn store = OracleAsyncADKMemoryStore(config) @@ -675,7 +1014,7 @@ def test_oracle_sync_adk_memory_hybrid_rrf_near_and_exact_fallback() -> None: """Sync ADK memory store forwards fts_operator='NEAR' and approximate=False.""" cursor = _RecordingCursor() conn = _RecordingConnection(cursor) - config = _config_for_storage(JSONStorageType.JSON_NATIVE) + config = _config_for_vectors() config.extension_config = {"adk": {"memory_use_fts": True, "fts_operator": "NEAR", "approximate": False}} config.provide_connection = lambda *_a, **_k: conn store = OracleSyncADKMemoryStore(config) @@ -702,7 +1041,7 @@ async def test_oracle_async_adk_memory_hybrid_rrf_near_and_exact_fallback() -> N """Async ADK memory store forwards fts_operator='NEAR' and approximate=False.""" cursor = _AsyncRecordingCursor() conn = _RecordingConnection(cursor) - config = _config_for_storage(JSONStorageType.JSON_NATIVE) + config = _config_for_vectors() config.extension_config = {"adk": {"memory_use_fts": True, "fts_operator": "NEAR", "approximate": False}} config.provide_connection = lambda *_a, **_k: conn store = OracleAsyncADKMemoryStore(config) diff --git a/tests/unit/adapters/test_oracledb/test_oracle_version_unification.py b/tests/unit/adapters/test_oracledb/test_oracle_version_unification.py index aa4f41fdc..ab3295513 100644 --- a/tests/unit/adapters/test_oracledb/test_oracle_version_unification.py +++ b/tests/unit/adapters/test_oracledb/test_oracle_version_unification.py @@ -242,3 +242,19 @@ def test_oracle_events_store_compression_partitioning_config() -> None: assert "TABLESPACE event_data" in sql assert "PARTITION BY RANGE (available_at)" in sql assert "NUMTODSINTERVAL(1, ''DAY'')" in sql + + +def test_resolve_oracle_connection_major_from_core() -> None: + """resolve_oracle_connection_major and cache are importable and callable directly from core.""" + from sqlspec.adapters.oracledb.core import ( + OracleVersionCache, + resolve_oracle_connection_compatible_major, + resolve_oracle_connection_major, + ) + + conn = MagicMock() + conn._sqlspec_oracle_major = 23 + conn._sqlspec_oracle_compatible_major = 23 + cache = OracleVersionCache() + assert resolve_oracle_connection_major(conn, cache) == 23 + assert resolve_oracle_connection_compatible_major(conn, cache) == 23 diff --git a/tests/unit/adapters/test_oracledb/test_stack_native_transaction.py b/tests/unit/adapters/test_oracledb/test_stack_native_transaction.py index a70038361..061e22be5 100644 --- a/tests/unit/adapters/test_oracledb/test_stack_native_transaction.py +++ b/tests/unit/adapters/test_oracledb/test_stack_native_transaction.py @@ -14,7 +14,7 @@ class _FakeOpResult: - def __init__(self, error: "Exception | None" = None) -> None: + def __init__(self, error: "Any | None" = None) -> None: self.error = error @@ -108,3 +108,101 @@ async def test_native_stack_inside_an_implicit_transaction_does_not_commit() -> assert connection.commit_count == 0 assert connection.rollback_count == 0 + + +@pytest.mark.anyio +async def test_native_stack_wraps_mapped_error() -> None: + """Stack execution errors wrap mapped SQLSpec exceptions, not raw driver errors.""" + import oracledb + + from sqlspec.exceptions import SQLSpecError, StackExecutionError + + native_err = oracledb.DatabaseError("ORA-00001: unique constraint violated") + results = [_FakeOpResult(native_err)] + connection = _FakePipelineConnection(results) + driver = OracleAsyncDriver(cast("Any", connection)) + + stack_results = await driver._execute_stack_native(_stack(), continue_on_error=True) + assert len(stack_results) == 1 + error = stack_results[0].error + assert isinstance(error, StackExecutionError) + assert isinstance(error.original_error, SQLSpecError) + + +@pytest.mark.anyio +@pytest.mark.parametrize("continue_on_error", [False, True]) +async def test_stack_maps_native_pipeline_error_records(continue_on_error: bool) -> None: + import oracledb + from oracledb.errors import _Error + + from sqlspec.exceptions import StackExecutionError, UniqueViolationError + + native_error = _Error("ORA-00001: unique constraint violated", code=1) + connection = _FakePipelineConnection([_FakeOpResult(), _FakeOpResult(native_error), _FakeOpResult()]) + driver = OracleAsyncDriver(cast("Any", connection)) + driver._pipeline_support = True + + if continue_on_error: + results = await driver.execute_stack(_stack(), continue_on_error=True) + assert len(results) == 3 + assert results[0].error is None + assert results[2].error is None + error = results[1].error + assert connection.commit_count == 1 + assert connection.rollback_count == 0 + else: + with pytest.raises(StackExecutionError) as exc_info: + await driver.execute_stack(_stack()) + error = exc_info.value + assert connection.commit_count == 0 + assert connection.rollback_count == 1 + + assert isinstance(error, StackExecutionError) + assert error.operation_index == 1 + assert isinstance(error.original_error, UniqueViolationError) + cause = error.original_error.__cause__ + assert isinstance(cause, oracledb.DatabaseError) + assert cause.args[0] is native_error + + +def test_sync_stack_executes_and_commits_without_a_storage_capability_flag() -> None: + from sqlspec.adapters.oracledb._typing import OracleSyncConnection + from sqlspec.adapters.oracledb.driver import OracleSyncDriver + + class _Cursor: + description = None + rowcount = 1 + + def __init__(self, connection: "_Connection") -> None: + self.connection = connection + + def execute(self, statement: str, parameters: object = None, **kwargs: object) -> None: + self.connection.executed.append((statement, parameters)) + + def close(self) -> None: + return None + + class _Connection: + transaction_in_progress = False + + def __init__(self) -> None: + self.executed: list[tuple[str, object]] = [] + self.commits = 0 + + def cursor(self) -> _Cursor: + return _Cursor(self) + + def commit(self) -> None: + self.commits += 1 + + connection = _Connection() + driver = OracleSyncDriver(cast("OracleSyncConnection", connection)) + stack = ( + StatementStack() + .push_execute("INSERT INTO items VALUES (:1)", (1,)) + .push_execute("INSERT INTO items VALUES (:1)", (2,)) + ) + results = driver.execute_stack(stack) + assert [result.rows_affected for result in results] == [1, 1] + assert connection.executed == [("INSERT INTO items VALUES (:1)", [1]), ("INSERT INTO items VALUES (:1)", [2])] + assert connection.commits == 1 diff --git a/tests/unit/adapters/test_oracledb/test_storage_features.py b/tests/unit/adapters/test_oracledb/test_storage_features.py index feb5aaca9..d1a1b495a 100644 --- a/tests/unit/adapters/test_oracledb/test_storage_features.py +++ b/tests/unit/adapters/test_oracledb/test_storage_features.py @@ -56,4 +56,14 @@ def test_oracle_range_partitioning_uses_interval() -> None: ) assert "PARTITION BY RANGE (timestamp)" in report["clause"] - assert "INTERVAL (NUMTOYMINTERVAL(1, ''MONTH''))" in report["clause"] + assert "INTERVAL (NUMTOYMINTERVAL(1, 'MONTH'))" in report["clause"] + assert "INTERVAL (NUMTOYMINTERVAL(1, ''MONTH''))" in report["plsql_clause"] + + +def test_validate_oracle_identifier_accepts_dollar_sign() -> None: + """validate_oracle_identifier accepts $, #, and _ in Oracle identifiers.""" + from sqlspec.adapters.oracledb._storage import validate_oracle_identifier + + assert validate_oracle_identifier("tbl$name", "table name") == "tbl$name" + assert validate_oracle_identifier("tbl#name", "table name") == "tbl#name" + assert validate_oracle_identifier("tbl_name", "table name") == "tbl_name" diff --git a/tests/unit/adapters/test_oracledb/test_transaction_state.py b/tests/unit/adapters/test_oracledb/test_transaction_state.py index 5f7078239..0662a6dbb 100644 --- a/tests/unit/adapters/test_oracledb/test_transaction_state.py +++ b/tests/unit/adapters/test_oracledb/test_transaction_state.py @@ -114,3 +114,63 @@ def test_execute_stack_does_not_commit_after_begin_without_dml() -> None: driver.commit() assert driver._connection_in_transaction() is False + + +def test_failed_commit_resets_transaction_active() -> None: + """A failed commit must reset _transaction_active in finally.""" + import oracledb + + from sqlspec.exceptions import SQLSpecError + + class _FailingCommitConn(_FakeSyncConnection): + def commit(self) -> None: + raise oracledb.DatabaseError("ORA-02091: transaction rolled back") + + driver = OracleSyncDriver(cast("Any", _FailingCommitConn())) + driver.begin() + assert driver._connection_in_transaction() is True + with pytest.raises(SQLSpecError): + driver.commit() + assert driver._connection_in_transaction() is False + + +def test_connection_context_rolls_back_on_block_exception() -> None: + """Connection context rolls back on unhandled block exceptions.""" + from unittest.mock import MagicMock + + from sqlspec.adapters.oracledb.config import OracleSyncConnectionContext + + conn = MagicMock() + config = MagicMock() + config.driver_features = {} + config.connection_instance.acquire.return_value = conn + ctx = OracleSyncConnectionContext(config) + + with pytest.raises(RuntimeError): + with ctx: + raise RuntimeError("boom") + + conn.rollback.assert_called_once() + + +@pytest.mark.anyio +async def test_async_connection_context_rolls_back_on_block_exception() -> None: + """Async connection context rolls back on unhandled block exceptions.""" + from unittest.mock import AsyncMock, MagicMock + + from sqlspec.adapters.oracledb.config import OracleAsyncConnectionContext + + conn = AsyncMock() + pool = MagicMock() + pool.acquire = AsyncMock(return_value=conn) + pool.release = AsyncMock() + config = MagicMock() + config.driver_features = {} + config.provide_pool = AsyncMock(return_value=pool) + ctx = OracleAsyncConnectionContext(config) + + with pytest.raises(RuntimeError): + async with ctx: + raise RuntimeError("boom") + + conn.rollback.assert_called_once() diff --git a/tests/unit/adapters/test_oracledb/test_uuid_handlers.py b/tests/unit/adapters/test_oracledb/test_uuid_handlers.py index 3a5d326dc..c7b37d7e9 100644 --- a/tests/unit/adapters/test_oracledb/test_uuid_handlers.py +++ b/tests/unit/adapters/test_oracledb/test_uuid_handlers.py @@ -252,3 +252,25 @@ def test_uuid_handlers_use_generic_chained_wrappers() -> None: assert callable(connection.outputtypehandler) assert len(inspect.signature(connection.outputtypehandler).parameters) == 2 assert len(inspect.signature(connection.inputtypehandler).parameters) == 3 + + +def test_output_type_handler_with_fetch_info_object() -> None: + """_output_type_handler handles FetchInfo objects with attributes.""" + import oracledb + + from sqlspec.adapters.oracledb._uuid_handlers import _output_type_handler + + cursor = Mock() + cursor.arraysize = 32 + cursor_var = Mock() + cursor.var = Mock(return_value=cursor_var) + + class _FetchInfoStub: + def __init__(self) -> None: + self.name = "ID" + self.type_code = oracledb.DB_TYPE_RAW + self.internal_size = 16 + + result = _output_type_handler(cursor, _FetchInfoStub()) + assert result is cursor_var + cursor.var.assert_called_once_with(oracledb.DB_TYPE_RAW, arraysize=32, outconverter=uuid_converter_out) diff --git a/tests/unit/adapters/test_oracledb/test_vector_handlers.py b/tests/unit/adapters/test_oracledb/test_vector_handlers.py index 728d7bdde..7aeb7b9be 100644 --- a/tests/unit/adapters/test_oracledb/test_vector_handlers.py +++ b/tests/unit/adapters/test_oracledb/test_vector_handlers.py @@ -86,18 +86,18 @@ def test_input_handler_claims_tuple_of_float() -> None: assert list(packed) == [0.5, 0.25, 0.125] -def test_input_handler_packs_int_in_int8_range_as_int8() -> None: - """``list[int]`` with all values in [-128, 127] is packed as int8.""" +def test_input_handler_packs_int_sequence_as_float32() -> None: + """``list[int]`` is deterministically packed as float32.""" from sqlspec.adapters.oracledb._vector_handlers import _input_type_handler # pyright: ignore[reportPrivateUsage] cursor = _mock_cursor() cursor.var.return_value = "var-token" - _input_type_handler(cursor, [-128, 0, 64, 127], 1) + _input_type_handler(cursor, [0, 1, 0], 1) - packed = cursor.var.call_args.kwargs["inconverter"]([-128, 0, 64, 127]) - assert packed.typecode == "b" - assert list(packed) == [-128, 0, 64, 127] + packed = cursor.var.call_args.kwargs["inconverter"]([0, 1, 0]) + assert packed.typecode == "f" + assert list(packed) == [0.0, 1.0, 0.0] def test_input_handler_packs_int_outside_int8_range_as_float32() -> None: @@ -382,12 +382,12 @@ def test_pack_sequence_converter_float_list() -> None: def test_pack_sequence_converter_int8_range() -> None: - """_pack_sequence_converter packs small int lists as int8.""" + """_pack_sequence_converter packs small int lists as float32.""" result = _pack_sequence_converter([-128, 0, 127]) assert isinstance(result, array.array) - assert result.typecode == "b" - assert list(result) == [-128, 0, 127] + assert result.typecode == "f" + assert list(result) == [-128.0, 0.0, 127.0] def test_input_handler_rejects_heterogeneous_list_and_pre_23ai_and_disabled_flag() -> None: diff --git a/tests/unit/adapters/test_postgres_migration_schema.py b/tests/unit/adapters/test_postgres_migration_schema.py index ced043c77..d249ad0a1 100644 --- a/tests/unit/adapters/test_postgres_migration_schema.py +++ b/tests/unit/adapters/test_postgres_migration_schema.py @@ -5,6 +5,7 @@ from typing import Any import pytest +from psycopg.sql import Composable from sqlspec.adapters.adbc.config import AdbcConfig from sqlspec.adapters.adbc.driver import AdbcDriver @@ -50,8 +51,8 @@ def __init__(self, schema_exists: bool = True) -> None: self.schema_exists = schema_exists self.executed: list[tuple[str, Any]] = [] - def execute(self, sql: str, parameters: Any | None = None) -> None: - self.executed.append((sql, parameters)) + def execute(self, sql: str | Composable, parameters: Any | None = None) -> None: + self.executed.append((sql.as_string() if isinstance(sql, Composable) else sql, parameters)) def fetchone(self) -> tuple[int] | None: return (1,) if self.schema_exists else None @@ -73,8 +74,8 @@ def __init__(self, schema_exists: bool = True) -> None: self.schema_exists = schema_exists self.executed: list[tuple[str, Any]] = [] - async def execute(self, sql: str, parameters: Any | None = None) -> None: - self.executed.append((sql, parameters)) + async def execute(self, sql: str | Composable, parameters: Any | None = None) -> None: + self.executed.append((sql.as_string() if isinstance(sql, Composable) else sql, parameters)) async def fetchone(self) -> tuple[int] | None: return (1,) if self.schema_exists else None diff --git a/tests/unit/adapters/test_postgres_native_options.py b/tests/unit/adapters/test_postgres_native_options.py index 8f76f8d18..27615733f 100644 --- a/tests/unit/adapters/test_postgres_native_options.py +++ b/tests/unit/adapters/test_postgres_native_options.py @@ -47,10 +47,12 @@ async def test_asyncpg_custom_codec_is_registered_after_internal_setup() -> None } -def test_asyncpg_pgbouncer_disables_prepared_cache_consistently() -> None: - config = AsyncpgConfig(connection_config={"pgbouncer": True, "statement_cache_size": 128}) +def test_asyncpg_pooler_mode_transaction_disables_prepared_cache_consistently() -> None: + config = AsyncpgConfig( + connection_config={"statement_cache_size": 128}, driver_features={"pooler_mode": "transaction"} + ) assert config.connection_config["statement_cache_size"] == 0 - assert config.driver_features["pgbouncer"] is True + assert config.driver_features["pooler_mode"] == "transaction" @pytest.mark.parametrize( @@ -62,7 +64,9 @@ async def test_null_pool_preserves_throttling_without_leaking_flag( ) -> None: constructor = MagicMock() monkeypatch.setattr(psycopg_module, pool_name, constructor) - config = config_type(connection_config={"null_pool": True, "open": False, "max_size": 7, "max_waiting": 2}) + config = config_type( + connection_config={"open": False, "max_size": 7, "max_waiting": 2}, driver_features={"enable_pooling": False} + ) if config_type is PsycopgAsyncConfig: await config._create_pool() else: @@ -70,8 +74,8 @@ async def test_null_pool_preserves_throttling_without_leaking_flag( args = constructor.call_args.kwargs assert args["max_size"] == 7 assert args["max_waiting"] == 2 - assert "null_pool" not in args["kwargs"] - assert "null_pool" not in config._connection_kwargs()[2] + assert "enable_pooling" not in args["kwargs"] + assert "enable_pooling" not in config._connection_kwargs()[2] @pytest.mark.parametrize("builder", [asyncpg_cockroach_config, psycopg_cockroach_config]) @@ -119,7 +123,7 @@ async def test_pgbouncer_stack_cancellation_keeps_transaction_cleanup() -> None: transaction.__aenter__ = AsyncMock() transaction.__aexit__ = AsyncMock(return_value=False) connection.transaction.return_value = transaction - driver = AsyncpgDriver(connection, driver_features={"pgbouncer": True}) + driver = AsyncpgDriver(connection, driver_features={"pooler_mode": "transaction"}) stack = StatementStack().push_execute("SELECT 1") with pytest.raises(asyncio.CancelledError): await driver.execute_stack(stack) diff --git a/tests/unit/adapters/test_psqlpy/test_adk_store.py b/tests/unit/adapters/test_psqlpy/test_adk_store.py index 0637c6411..8a1e56636 100644 --- a/tests/unit/adapters/test_psqlpy/test_adk_store.py +++ b/tests/unit/adapters/test_psqlpy/test_adk_store.py @@ -3,13 +3,14 @@ from datetime import UTC, datetime from typing import Any, NotRequired, Self, cast, get_args, get_origin from unittest.mock import MagicMock +from uuid import UUID import pytest from psqlpy.exceptions import ConnectionExecuteError, DatabaseError +from sqlspec.adapters.psqlpy import PsqlpyConfig from sqlspec.adapters.psqlpy.adk import PsqlpyADKConfig, PsqlpyADKMemoryStore, PsqlpyADKStore from sqlspec.adapters.psqlpy.adk.store import _is_psqlpy_database_error, _is_table_missing_error, decode_postgres_vector -from sqlspec.adapters.psqlpy.core import detect_postgres_extensions from sqlspec.config import ADKConfig @@ -298,17 +299,24 @@ async def test_psqlpy_memory_store_normalizes_tuple_embeddings_on_insert_and_sea results = await store.search_entries("", "app", "user-1", embedding=(0.1, 0.2)) assert inserted == 1 - assert conn.execute_calls[0][1][8] == [0.1, 0.2] - assert conn.fetch_calls[0][1][0] == [0.1, 0.2] + assert conn.fetch_calls[0][1][-1] == ["[0.1,0.2]"] + assert conn.fetch_calls[1][1][0] == [0.1, 0.2] assert results[0]["embedding"] == [0.1, 0.2] -async def test_psqlpy_detect_postgres_extensions_sets_pg_textsearch_available() -> None: - """detect_postgres_extensions detects vector, pg_search, and pg_textsearch for psqlpy connections.""" - conn = _FakeMemoryConnection([{"extname": "vector"}, {"extname": "pg_search"}, {"extname": "pg_textsearch"}]) - pgvector, paradedb, pg_textsearch = await detect_postgres_extensions( - conn, enable_pgvector=True, enable_paradedb=True, enable_pg_textsearch=True - ) - assert pgvector is True - assert paradedb is True - assert pg_textsearch is True +@pytest.mark.parametrize(("owner_type", "owner_id"), [("INTEGER", 42), ("UUID", UUID(int=1))]) +async def test_psqlpy_memory_batch_preserves_owner_type( + owner_type: str, owner_id: object, monkeypatch: pytest.MonkeyPatch +) -> None: + """A scalar owner bind lets the destination column infer its native type.""" + conn = _FakeMemoryConnection([_sample_memory_row(None)]) + monkeypatch.setattr(PsqlpyConfig, "provide_connection", lambda *_a, **_k: conn) + config = PsqlpyConfig(extension_config={"adk": {"owner_id_column": f"owner_id {owner_type}"}}) + store = PsqlpyADKMemoryStore(config) + + assert await store.insert_memory_entries([cast("Any", _sample_memory_row(None))], owner_id) == 1 + + sql, params = conn.fetch_calls[0] + assert "author, $14, timestamp" in _normalized(sql) + assert params[-1] == owner_id + assert len(params[0]) == 1 diff --git a/tests/unit/adapters/test_psqlpy/test_config.py b/tests/unit/adapters/test_psqlpy/test_config.py index e3494c315..993d0bdc9 100644 --- a/tests/unit/adapters/test_psqlpy/test_config.py +++ b/tests/unit/adapters/test_psqlpy/test_config.py @@ -287,6 +287,8 @@ async def test_psqlpy_session_context_rolls_back_open_transaction_on_exit() -> N @pytest.mark.anyio async def test_psqlpy_pooler_mode_transaction_disables_prepared_statements() -> None: """Transaction pooler mode should construct a valid ConnectionPool and pass prepared=False on execution.""" + assert "pgbouncer" not in PsqlpyDriverFeatures.__annotations__ + assert "pgbouncer_mode" not in PsqlpyDriverFeatures.__annotations__ config = PsqlpyConfig( connection_config={"dsn": "postgresql://localhost:6432/test"}, driver_features={"pooler_mode": "transaction"} ) @@ -294,13 +296,6 @@ async def test_psqlpy_pooler_mode_transaction_disables_prepared_statements() -> pool = await config.create_pool() pool.close() - alias_config = PsqlpyConfig( - connection_config={"dsn": "postgresql://localhost:6432/test"}, driver_features={"pgbouncer": True} - ) - assert alias_config.driver_features["pooler_mode"] == "transaction" - alias_pool = await alias_config.create_pool() - alias_pool.close() - conn = AsyncMock() conn.fetch.return_value = _ExtensionQueryResult(set()) conn.execute.return_value = "CREATE TABLE" @@ -335,12 +330,12 @@ async def test_psqlpy_derive_direct_config_and_columnar_engine() -> None: ) assert auto_direct_config.connection_config["options"] == "-c google_columnar_engine.enable_columnar_scan=on" - auto_pgbouncer_config = PsqlpyConfig( + tx_columnar_config = PsqlpyConfig( connection_config={"dsn": "postgresql://pooler:6432/test"}, - driver_features={"pooler_mode": "auto", "pgbouncer": True, "enable_alloydb_columnar": True}, + driver_features={"pooler_mode": "transaction", "enable_alloydb_columnar": True}, ) - assert auto_pgbouncer_config.driver_features["pooler_mode"] == "transaction" - assert "options" not in auto_pgbouncer_config.connection_config + assert tx_columnar_config.driver_features["pooler_mode"] == "transaction" + assert "options" not in tx_columnar_config.connection_config with pytest.raises(ImproperConfigurationError, match="Invalid pooler_mode"): PsqlpyConfig( @@ -375,3 +370,42 @@ async def test_psqlpy_derive_direct_config_and_columnar_engine() -> None: await driver.begin() conn.execute.assert_any_await("BEGIN") conn.execute.assert_any_await("SET LOCAL google_columnar_engine.enable_columnar_scan = 'on'") + + +@pytest.mark.anyio +async def test_psqlpy_auth_token_provider_evaluated_on_create_pool_not_init() -> None: + """Synchronous and async auth_token_provider is evaluated when creating a pool, not during __init__.""" + call_count = 0 + + def token_provider() -> str: + nonlocal call_count + call_count += 1 + return "dynamic_token" + + config = PsqlpyConfig( + connection_config={"dsn": "postgresql://localhost:5432/test", "auth_token_provider": token_provider} + ) + assert call_count == 0 + + pool = await config.create_pool() + assert call_count == 1 + pool.close() + + +@pytest.mark.anyio +async def test_psqlpy_sync_on_connection_create_callback() -> None: + """Synchronous callback for on_connection_create should execute without awaiting.""" + hook_called = False + + def sync_hook(conn: Any) -> None: + nonlocal hook_called + hook_called = True + + config = PsqlpyConfig( + connection_config={"dsn": "postgresql://localhost:5432/test"}, + driver_features={"on_connection_create": sync_hook}, + ) + mock_conn = AsyncMock() + mock_conn.fetch = AsyncMock(return_value=None) + await config._ensure_connection(mock_conn) # pyright: ignore[reportPrivateUsage] + assert hook_called is True diff --git a/tests/unit/adapters/test_psqlpy/test_core.py b/tests/unit/adapters/test_psqlpy/test_core.py index d8e03eaeb..80d44a07f 100644 --- a/tests/unit/adapters/test_psqlpy/test_core.py +++ b/tests/unit/adapters/test_psqlpy/test_core.py @@ -67,6 +67,15 @@ def test_create_mapped_exception_leaves_not_supported_error_generic() -> None: assert type(mapped) is SQLSpecError +def test_create_mapped_exception_syntax_error_containing_connection_maps_to_sql_parsing_error() -> None: + """Syntax error containing the substring 'connection' maps to SQLParsingError, not DatabaseConnectionError.""" + from sqlspec.exceptions import SQLParsingError + + error = psqlpy_exceptions.Error('syntax error at or near "connection"') + mapped = psqlpy_core.create_mapped_exception(error) + assert type(mapped) is SQLParsingError + + def test_format_execute_many_parameters_no_coercion_reuses_list_rows() -> None: """Formatting should preserve list rows when no numeric coercion is requested.""" records = [[1, "a"], [2, "b"]] diff --git a/tests/unit/adapters/test_psqlpy/test_driver.py b/tests/unit/adapters/test_psqlpy/test_driver.py index cd7b47b85..8ed910ba7 100644 --- a/tests/unit/adapters/test_psqlpy/test_driver.py +++ b/tests/unit/adapters/test_psqlpy/test_driver.py @@ -1,15 +1,36 @@ """Unit tests for PsqlpyDriver wire-level sqlcommenter and _dml_count_query cache preservation.""" +from pathlib import Path from types import SimpleNamespace from typing import Any, cast from unittest.mock import AsyncMock, patch import pytest +from sqlspec.adapters.psqlpy import PsqlpyConfig from sqlspec.adapters.psqlpy import core as psqlpy_core from sqlspec.adapters.psqlpy.core import build_statement_config from sqlspec.adapters.psqlpy.driver import PsqlpyDriver +from sqlspec.core import SQL from sqlspec.core.sqlcommenter import SQLCommenterContext +from sqlspec.exceptions import UniqueViolationError + + +@pytest.mark.anyio +async def test_storage_bridge_writes_client_converted_rows_without_native_export(tmp_path: Path) -> None: + """Client-side storage export works without claiming native Arrow export.""" + parquet = pytest.importorskip("pyarrow.parquet") + connection = AsyncMock() + connection.fetch.return_value = SimpleNamespace(result=lambda: [{"name": "Ada", "value": 1}]) + config = PsqlpyConfig() + driver = PsqlpyDriver(cast("Any", connection), driver_features=config.driver_features) + target = tmp_path / "psqlpy.parquet" + + job = await driver.select_to_storage("SELECT name, value FROM storage_rows", target, format_hint="parquet") + + assert config.storage_capabilities()["arrow_export_enabled"] is False + assert job.telemetry["rows_processed"] == 1 + assert parquet.read_table(target).to_pylist() == [{"name": "Ada", "value": 1}] @pytest.mark.anyio @@ -50,3 +71,155 @@ def _fake_fetch(query: str, _params: Any) -> _FakeResult: assert len(fetched_queries) == 5 assert "0000000000000001" in fetched_queries[0] assert "0000000000000005" in fetched_queries[4] + + +@pytest.mark.anyio +async def test_psqlpy_dispatch_execute_script_pooler_mode_disables_prepared() -> None: + """dispatch_execute_script in pooler mode passes prepared=False and emits script observer events.""" + from sqlspec.observability import ObservabilityConfig, ObservabilityRuntime + + cursor = AsyncMock() + cursor.execute = AsyncMock() + driver = PsqlpyDriver( + connection=cast("Any", cursor), + statement_config=build_statement_config(), + driver_features={"pooler_mode": "transaction"}, + ) + observed_events: list[Any] = [] + obs_config = ObservabilityConfig(statement_observers=(observed_events.append,)) + driver._observability = ObservabilityRuntime(obs_config) # pyright: ignore[reportPrivateUsage] + + stmt = SQL("SELECT 1; SELECT 2", is_script=True, statement_config=build_statement_config()) + result = await driver.dispatch_execute_script(cast("Any", cursor), stmt) + + assert result.statement_count == 2 + assert result.successful_statements == 2 + assert cursor.execute.await_count == 2 + for call_args in cursor.execute.await_args_list: + assert call_args.kwargs.get("prepared") is False + + assert len(observed_events) == 2 + assert observed_events[0].operation == "SCRIPT" + assert observed_events[0].script_statement_count == 2 + + +@pytest.mark.anyio +async def test_psqlpy_transaction_methods_raise_mapped_exceptions() -> None: + """begin, commit, and rollback wrap driver errors in mapped exceptions.""" + from psqlpy.exceptions import OperationalError as PsqlpyOperationalError + + from sqlspec.exceptions import DatabaseConnectionError + + conn = AsyncMock() + conn.in_transaction = False + conn.execute = AsyncMock(side_effect=PsqlpyOperationalError("server closed the connection unexpectedly")) + + driver = PsqlpyDriver(connection=cast("Any", conn), statement_config=build_statement_config()) + + with pytest.raises(DatabaseConnectionError): + await driver.begin() + + driver._transaction_active = True # pyright: ignore[reportPrivateUsage] + with pytest.raises(DatabaseConnectionError): + await driver.commit() + + driver._transaction_active = True # pyright: ignore[reportPrivateUsage] + with pytest.raises(DatabaseConnectionError): + await driver.rollback() + + +@pytest.mark.anyio +async def test_psqlpy_prepare_parameters_with_casts_zero_allocation() -> None: + """prepare_parameters_with_casts returns original sequence when no coercion or cast is needed.""" + from sqlspec.adapters.psqlpy.core import _EMPTY_PARAMETER_CASTS, prepare_parameters_with_casts + + params = (1, "two", 3.0) + cfg = build_statement_config() + result = prepare_parameters_with_casts(params, _EMPTY_PARAMETER_CASTS, cfg) + assert result is params + + +@pytest.mark.anyio +async def test_psqlpy_execute_cache_hit_direct_fast_path() -> None: + """_execute_cache_hit executes directly on connection without calling _execute_cached_statement.""" + conn = AsyncMock() + conn.fetch = AsyncMock(return_value=[{"col": 1}]) + driver = PsqlpyDriver(connection=cast("Any", conn), statement_config=build_statement_config()) + statement = SQL("SELECT 1 AS col", statement_config=driver.statement_config) + statement.compile() + driver._cache_statement(statement) + cached = driver._stmt_cache.get(statement.raw_sql) + assert cached is not None + + result = await driver._execute_cache_hit("SELECT 1 AS col", (), cached) + assert result.data == [{"col": 1}] + assert conn.fetch.await_count == 1 + + +@pytest.mark.parametrize("query", ["UPDATE items SET qty = 1", "UPDATE items SET qty = ?"]) +async def test_psqlpy_cached_dml_preserves_rowcount(query: str) -> None: + """DML without RETURNING retains its affected-row count after the query is cached.""" + native_result = SimpleNamespace(result=lambda: [{psqlpy_core._DML_COUNT_COLUMN: 3}]) + connection = SimpleNamespace(fetch=AsyncMock(return_value=native_result), execute=AsyncMock()) + driver = PsqlpyDriver(connection=cast("Any", connection)) + parameters = (1,) if "?" in query else () + + results = [await driver.execute(query, *parameters) for _ in range(2)] + + assert [result.rows_affected for result in results] == [3, 3] + assert connection.fetch.await_count == 2 + + +async def test_psqlpy_cached_query_maps_native_errors() -> None: + """A warmed cache retains the same exception translation and cause as ordinary execution.""" + from psqlpy.exceptions import ConnectionExecuteError + + native_result = SimpleNamespace(result=lambda: [{"value": 1}]) + native_error = ConnectionExecuteError("duplicate key value violates unique constraint items_pk") + connection = SimpleNamespace(fetch=AsyncMock(side_effect=[native_result, native_error])) + driver = PsqlpyDriver(connection=cast("Any", connection)) + + assert (await driver.execute("SELECT 1 AS value")).get_data() == [{"value": 1}] + with pytest.raises(UniqueViolationError) as error: + await driver.execute("SELECT 1 AS value") + assert error.value.__cause__ is native_error + + +@pytest.mark.anyio +@pytest.mark.parametrize("binary_copy", [False, True], ids=["records", "binary"]) +async def test_psqlpy_load_from_arrow_stops_after_first_copy_failure(binary_copy: bool) -> None: + """A failed COPY stops the batch stream and preserves the first native cause.""" + import pyarrow as pa + from psqlpy.exceptions import ConnectionExecuteError + + first_error = ConnectionExecuteError("duplicate key value violates unique constraint items_pk") + later_error = ConnectionExecuteError("later copy network drop") + copy = AsyncMock(side_effect=[first_error, later_error]) + conn = SimpleNamespace( + binary_copy_to_table=copy, + copy_records_to_table=copy, + fetch=AsyncMock(return_value=SimpleNamespace(result=lambda: [])), + ) + driver = PsqlpyDriver( + connection=cast("Any", conn), + statement_config=build_statement_config(), + driver_features={"enable_pgcopy_binary": binary_copy, "storage_capabilities": {"arrow_import_enabled": True}}, + ) + batches = [pa.record_batch({"id": [value]}) for value in (1, 2)] + reader = pa.RecordBatchReader.from_batches(batches[0].schema, batches) + + with pytest.raises(UniqueViolationError) as exc_info: + await driver.load_from_arrow("items", reader) + + assert copy.await_count == 1 + assert exc_info.value.__cause__ is first_error + + +def test_psqlpy_driver_class_cached_in_resolve() -> None: + """_resolve_psqlpy_driver caches the resolved driver class.""" + from sqlspec.adapters.psqlpy._typing import _resolve_psqlpy_driver + + cls1 = _resolve_psqlpy_driver() + cls2 = _resolve_psqlpy_driver() + assert cls1 is cls2 + assert cls1 is PsqlpyDriver diff --git a/tests/unit/adapters/test_psqlpy/test_transaction_state.py b/tests/unit/adapters/test_psqlpy/test_transaction_state.py index 9c48ea87f..6faedab9f 100644 --- a/tests/unit/adapters/test_psqlpy/test_transaction_state.py +++ b/tests/unit/adapters/test_psqlpy/test_transaction_state.py @@ -7,7 +7,7 @@ from sqlspec.adapters.psqlpy._typing import PsqlpyDatabaseError from sqlspec.adapters.psqlpy.config import PsqlpyConfig -from sqlspec.adapters.psqlpy.core import PsqlpyStreamSource +from sqlspec.adapters.psqlpy.core import PsqlpyStreamSource, build_postgres_storage_params_sql from sqlspec.adapters.psqlpy.driver import PsqlpyDriver from sqlspec.exceptions import SQLSpecError @@ -205,3 +205,112 @@ async def test_load_from_arrow_decodes_json_text_for_json_columns() -> None: _table_name, records, _kwargs = connection.copy_calls[0] assert records == [(1, {"name": "alpha"}, '{"not": "json"}')] + + +def test_psqlpy_build_postgres_storage_params_sql() -> None: + """build_postgres_storage_params_sql consolidates storage and autovacuum clauses.""" + settings = {"fillfactor": 85, "autovacuum_vacuum_scale_factor": 0.05, "autovacuum_analyze_scale_factor": 0.02} + clause = build_postgres_storage_params_sql(settings, include_autovacuum=True) + assert "WITH (" in clause + assert "fillfactor = 85" in clause + assert "autovacuum_vacuum_scale_factor = 0.05" in clause + assert "autovacuum_analyze_scale_factor = 0.02" in clause + + +def test_psqlpy_events_validate_listener_pool_capacity_skips_when_direct_dsn() -> None: + """_validate_listener_pool_capacity should not enforce pool size when direct_dsn is configured.""" + from sqlspec.adapters.psqlpy.events.backend import _validate_listener_pool_capacity + + config = PsqlpyConfig( + connection_config={"dsn": "postgresql://localhost/test", "max_db_pool_size": 1}, + driver_features={"direct_dsn": "postgresql://localhost/direct"}, + ) + _validate_listener_pool_capacity(config) + + +def test_psqlpy_events_parse_notify_payload_handles_malformed_json() -> None: + """Malformed NOTIFY payload does not raise and returns None safely.""" + from sqlspec.adapters.psqlpy.events.backend import ( + _extract_batch_size, + _extract_event_id, + _extract_marker_id, + _parse_notify_payload, + ) + + bad_payload = "not valid json {{" + assert _parse_notify_payload(bad_payload) is None + assert _extract_event_id(bad_payload) is None + assert _extract_batch_size(bad_payload) == 0 + assert _extract_marker_id(bad_payload) is None + + good_payload = '{"event_id": "e1", "batch_size": 5, "marker_id": "m1"}' + parsed = _parse_notify_payload(good_payload) + assert parsed is not None + assert _extract_event_id(parsed) == "e1" + assert _extract_batch_size(parsed) == 5 + assert _extract_marker_id(parsed) == "m1" + + +def test_psqlpy_events_resolve_listener_config_allows_mapping_direct_dsn() -> None: + """_resolve_listener_config accepts Mapping for direct_dsn without raising.""" + from sqlspec.adapters.psqlpy.events._hub import PsqlpyListenerHub + + config = PsqlpyConfig( + connection_config={"dsn": "postgresql://localhost/test"}, + driver_features={"pooler_mode": "transaction", "direct_dsn": {"host": "localhost", "database": "direct_db"}}, + ) + hub = PsqlpyListenerHub(config) + resolved = hub._resolve_listener_config() + assert resolved is not None + + +def test_psqlpy_events_backend_logger_namespace() -> None: + """Psqlpy events backend logger namespace follows standard naming.""" + from sqlspec.adapters.psqlpy.events.backend import logger + + assert logger.name == "sqlspec.adapters.psqlpy.events.backend" + + +async def test_psqlpy_litestar_store_operations_handle_none_query_result() -> None: + """exists, delete, delete_expired handle None query result safely.""" + from unittest.mock import AsyncMock + + from sqlspec.adapters.psqlpy.litestar.store import PsqlpyStore + + mock_conn = AsyncMock() + mock_conn.fetch = AsyncMock(return_value=None) + mock_conn.execute = AsyncMock(return_value=None) + + config = PsqlpyConfig(connection_config={"dsn": "postgresql://localhost/test"}) + store = PsqlpyStore(config) + object.__setattr__( + config, + "provide_connection", + lambda: AsyncMock(__aenter__=AsyncMock(return_value=mock_conn), __aexit__=AsyncMock()), + ) + + assert await store.exists("sess_1") is False + assert await store.expires_in("sess_1") is None + assert await store.delete_expired() == 0 + + +async def test_psqlpy_get_indexes_with_none_schema_resolves_via_search_path() -> None: + """get_indexes passes schema_name=None when schema is None to resolve table on search_path.""" + from unittest.mock import AsyncMock, MagicMock + + from sqlspec.adapters.psqlpy.data_dictionary import PsqlpyDataDictionary + + calls: list[dict[str, Any]] = [] + driver = MagicMock() + + async def fake_select(query: Any, **kwargs: Any) -> list[Any]: + calls.append(kwargs) + return [] + + driver.select = AsyncMock(side_effect=fake_select) + + dd = PsqlpyDataDictionary() + await dd.get_indexes(driver, table="users", schema=None) + assert len(calls) == 1 + assert calls[0]["schema_name"] is None + assert calls[0]["table_name"] == "users" diff --git a/tests/unit/adapters/test_psycopg/test_config.py b/tests/unit/adapters/test_psycopg/test_config.py index af34e66a6..4c592726f 100644 --- a/tests/unit/adapters/test_psycopg/test_config.py +++ b/tests/unit/adapters/test_psycopg/test_config.py @@ -13,18 +13,16 @@ from sqlspec.adapters.psycopg._typing import PsycopgAsyncSessionContext, PsycopgSyncSessionContext from sqlspec.adapters.psycopg.config import ( PsycopgAsyncConfig, + PsycopgConnectionParams, + PsycopgDriverFeatures, PsycopgPoolParams, PsycopgSyncConfig, build_connection_config, ) -from sqlspec.adapters.psycopg.core import ( - build_postgres_extension_probe_names, - build_statement_config, - default_statement_config, - resolve_postgres_extension_state, -) +from sqlspec.adapters.psycopg.core import build_statement_config, default_statement_config from sqlspec.adapters.psycopg.driver import PsycopgAsyncDriver, PsycopgSyncDriver from sqlspec.core import SQL, StatementConfig +from sqlspec.core.config_runtime import build_postgres_extension_probe_names, resolve_postgres_extension_state from sqlspec.exceptions import ImproperConfigurationError @@ -631,13 +629,14 @@ def test_async_config_normalizes_aliases() -> None: assert "db" not in config.connection_config -def test_psycopg_sync_configure_connection_supports_dict_row_and_rolls_back_after_hook() -> None: - """_configure_connection should handle dict_row probe results and roll back after user hooks when autocommit=False.""" +def test_psycopg_sync_configure_connection_commits_hook_after_extension_probe() -> None: + """The probe and SQL hook both leave a non-autocommit connection idle.""" events: list[str] = [] connection = MagicMock() connection.autocommit = False connection.execute.return_value.fetchall.return_value = [{"extname": "pg_textsearch"}] connection.rollback.side_effect = lambda: events.append("rollback") + connection.commit.side_effect = lambda: events.append("commit") config = PsycopgSyncConfig( driver_features={ @@ -652,12 +651,12 @@ def test_psycopg_sync_configure_connection_supports_dict_row_and_rolls_back_afte assert config.pg_textsearch_available is True assert config.statement_config.dialect == "pg_textsearch" - assert events == ["hook", "rollback"] + assert events == ["rollback", "hook", "commit"] @pytest.mark.anyio -async def test_psycopg_async_configure_connection_supports_dict_row_and_rolls_back_after_hook() -> None: - """_configure_async_connection should handle dict_row probe results and roll back after user hooks when autocommit=False.""" +async def test_psycopg_async_configure_connection_commits_hook_after_extension_probe() -> None: + """The probe and SQL hook both leave a non-autocommit connection idle.""" events: list[str] = [] connection = AsyncMock() connection.autocommit = False @@ -672,6 +671,7 @@ async def record_hook(_: object) -> None: events.append("hook") connection.rollback = AsyncMock(side_effect=record_rollback) + connection.commit = AsyncMock(side_effect=lambda: events.append("commit")) config = PsycopgAsyncConfig( driver_features={ @@ -686,7 +686,57 @@ async def record_hook(_: object) -> None: assert config.pg_textsearch_available is True assert config.statement_config.dialect == "pg_textsearch" - assert events == ["hook", "rollback"] + assert events == ["rollback", "hook", "commit"] + + +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.parametrize("failure_phase", ["hook", "commit"]) +async def test_failed_connection_initialization_rolls_back(is_async: bool, failure_phase: str) -> None: + """A failed hook or setup commit rolls its transaction back before propagating the error.""" + connection = AsyncMock() if is_async else MagicMock() + connection.autocommit = False + native_error = RuntimeError("initialization failed") + hook = AsyncMock() if is_async else MagicMock() + if failure_phase == "hook": + hook.side_effect = native_error + else: + connection.commit.side_effect = native_error + config = ( + PsycopgAsyncConfig(driver_features={"on_connection_create": hook}) + if is_async + else PsycopgSyncConfig(driver_features={"on_connection_create": hook}) + ) + config._pgvector_available = False # pyright: ignore[reportPrivateUsage] + + with pytest.raises(RuntimeError, match="initialization failed"): + if isinstance(config, PsycopgAsyncConfig): + await config._configure_async_connection(connection) # pyright: ignore[reportPrivateUsage] + else: + config._configure_connection(connection) # pyright: ignore[reportPrivateUsage] + assert connection.rollback.call_count == 1 + + +@pytest.mark.parametrize("is_async", [False, True]) +async def test_direct_config_preserves_connection_initialization_hook(is_async: bool) -> None: + """Bypassing a pooler keeps the configured physical-connection initialization.""" + hook = AsyncMock() if is_async else MagicMock() + config = ( + PsycopgAsyncConfig( + driver_features={"direct_dsn": "postgresql://localhost/direct", "on_connection_create": hook} + ) + if is_async + else PsycopgSyncConfig( + driver_features={"direct_dsn": "postgresql://localhost/direct", "on_connection_create": hook} + ) + ) + direct = config.derive_direct_config() + connection = AsyncMock() if is_async else MagicMock() + connection.autocommit = True + if isinstance(direct, PsycopgAsyncConfig): + await direct._configure_async_connection(connection) # pyright: ignore[reportPrivateUsage] + else: + direct._configure_connection(connection) # pyright: ignore[reportPrivateUsage] + hook.assert_called_once_with(connection) def test_psycopg_pooler_mode_transaction_sets_prepare_threshold_none_and_alloydb_class( @@ -705,11 +755,11 @@ def test_psycopg_pooler_mode_transaction_sets_prepare_threshold_none_and_alloydb assert async_config.connection_config["prepare_threshold"] is None assert async_config.driver_features["pooler_mode"] == "statement" - alias_config = PsycopgSyncConfig( - connection_config={"conninfo": "postgresql://localhost/db"}, driver_features={"pgbouncer": True} - ) - assert alias_config.connection_config["prepare_threshold"] is None - assert alias_config.driver_features["pooler_mode"] == "transaction" + assert "pgbouncer" not in PsycopgDriverFeatures.__annotations__ + assert "pgbouncer_mode" not in PsycopgDriverFeatures.__annotations__ + assert "null_pool" not in PsycopgDriverFeatures.__annotations__ + assert "null_pool" not in PsycopgConnectionParams.__annotations__ + assert "enable_pooling" in PsycopgDriverFeatures.__annotations__ mock_connector = MagicMock() expected_conn = MagicMock() @@ -754,15 +804,13 @@ async def test_psycopg_enable_alloydb_columnar_and_derive_direct_config() -> Non ) assert "-c google_columnar_engine.enable_columnar_scan=on" in auto_sync.connection_config["options"] - auto_pgbouncer_sync = PsycopgSyncConfig( + tx_sync = PsycopgSyncConfig( connection_config={"conninfo": "postgresql://pooler:6432/db"}, - driver_features={"enable_alloydb_columnar": True, "pooler_mode": "auto", "pgbouncer": True}, + driver_features={"enable_alloydb_columnar": True, "pooler_mode": "transaction"}, ) - assert auto_pgbouncer_sync.driver_features["pooler_mode"] == "transaction" - assert auto_pgbouncer_sync.driver_features["pgbouncer"] is True - assert auto_pgbouncer_sync.driver_features["pgbouncer_mode"] is True - assert "options" not in auto_pgbouncer_sync.connection_config - assert auto_pgbouncer_sync.connection_config["prepare_threshold"] is None + assert tx_sync.driver_features["pooler_mode"] == "transaction" + assert "options" not in tx_sync.connection_config + assert tx_sync.connection_config["prepare_threshold"] is None with pytest.raises(ImproperConfigurationError, match="Invalid pooler_mode"): PsycopgSyncConfig( @@ -828,3 +876,23 @@ async def test_psycopg_enable_alloydb_columnar_and_derive_direct_config() -> Non ) await async_driver.begin() async_conn.execute.assert_awaited_once_with("SET LOCAL google_columnar_engine.enable_columnar_scan = 'on'") + + +def test_psycopg_driver_features_enable_pgcopy_binary() -> None: + """PsycopgDriverFeatures accepts enable_pgcopy_binary and apply_driver_features provides default.""" + from sqlspec.adapters.psycopg.config import PsycopgDriverFeatures + from sqlspec.adapters.psycopg.core import apply_driver_features + + features: PsycopgDriverFeatures = {"enable_pgcopy_binary": True} + assert features["enable_pgcopy_binary"] is True + + _, resolved = apply_driver_features(StatementConfig(), {}) + assert resolved["enable_pgcopy_binary"] is False + + +def test_psycopg_connection_type_aliases() -> None: + """PsycopgSyncConnection and PsycopgAsyncConnection type aliases allow generic connections.""" + from sqlspec.adapters.psycopg._typing import PsycopgAsyncConnection, PsycopgSyncConnection + + assert PsycopgSyncConnection is not None + assert PsycopgAsyncConnection is not None diff --git a/tests/unit/adapters/test_psycopg/test_driver.py b/tests/unit/adapters/test_psycopg/test_driver.py index 4cdf238c7..b4620935b 100644 --- a/tests/unit/adapters/test_psycopg/test_driver.py +++ b/tests/unit/adapters/test_psycopg/test_driver.py @@ -1,9 +1,10 @@ """Unit tests for psycopg driver transaction behavior.""" +import inspect from asyncio import CancelledError from contextlib import nullcontext from types import SimpleNamespace -from typing import TYPE_CHECKING, cast +from typing import TYPE_CHECKING, Any, cast from unittest.mock import AsyncMock, MagicMock import psycopg @@ -12,7 +13,7 @@ from sqlspec.adapters.psycopg.config import PsycopgAsyncConfig, PsycopgSyncConfig from sqlspec.adapters.psycopg.driver import PsycopgAsyncDriver, PsycopgSyncDriver from sqlspec.core import SQL -from sqlspec.exceptions import SerializationConflictError, SQLSpecError +from sqlspec.exceptions import SerializationConflictError, SQLSpecError, UniqueViolationError if TYPE_CHECKING: from sqlspec.adapters.psycopg._typing import PsycopgAsyncConnection, PsycopgSyncConnection @@ -119,6 +120,7 @@ def commit(self) -> None: def rollback(self) -> None: self.rollback_calls += 1 + self.info.transaction_status = 0 if self.rollback_error is not None: raise self.rollback_error @@ -158,6 +160,7 @@ async def commit(self) -> None: async def rollback(self) -> None: self.rollback_calls += 1 + self.info.transaction_status = 0 if self.rollback_error is not None: raise self.rollback_error @@ -255,9 +258,9 @@ def test_sync_rollback_restores_changed_autocommit() -> None: assert connection.autocommit_calls == [False, True] assert connection.autocommit is True - assert driver._connection_in_transaction() is False # pyright: ignore[reportPrivateUsage] +@pytest.mark.anyio async def test_async_rollback_restores_changed_autocommit() -> None: connection = _AsyncTransactionConnection() driver = _async_driver(connection) @@ -267,6 +270,59 @@ async def test_async_rollback_restores_changed_autocommit() -> None: assert connection.autocommit_calls == [False, True] assert connection.autocommit is True + + +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.parametrize( + ("query", "parameters", "wire_sql"), + [ + ("SELECT '100%' AS value", (), "SELECT '100%' AS value"), + ("SELECT 10 % 3 AS value", (), "SELECT 10 % 3 AS value"), + ("SELECT %s AS value WHERE '100%' LIKE '100%'", (1,), "SELECT %s AS value WHERE '100%%' LIKE '100%%'"), + ], +) +async def test_cached_percent_sql_preserves_native_binding( + is_async: bool, query: str, parameters: tuple[int, ...], wire_sql: str +) -> None: + """Cached and uncached execution apply percent escaping only with native parameters.""" + cursor = MagicMock() + cursor.description = [SimpleNamespace(name="value", type_code=23)] + cursor.rowcount = 1 + if is_async: + cursor.execute = AsyncMock() + cursor.fetchall = AsyncMock(return_value=[(1,)]) + cursor.close = AsyncMock() + else: + cursor.fetchall.return_value = [(1,)] + connection = MagicMock() + connection.cursor.return_value = cursor + driver = PsycopgAsyncDriver(connection) if is_async else PsycopgSyncDriver(connection) + + for _ in range(2): + result = driver.execute(query, *parameters) + if inspect.isawaitable(result): + result = await result + assert result.get_data() == [{"value": 1}] + + assert [call.args[0] for call in cursor.execute.call_args_list] == [wire_sql, wire_sql] + bound = [tuple(call.args[1]) if len(call.args) > 1 else () for call in cursor.execute.call_args_list] + assert bound == [parameters, parameters] + + +def test_psycopg_sync_stream_source_runs_with_exception_handler() -> None: + """PsycopgSyncStreamSource delegates to _run_with_exception_handler inside active except block.""" + from sqlspec.adapters.psycopg.core import PsycopgSyncStreamSource + + connection = MagicMock() + driver = PsycopgSyncDriver(connection=connection) + source = PsycopgSyncStreamSource(driver, "SELECT 1", (), 10) + + try: + raise ValueError("ambient error") + except ValueError: + source.start() + + assert source._cursor is not None assert driver._connection_in_transaction() is False # pyright: ignore[reportPrivateUsage] @@ -409,11 +465,11 @@ async def test_async_logical_transaction_is_reported_while_libpq_is_idle() -> No def test_sync_inerror_status_is_not_active_transaction_and_rolls_back_on_begin() -> None: - """TRANSACTION_STATUS_INERROR (3) should not report an active transaction and should roll back on begin().""" + """TRANSACTION_STATUS_INERROR (3) is reported in transaction state and rolls back on begin().""" connection = _SyncTransactionConnection(transaction_status=3) driver = _sync_driver(connection) - assert driver._connection_in_transaction() is False # pyright: ignore[reportPrivateUsage] + assert driver._connection_in_transaction() is True # pyright: ignore[reportPrivateUsage] driver.begin() assert connection.rollback_calls == 1 @@ -422,11 +478,11 @@ def test_sync_inerror_status_is_not_active_transaction_and_rolls_back_on_begin() async def test_async_inerror_status_is_not_active_transaction_and_rolls_back_on_begin() -> None: - """TRANSACTION_STATUS_INERROR (3) should not report an active transaction and should roll back on async begin().""" + """TRANSACTION_STATUS_INERROR (3) is reported in transaction state and rolls back on async begin().""" connection = _AsyncTransactionConnection(transaction_status=3) driver = _async_driver(connection) - assert driver._connection_in_transaction() is False # pyright: ignore[reportPrivateUsage] + assert driver._connection_in_transaction() is True # pyright: ignore[reportPrivateUsage] await driver.begin() assert connection.rollback_calls == 1 @@ -460,3 +516,158 @@ def test_sync_copy_to_stdout_decodes_memoryview_chunks_and_reports_rowcount() -> assert result is not None assert result.get_data() == [{"copy_output": "1\talice\n2\tbob\n"}] assert result.rows_affected == 2 + + +def test_sync_session_exit_rolls_back_uncommitted_implicit_transaction() -> None: + """Session exit rolls back uncommitted implicit autocommit=False transactions.""" + connection = _SyncTransactionConnection(transaction_status=2) + config = PsycopgSyncConfig() + context = config.provide_session() + object.__setattr__(context, "_acquire_connection", lambda: connection) + with context: + pass + assert connection.rollback_calls == 1 + + +@pytest.mark.anyio +async def test_async_session_exit_rolls_back_uncommitted_implicit_transaction() -> None: + """Async session exit rolls back uncommitted implicit autocommit=False transactions.""" + connection = _AsyncTransactionConnection(transaction_status=2) + config = PsycopgAsyncConfig() + context = config.provide_session() + object.__setattr__(context, "_acquire_connection", AsyncMock(return_value=connection)) + async with context: + pass + assert connection.rollback_calls == 1 + + +@pytest.mark.anyio +async def test_psycopg_async_load_from_arrow_runs_with_exception_handler() -> None: + """PsycopgAsyncDriver.load_from_arrow routes operations through exception handler inside active except block.""" + import pyarrow as pa + + connection = MagicMock() + cursor_mock = MagicMock() + cursor_mock.execute = AsyncMock() + copy_ctx = AsyncMock() + cursor_mock.copy.return_value = copy_ctx + connection.cursor.return_value = cursor_mock + + driver = PsycopgAsyncDriver( + connection=connection, driver_features={"storage_capabilities": {"arrow_import_enabled": True}} + ) + table = pa.table({"id": [1, 2], "val": ["a", "b"]}) + + try: + raise ValueError("ambient error") + except ValueError: + job = await driver.load_from_arrow("test_table", table, overwrite=True) + + assert job.telemetry["rows_processed"] == 2 + + +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.parametrize("binary", [False, True]) +@pytest.mark.parametrize("failure_phase", ["enter", "exit"]) +async def test_copy_lifecycle_maps_native_errors(is_async: bool, binary: bool, failure_phase: str) -> None: + """COPY setup and finalization errors retain SQLSpec mapping even inside an except block.""" + import pyarrow as pa + + native_error = psycopg.errors.UniqueViolation("duplicate key value violates unique constraint items_pk") + copy_context = AsyncMock() if is_async else MagicMock() + enter = copy_context.__aenter__ if is_async else copy_context.__enter__ + exit_ = copy_context.__aexit__ if is_async else copy_context.__exit__ + enter.return_value = copy_context + if failure_phase == "enter": + enter.side_effect = native_error + else: + exit_.side_effect = native_error + cursor = MagicMock() + cursor.copy.return_value = copy_context + if is_async: + cursor.close = AsyncMock() + connection = MagicMock() + connection.cursor.return_value = cursor + features = {"storage_capabilities": {"arrow_import_enabled": True}, "enable_pgcopy_binary": binary} + driver = ( + PsycopgAsyncDriver(connection, driver_features=features) + if is_async + else PsycopgSyncDriver(connection, driver_features=features) + ) + + try: + raise ValueError("ambient exception") + except ValueError: + with pytest.raises(UniqueViolationError) as error: + result = driver.load_from_arrow("items", pa.table({"id": [1, 1]})) + if inspect.isawaitable(result): + await result + assert error.value.__cause__ is native_error + + +def test_sync_dispatch_execute_script_emits_statement_observer_events() -> None: + """PsycopgSyncDriver.dispatch_execute_script emits script statement events to observers.""" + from sqlspec.core import SQL + from sqlspec.observability import ObservabilityConfig, ObservabilityRuntime + + events: list[Any] = [] + connection = MagicMock() + cursor = MagicMock() + connection.cursor.return_value = cursor + + runtime = ObservabilityRuntime(ObservabilityConfig(statement_observers=(events.append,))) + driver = PsycopgSyncDriver(connection=connection) + driver._observability = runtime + + stmt = SQL("SELECT 1; SELECT 2;") + result = driver.dispatch_execute_script(cursor, stmt) + + assert result.statement_count == 2 + assert len(events) == 2 + assert events[0].script_ordinal == 0 + assert events[1].script_ordinal == 1 + + +@pytest.mark.anyio +async def test_async_dispatch_execute_script_emits_statement_observer_events() -> None: + """PsycopgAsyncDriver.dispatch_execute_script emits script statement events to observers.""" + from sqlspec.core import SQL + from sqlspec.observability import ObservabilityConfig, ObservabilityRuntime + + events: list[Any] = [] + connection = MagicMock() + cursor = AsyncMock() + connection.cursor.return_value = cursor + + runtime = ObservabilityRuntime(ObservabilityConfig(statement_observers=(events.append,))) + driver = PsycopgAsyncDriver(connection=connection) + driver._observability = runtime + + stmt = SQL("SELECT 1; SELECT 2;") + result = await driver.dispatch_execute_script(cursor, stmt) + + assert result.statement_count == 2 + assert len(events) == 2 + assert events[0].script_ordinal == 0 + assert events[1].script_ordinal == 1 + + +def test_sync_load_from_arrow_preserves_cause_on_failure() -> None: + """PsycopgSyncDriver.load_from_arrow preserves underlying __cause__ on execution failure.""" + import psycopg + import pyarrow as pa + + connection = MagicMock() + cursor_mock = MagicMock() + cursor_mock.execute.side_effect = psycopg.OperationalError("db error") + connection.cursor.return_value = cursor_mock + + driver = PsycopgSyncDriver( + connection=connection, driver_features={"storage_capabilities": {"arrow_import_enabled": True}} + ) + table = pa.table({"id": [1, 2]}) + + with pytest.raises(Exception) as exc_info: + driver.load_from_arrow("test_table", table, overwrite=True) + + assert exc_info.value.__cause__ is not None diff --git a/tests/unit/adapters/test_psycopg/test_type_handlers.py b/tests/unit/adapters/test_psycopg/test_type_handlers.py index 9d959e557..881b956fe 100644 --- a/tests/unit/adapters/test_psycopg/test_type_handlers.py +++ b/tests/unit/adapters/test_psycopg/test_type_handlers.py @@ -12,7 +12,7 @@ def test_register_pgvector_sync_with_pgvector_installed(monkeypatch: pytest.Monk mock_connection = MagicMock() mock_pgvector = MagicMock() - monkeypatch.setattr(type_converter, "_pgvector_psycopg", mock_pgvector) + monkeypatch.setattr(type_converter, "_pgvector_psycopg_module", lambda: mock_pgvector) register_pgvector_sync(mock_connection) @@ -23,7 +23,7 @@ def test_register_pgvector_sync_without_pgvector(monkeypatch: pytest.MonkeyPatch """Test register_pgvector_sync gracefully handles pgvector not installed.""" import sqlspec.adapters.psycopg.type_converter - monkeypatch.setattr(sqlspec.adapters.psycopg.type_converter, "_pgvector_psycopg", None) + monkeypatch.setattr(sqlspec.adapters.psycopg.type_converter, "_pgvector_psycopg_module", lambda: None) from sqlspec.adapters.psycopg.type_converter import register_pgvector_sync @@ -41,7 +41,7 @@ async def test_register_pgvector_async_with_pgvector_installed(monkeypatch: pyte mock_connection = AsyncMock() mock_pgvector = MagicMock() mock_pgvector.register_vector_async = AsyncMock() - monkeypatch.setattr(type_converter, "_pgvector_psycopg", mock_pgvector) + monkeypatch.setattr(type_converter, "_pgvector_psycopg_module", lambda: mock_pgvector) await register_pgvector_async(mock_connection) @@ -52,7 +52,7 @@ async def test_register_pgvector_async_without_pgvector(monkeypatch: pytest.Monk """Test register_pgvector_async gracefully handles pgvector not installed.""" import sqlspec.adapters.psycopg.type_converter - monkeypatch.setattr(sqlspec.adapters.psycopg.type_converter, "_pgvector_psycopg", None) + monkeypatch.setattr(sqlspec.adapters.psycopg.type_converter, "_pgvector_psycopg_module", lambda: None) from sqlspec.adapters.psycopg.type_converter import register_pgvector_async @@ -70,7 +70,7 @@ def test_register_pgvector_sync_handles_registration_failure(monkeypatch: pytest mock_connection = MagicMock() mock_pgvector = MagicMock() mock_pgvector.register_vector.side_effect = Exception("Registration failed") - monkeypatch.setattr(type_converter, "_pgvector_psycopg", mock_pgvector) + monkeypatch.setattr(type_converter, "_pgvector_psycopg_module", lambda: mock_pgvector) register_pgvector_sync(mock_connection) @@ -83,13 +83,57 @@ async def test_register_pgvector_async_handles_registration_failure(monkeypatch: mock_connection = AsyncMock() mock_pgvector = MagicMock() mock_pgvector.register_vector_async = AsyncMock(side_effect=Exception("Registration failed")) - monkeypatch.setattr(type_converter, "_pgvector_psycopg", mock_pgvector) + monkeypatch.setattr(type_converter, "_pgvector_psycopg_module", lambda: mock_pgvector) await register_pgvector_async(mock_connection) -def test_pgvector_module_reference_is_cached() -> None: - import sqlspec.adapters.psycopg.type_converter as type_converter +def test_operational_error_with_none_sqlstate_maps_to_database_connection_error() -> None: + """OperationalError with sqlstate=None maps to DatabaseConnectionError.""" + import psycopg + + from sqlspec.adapters.psycopg.core import create_mapped_exception + from sqlspec.exceptions import DatabaseConnectionError + + err = psycopg.OperationalError("connection failed") + mapped = create_mapped_exception(err) + assert type(mapped) is DatabaseConnectionError + + +def test_lock_not_available_maps_to_query_timeout_error() -> None: + """pg_errors.LockNotAvailable (55P03) maps to QueryTimeoutError.""" + from psycopg import errors as pg_errors + + from sqlspec.adapters.psycopg.core import create_mapped_exception + from sqlspec.exceptions import QueryTimeoutError + + err = pg_errors.LockNotAvailable("could not obtain lock on relation") + mapped = create_mapped_exception(err) + assert type(mapped) is QueryTimeoutError + + +def test_sqlstate_class_25_maps_to_transaction_error() -> None: + """Class 25 SQLSTATE maps to TransactionError.""" + from psycopg import OperationalError + + from sqlspec.adapters.psycopg.core import create_mapped_exception + from sqlspec.exceptions import TransactionError + + err = OperationalError("active sql transaction") + err.sqlstate = "25000" + mapped = create_mapped_exception(err) + assert type(mapped) is TransactionError + + +def test_exclusion_violation_23p01_preserves_sqlstate() -> None: + """ExclusionViolation (23P01) preserves its exact SQLSTATE code in error message.""" + from psycopg import errors as pg_errors + + from sqlspec.adapters.psycopg.core import create_mapped_exception + from sqlspec.exceptions import IntegrityError - assert not hasattr(type_converter, "importlib") - assert hasattr(type_converter, "_pgvector_psycopg") + err = pg_errors.ExclusionViolation("conflicting key value violates exclusion constraint") + err.sqlstate = "23P01" + mapped = create_mapped_exception(err) + assert isinstance(mapped, IntegrityError) + assert "[23P01]" in str(mapped) diff --git a/tests/unit/adapters/test_pymssql/_fakes.py b/tests/unit/adapters/test_pymssql/_fakes.py index 8fcd4ec1d..b719ccc5c 100644 --- a/tests/unit/adapters/test_pymssql/_fakes.py +++ b/tests/unit/adapters/test_pymssql/_fakes.py @@ -19,11 +19,16 @@ class FakeCursor: """Minimal DB-API cursor for pymssql unit tests.""" def __init__( - self, rows: "list[Any] | None" = None, description: "list[tuple[str, ...]] | None" = None, rowcount: int = -1 + self, + rows: "list[Any] | None" = None, + description: "list[tuple[str, ...]] | None" = None, + rowcount: int = -1, + error: BaseException | None = None, ) -> None: self.rows = rows or [] self.description = description self.rowcount = rowcount + self.error = error self.closed = False self.calls: list[tuple[str, Any]] = [] self.fetchmany_sizes: list[int] = [] @@ -31,6 +36,8 @@ def __init__( def execute(self, sql: str, parameters: Any = None) -> None: self.calls.append((sql, parameters)) + if self.error is not None: + raise self.error def executemany(self, sql: str, parameters: Any = None) -> None: self.many_calls.append((sql, parameters)) diff --git a/tests/unit/adapters/test_pymssql/test_config.py b/tests/unit/adapters/test_pymssql/test_config.py index 07a9e0a49..d4593b23f 100644 --- a/tests/unit/adapters/test_pymssql/test_config.py +++ b/tests/unit/adapters/test_pymssql/test_config.py @@ -5,7 +5,7 @@ import pytest from sqlspec.adapters.pymssql import build_connection_config -from sqlspec.adapters.pymssql.config import PymssqlConfig, PymssqlConnectionParams +from sqlspec.adapters.pymssql.config import PymssqlConfig, PymssqlConnectionParams, PymssqlDriverFeatures from sqlspec.adapters.pymssql.driver import PymssqlDriver from sqlspec.adapters.pymssql.pool import PymssqlConnectionPool @@ -158,3 +158,19 @@ def test_pymssql_session_clean_exit_rolls_back_active_transaction() -> None: assert conn.rollbacks == 1 assert calls == ["release"] + + +def test_use_datetime2_default() -> None: + """PymssqlConfig connection_config supports use_datetime2.""" + config = PymssqlConfig(connection_config={"server": "localhost", "use_datetime2": True}) + assert config.connection_config.get("use_datetime2") is True + + +def test_driver_features_mssql_parity() -> None: + """PymssqlConfig defaults enable_pooling to True and harmonizes driver features.""" + assert "use_pool" not in PymssqlDriverFeatures.__annotations__ + assert "enable_pooling" in PymssqlDriverFeatures.__annotations__ + config = PymssqlConfig() + assert config.driver_features.get("enable_pooling") is True + assert config.driver_features.get("events_backend") == "poll_queue" + assert config.driver_features.get("enable_uuid_conversion") is False diff --git a/tests/unit/adapters/test_pymssql/test_core.py b/tests/unit/adapters/test_pymssql/test_core.py index aac791580..da4b88405 100644 --- a/tests/unit/adapters/test_pymssql/test_core.py +++ b/tests/unit/adapters/test_pymssql/test_core.py @@ -16,6 +16,7 @@ format_identifier, normalize_execute_many_parameters, normalize_execute_parameters, + quote_bracket_identifier, ) from sqlspec.core import SQL, ParameterStyle from sqlspec.exceptions import ( @@ -252,3 +253,36 @@ def test_build_connection_config_merges_extra_mapping() -> None: assert cfg["appname"] == "analytics" assert "extra" not in cfg assert "ignored" not in cfg + + +def test_freetds_error_code_mapping() -> None: + """FreeTDS 20000-series codes and InterfaceError map to appropriate SQLSpec errors.""" + + class InterfaceError(Exception): + pass + + assert isinstance(create_mapped_exception(Exception(20003, "DBPROCESS is dead")), QueryTimeoutError) + assert isinstance( + create_mapped_exception(Exception(20002, "Adaptive Server connection failed")), DatabaseConnectionError + ) + assert isinstance(create_mapped_exception(Exception(20009, "Server is unavailable")), DatabaseConnectionError) + assert isinstance(create_mapped_exception(Exception(20047, "Connection closed")), DatabaseConnectionError) + assert isinstance(create_mapped_exception(InterfaceError("Connection broken")), DatabaseConnectionError) + + +def test_percent_escaping_performance() -> None: + """escape_literal_percent correctly handles string literals and placeholders.""" + validator = default_statement_config.parameter_validator + sql = "SELECT [id] FROM [t] WHERE [name] LIKE '50%' AND [cat] = %s" + escaped = escape_literal_percent(sql, ("books",), validator) + assert escaped == "SELECT [id] FROM [t] WHERE [name] LIKE '50%%' AND [cat] = %s" + + +def test_quote_bracket_identifier() -> None: + """quote_bracket_identifier cleanly handles nested quotes and brackets.""" + assert quote_bracket_identifier("users") == "[users]" + assert quote_bracket_identifier("[users]") == "[users]" + assert quote_bracket_identifier('"users"') == "[users]" + assert quote_bracket_identifier('user"name') == '[user"name]' + assert quote_bracket_identifier("[user]]name]") == "[user]]name]" + assert quote_bracket_identifier("user]name") == "[user]]name]" diff --git a/tests/unit/adapters/test_pymssql/test_data_dictionary.py b/tests/unit/adapters/test_pymssql/test_data_dictionary.py index b9c4aafe2..e0ee85f11 100644 --- a/tests/unit/adapters/test_pymssql/test_data_dictionary.py +++ b/tests/unit/adapters/test_pymssql/test_data_dictionary.py @@ -129,3 +129,27 @@ def test_sync_data_dictionary_explicit_schema_skips_connection_lookup() -> None: data_dictionary.get_tables(cast(Any, driver), schema="custom") assert driver.select_calls[0][1]["schema_name"] == "custom" assert len(driver.executed) == 0 + + +def test_data_dictionary_capabilities() -> None: + """The sync data dictionary implements schemas, objects, details, routines, privileges, and filtered dependencies.""" + driver = FakeSyncDriver() + data_dictionary = PymssqlSyncDataDictionary() + + schemas_result = data_dictionary.get_schemas(cast(Any, driver)) + assert schemas_result.domain == "schemas" + + objects_result = data_dictionary.get_objects(cast(Any, driver), schema="sales") + assert objects_result.domain == "objects" + + details_result = data_dictionary.get_table_details(cast(Any, driver), "orders", schema="sales") + assert details_result.domain == "tables" + + routines_result = data_dictionary.get_routines(cast(Any, driver), schema="sales") + assert routines_result.domain == "routines" + + privileges_result = data_dictionary.get_privileges(cast(Any, driver), object_name="orders", schema="sales") + assert privileges_result.domain == "privileges" + + deps_result = data_dictionary.get_dependencies(cast(Any, driver), object_name="orders", schema="sales") + assert deps_result.domain == "dependencies" diff --git a/tests/unit/adapters/test_pymssql/test_driver.py b/tests/unit/adapters/test_pymssql/test_driver.py index 452b462ae..54b117665 100644 --- a/tests/unit/adapters/test_pymssql/test_driver.py +++ b/tests/unit/adapters/test_pymssql/test_driver.py @@ -6,7 +6,9 @@ from pymssql import IntegrityError as PymssqlIntegrityError from sqlspec import StatementStack -from sqlspec.adapters.pymssql._typing import PymssqlConnection, PymssqlRawCursor +from sqlspec.adapters.pymssql._typing import PymssqlConnection, PymssqlRawCursor, PymssqlSessionContext +from sqlspec.adapters.pymssql.core import default_statement_config +from sqlspec.adapters.pymssql.driver import PymssqlDriver from sqlspec.core import SQL from sqlspec.exceptions import SQLSpecError, StackExecutionError, TransactionError, UniqueViolationError from tests.unit.adapters.test_pymssql._fakes import FakeConnection, FakeCursor @@ -89,6 +91,23 @@ def test_dispatch_execute_select_compiles_to_pyformat_and_collects_rows() -> Non assert result.data_row_count == 1 +def test_execute_output_into_does_not_fetch_client_rows() -> None: + """OUTPUT INTO must not fetch a result set that the native cursor does not expose.""" + + class OutputIntoCursor(FakeCursor): + def fetchall(self) -> list[Any]: + msg = "OUTPUT INTO does not expose a client result set" + raise AssertionError(msg) + + cursor = OutputIntoCursor(rowcount=1) + driver = PymssqlDriver(cast("PymssqlConnection", FakeConnection(cursor))) + + result = driver.execute("INSERT INTO dbo.users OUTPUT inserted.id INTO @ids DEFAULT VALUES") + + assert result.rows_affected == 1 + assert result.get_data() == [] + + def test_dispatch_execute_many_uses_executemany_and_rowcount() -> None: """execute_many dispatch should forward batch parameters to pymssql.""" from sqlspec.adapters.pymssql.core import default_statement_config @@ -340,3 +359,69 @@ def test_pymssql_load_from_arrow_truncates_and_inserts_batches() -> None: assert cursor.many_calls == [ ("INSERT INTO [dbo].[metrics] ([id], [label]) VALUES (%s, %s)", [(1, "50%"), (2, "100%")]) ] + + +def test_pymssql_transaction_lifecycle() -> None: + """begin(), commit(), and rollback() manage transaction active flags and are idempotent.""" + cursor = FakeCursor() + conn = FakeConnection(cursor) + conn.autocommit_state = True + driver = PymssqlDriver(cast("Any", conn)) + + driver.begin() + assert driver._connection_in_transaction() is True + driver.begin() + assert driver._connection_in_transaction() is True + + driver.commit() + assert driver._connection_in_transaction() is False + driver.commit() + assert driver._connection_in_transaction() is False + + driver.begin() + assert driver._connection_in_transaction() is True + driver.rollback() + assert driver._connection_in_transaction() is False + driver.rollback() + assert driver._connection_in_transaction() is False + + +def test_cache_hit_fast_path_percent_escape() -> None: + """Parameterized queries without literal percents take the cache-hit fast path.""" + from types import SimpleNamespace + + from sqlspec.core.statement import ProcessedState + + cursor = FakeCursor(description=[("id",), ("name",)]) + cursor.rows = [(1, "Ada")] + conn = FakeConnection(cursor) + driver = PymssqlDriver(cast("Any", conn)) + cached = SimpleNamespace( + compiled_sql="SELECT [id], [name] FROM [dbo].[users] WHERE [id] = %s", + expression=None, + operation_type="SELECT", + operation_profile=SimpleNamespace(returns_rows=True, modifies_rows=False), + input_named_parameters=(), + processed_state=ProcessedState("", ()), + ) + + result = driver._execute_cache_hit( + "SELECT [id], [name] FROM [dbo].[users] WHERE [id] = %s", (1,), cast("Any", cached) + ) + assert result.get_data() == [{"id": 1, "name": "Ada"}] + + +def test_session_context_initialization() -> None: + """PymssqlSessionContext accepts explicit driver_cls factory without dynamic import.""" + cursor = FakeCursor() + conn = FakeConnection(cursor) + context = PymssqlSessionContext( + acquire_connection=lambda: conn, + release_connection=lambda *args, **kwargs: None, + statement_config=default_statement_config, + driver_features={}, + prepare_driver=lambda d: d, + driver_cls=PymssqlDriver, + ) + with context as driver: + assert isinstance(driver, PymssqlDriver) diff --git a/tests/unit/adapters/test_pymssql/test_extensions.py b/tests/unit/adapters/test_pymssql/test_extensions.py index ab8787ecd..d76e4fa8a 100644 --- a/tests/unit/adapters/test_pymssql/test_extensions.py +++ b/tests/unit/adapters/test_pymssql/test_extensions.py @@ -1,11 +1,12 @@ """pymssql extension package tests.""" from typing import Any -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, MagicMock import pytest from sqlspec.adapters.pymssql.config import PymssqlConfig +from sqlspec.adapters.pymssql.litestar.store import PymssqlStore def test_event_store_uses_tsql_column_types_and_idempotent_wrappers() -> None: @@ -126,3 +127,19 @@ def test_adk_store_can_force_fallback_json_column_type() -> None: store = PymssqlADKStore(PymssqlConfig(extension_config={"adk": {"native_json": False}})) assert "state NVARCHAR(MAX) NOT NULL" in store._sessions_table_ddl() + + +def test_litestar_merge_holdlock() -> None: + """Litestar session store MERGE statement uses WITH (HOLDLOCK).""" + store = PymssqlStore( + PymssqlConfig(extension_config={"litestar": {"session_table": "litestar_session", "manage_schema": False}}) + ) + mock_cursor = MagicMock() + mock_conn = MagicMock() + mock_conn.cursor.return_value = mock_cursor + store._config = MagicMock() + store._config.provide_connection.return_value.__enter__.return_value = mock_conn + + store._set("sess-1", b"data", 60) + executed_sql = mock_cursor.execute.call_args[0][0] + assert "MERGE INTO [dbo].[litestar_session] WITH (HOLDLOCK) AS target" in executed_sql diff --git a/tests/unit/adapters/test_pymssql/test_pool.py b/tests/unit/adapters/test_pymssql/test_pool.py index 617cd6e28..f9665b1b0 100644 --- a/tests/unit/adapters/test_pymssql/test_pool.py +++ b/tests/unit/adapters/test_pymssql/test_pool.py @@ -3,7 +3,9 @@ import threading from typing import Any, cast -from tests.unit.adapters.test_pymssql._fakes import FakeConnection, FakePymssqlModule +import pytest + +from tests.unit.adapters.test_pymssql._fakes import FakeConnection, FakeCursor, FakePymssqlModule def test_pool_connects_with_config_and_runs_hook(monkeypatch) -> None: @@ -36,7 +38,7 @@ def test_pool_recycles_failed_health_check(monkeypatch) -> None: import sqlspec.adapters.pymssql.pool as pool_module from sqlspec.adapters.pymssql.pool import PymssqlConnectionPool - first = FakeConnection() + first = FakeConnection(FakeCursor(error=OSError("connection lost"))) second = FakeConnection() connections = [first, second] @@ -47,7 +49,6 @@ def connect(self, **kwargs: Any) -> FakeConnection: fake_module = SequencedModule() monkeypatch.setattr(pool_module, "pymssql", fake_module) - monkeypatch.setattr(PymssqlConnectionPool, "_is_connection_alive", lambda *_: False) pool = PymssqlConnectionPool({"server": "sql.example.test"}, health_check_interval=0.5) @@ -55,6 +56,7 @@ def connect(self, **kwargs: Any) -> FakeConnection: pool._thread_local.last_used = 0.0 # pyright: ignore[reportPrivateUsage] assert cast(object, pool.acquire()) is second assert first.closed is True + assert first.cursor_obj.closed is True assert len(fake_module.connect_calls) == 2 @@ -63,23 +65,15 @@ def test_pool_zero_health_check_interval_disables_health_check(monkeypatch) -> N import sqlspec.adapters.pymssql.pool as pool_module from sqlspec.adapters.pymssql.pool import PymssqlConnectionPool - connection = FakeConnection() + connection = FakeConnection(FakeCursor(error=OSError("health checks should be disabled"))) fake_module = FakePymssqlModule(connection) monkeypatch.setattr(pool_module, "pymssql", fake_module) - health_checks = 0 - - def _fail_if_called(*_args: Any) -> bool: - nonlocal health_checks - health_checks += 1 - return False - - monkeypatch.setattr(PymssqlConnectionPool, "_is_connection_alive", _fail_if_called) pool = PymssqlConnectionPool({"server": "sql.example.test"}, health_check_interval=0) assert cast(object, pool.acquire()) is connection pool._thread_local.last_used = 0.0 # pyright: ignore[reportPrivateUsage] assert cast(object, pool.acquire()) is connection - assert health_checks == 0 + assert connection.cursor_obj.calls == [] assert len(fake_module.connect_calls) == 1 @@ -169,18 +163,96 @@ def test_pool_registry_does_not_grow_across_replacements(monkeypatch) -> None: class SequencedModule(FakePymssqlModule): def connect(self, **kwargs: Any) -> FakeConnection: self.connect_calls.append(kwargs) - return FakeConnection() + return FakeConnection(FakeCursor(error=OSError("connection lost"))) monkeypatch.setattr(pool_module, "pymssql", SequencedModule()) - monkeypatch.setattr(PymssqlConnectionPool, "_is_connection_alive", lambda *_: False) pool = PymssqlConnectionPool({"server": "sql.example.test"}, health_check_interval=0.5) + connections = [] for _ in range(4): - pool.acquire() + connections.append(cast("FakeConnection", pool.acquire())) pool._thread_local.last_used = 0.0 # pyright: ignore[reportPrivateUsage] - assert len(pool._connection_registry) == 1 # pyright: ignore[reportPrivateUsage] + assert pool.size() == 1 + assert [connection.closed for connection in connections] == [True, True, True, False] + + pool.close() + + assert pool.size() == 0 + assert all(connection.closed for connection in connections) + + +def test_pool_reentrant_thread_checkout(monkeypatch: Any) -> None: + """Re-entrant get_connection calls on the same thread do not close the connection on inner error.""" + import sqlspec.adapters.pymssql.pool as pool_module + from sqlspec.adapters.pymssql.pool import PymssqlConnectionPool + + connection = FakeConnection() + fake_module = FakePymssqlModule(connection) + monkeypatch.setattr(pool_module, "pymssql", fake_module) + pool = PymssqlConnectionPool({"server": "sql.example.test"}, recycle_seconds=0) + + with pool.get_connection() as c1: + assert cast(object, c1) is connection + try: + with pool.get_connection() as c2: + assert cast(object, c2) is connection + raise ValueError("inner failure") + except ValueError: + pass + assert not connection.closed + + assert not connection.closed pool.close() + assert connection.closed + + +def test_pool_recovers_after_failed_acquisition_and_retires_failed_checkout(monkeypatch: pytest.MonkeyPatch) -> None: + """A failed connect must not prevent a later failed checkout from retiring its connection.""" + import sqlspec.adapters.pymssql.pool as pool_module + from sqlspec.adapters.pymssql.pool import PymssqlConnectionPool + + dirty_connection = FakeConnection() + clean_connection = FakeConnection() + attempts: list[OSError | FakeConnection] = [ + OSError("temporary connect failure"), + dirty_connection, + clean_connection, + ] + + class RecoveringModule(FakePymssqlModule): + def connect(self, **kwargs: Any) -> FakeConnection: + outcome = attempts.pop(0) + if isinstance(outcome, OSError): + raise outcome + return outcome + + monkeypatch.setattr(pool_module, "pymssql", RecoveringModule()) + pool = PymssqlConnectionPool({"server": "sql.example.test"}, recycle_seconds=0, health_check_interval=0) + + try: + with pytest.raises(OSError, match="temporary connect failure"), pool.get_connection(): + pytest.fail("Failed acquisition must not enter the checkout body") + + assert pool.checked_out() == 0 + assert pool.size() == 0 + + with pytest.raises(ValueError, match="failed transaction"), pool.get_connection() as connection: + assert cast(object, connection) is dirty_connection + dirty_connection.autocommit(False) + raise ValueError("failed transaction") + + assert dirty_connection.closed is True + assert pool.checked_out() == 0 + assert pool.size() == 0 + + with pool.get_connection() as replacement: + assert cast(object, replacement) is clean_connection + assert clean_connection.closed is False + assert clean_connection.autocommit_state is True + assert pool.checked_out() == 1 - assert pool._connection_registry == set() # pyright: ignore[reportPrivateUsage] + assert pool.checked_out() == 0 + finally: + pool.close() diff --git a/tests/unit/adapters/test_pymysql/test_adk_store.py b/tests/unit/adapters/test_pymysql/test_adk_store.py index 76b3f2732..8f938ecf4 100644 --- a/tests/unit/adapters/test_pymysql/test_adk_store.py +++ b/tests/unit/adapters/test_pymysql/test_adk_store.py @@ -286,18 +286,40 @@ def test_pymysql_adk_memory_store_insert_and_search_json_handling() -> None: def test_pymysql_stream_source_closes_cursor_when_execute_raises() -> None: - """PymysqlStreamSource.start() closes the cursor if execute() fails.""" - from sqlspec.adapters.pymysql.core import PymysqlStreamSource + """PyMysqlStreamSource.start() closes the cursor if execute() fails.""" + from sqlspec.adapters.pymysql.core import PyMysqlStreamSource cursor = MagicMock() cursor.execute.side_effect = RuntimeError("execute boom") connection = MagicMock() connection.cursor.return_value = cursor driver = MagicMock(connection=connection) - source = PymysqlStreamSource(driver, "SELECT 1", (), 100, set()) + driver._run_with_exception_handler = lambda _handler, fn: fn() + source = PyMysqlStreamSource(driver, "SELECT 1", (), 100, set()) with pytest.raises(RuntimeError, match="execute boom"): source.start() cursor.close.assert_called_once() assert source._cursor is None + + +def test_pymysql_adk_memory_store_missing_table_returns_empty_or_zero() -> None: + """_search_entries and deletes return empty/0 when table is missing (errno 1146).""" + from pymysql.err import MySQLError + + missing_err = MySQLError(1146, "Table 'test.mem' doesn't exist") + + cursor = MagicMock() + cursor.execute.side_effect = missing_err + conn = MagicMock() + conn.__enter__.return_value = conn + conn.__exit__.return_value = None + conn.cursor.return_value = cursor + config = _mock_config() + config.provide_connection = lambda *_a, **_k: conn + + store = PyMysqlADKMemoryStore(config) + assert store.search_entries("test", "app", "u1") == [] + assert store.delete_entries_by_session("sess-1") == 0 + assert store.delete_entries_older_than(30) == 0 diff --git a/tests/unit/adapters/test_pymysql/test_config.py b/tests/unit/adapters/test_pymysql/test_config.py index 22ceb0f23..cb25c4ba0 100644 --- a/tests/unit/adapters/test_pymysql/test_config.py +++ b/tests/unit/adapters/test_pymysql/test_config.py @@ -294,3 +294,45 @@ def _bad_hook(_conn: object) -> None: raw_conn.close.assert_called_once() finally: pool.close() + + +def test_pool_prunes_dead_thread_connections() -> None: + """PyMysqlConnectionPool closes and prunes connections when worker threads die.""" + import threading + + pool = PyMysqlConnectionPool({}, connection_factory=lambda: MagicMock(server_status=0)) + + def _worker() -> None: + pool.acquire() + + t = threading.Thread(target=_worker) + t.start() + t.join() + + assert pool.size() == 0 + + +def test_pool_get_connection_keeps_connection_on_exception() -> None: + """get_connection keeps healthy connection in pool when caller exception occurs.""" + raw_conn = MagicMock(server_status=0) + pool = PyMysqlConnectionPool({}, connection_factory=lambda: raw_conn) + try: + with pytest.raises(ValueError, match="caller error"): + with pool.get_connection() as conn: + assert conn is raw_conn + raise ValueError("caller error") + assert not raw_conn.close.called + assert pool.size() == 1 + finally: + pool.close() + + +def test_pool_health_check_bypassed_when_interval_zero() -> None: + """health check is bypassed when health_check_interval is 0.""" + raw_conn = MagicMock(server_status=0) + pool = PyMysqlConnectionPool({}, health_check_interval=0, connection_factory=lambda: raw_conn) + try: + pool.acquire() + assert not raw_conn.ping.called + finally: + pool.close() diff --git a/tests/unit/adapters/test_pymysql/test_driver.py b/tests/unit/adapters/test_pymysql/test_driver.py index 09b3271af..0d75c2c4a 100644 --- a/tests/unit/adapters/test_pymysql/test_driver.py +++ b/tests/unit/adapters/test_pymysql/test_driver.py @@ -5,7 +5,7 @@ import pytest from pymysql.constants import SERVER_STATUS -from sqlspec import StatementStack +from sqlspec import SQL, StatementStack from sqlspec.adapters.pymysql import PyMysqlDriver from sqlspec.adapters.pymysql._typing import PyMysqlMySQLError from sqlspec.exceptions import DeadlockError @@ -25,6 +25,9 @@ def close(self) -> None: def execute(self, statement: str, parameters: Any = None) -> None: self.connection.executed.append((statement, parameters)) + def executemany(self, statement: str, parameters: Any = None) -> None: + self.connection.executed.append((statement, parameters)) + def fetchall(self) -> "list[Any]": return [] @@ -73,8 +76,30 @@ def test_execute_stack_commits_when_autocommit_disabled_without_server_transacti assert connection.begins == 1 assert all(statement != "BEGIN" for statement, _ in connection.executed) - assert connection.commits == 1 - assert connection.rollbacks == 0 + + +def test_pymysql_sync_stream_source_in_active_except_block() -> None: + """PyMysqlStreamSource safely handles execution inside an active except block.""" + from unittest.mock import MagicMock + + from sqlspec.adapters.pymysql.core import PyMysqlStreamSource + + cursor = MagicMock() + cursor.fetchmany.return_value = [{"id": 1}] + cursor.description = [("id", 3, None, None, None, None, None)] + connection = MagicMock() + connection.cursor.return_value = cursor + driver = PyMysqlDriver(connection) + + try: + raise ValueError("simulated active exception") + except ValueError: + source = PyMysqlStreamSource(driver, "SELECT 1", (), 10, set()) + source.start() + chunk = source.fetch_chunk() + assert len(chunk) == 1 + source.close() + connection.rollback.assert_not_called() def test_execute_stack_preserves_caller_transaction_when_server_status_reports_transaction() -> None: @@ -105,3 +130,59 @@ def _fail_commit() -> None: driver = _driver(connection) with pytest.raises(DeadlockError): driver.commit() + + +def test_pymysql_stream_source_caches_deserializer() -> None: + """PyMysqlStreamSource caches deserializer in __init__.""" + from unittest.mock import MagicMock + + from sqlspec.adapters.pymysql.core import PyMysqlStreamSource + + def custom_deserializer(v: Any) -> Any: + return v + + driver = MagicMock( + _json_deserializer=custom_deserializer, driver_features={"json_deserializer": custom_deserializer} + ) + + source = PyMysqlStreamSource(driver, "SELECT 1", (), 10, set()) + assert source._deserializer is custom_deserializer + + +def test_pymysql_dispatch_execute_many_escapes_entire_sql() -> None: + """dispatch_execute_many escapes % across the entire SQL statement.""" + conn = _FakeConnection() + cursor = conn.cursor() + driver = PyMysqlDriver(cast("Any", conn)) + statement = SQL( + "INSERT INTO demo (val) VALUES (%s) ON DUPLICATE KEY UPDATE val = val % 2", + [(1,)], + statement_config=driver.statement_config, + is_many=True, + ) + driver.dispatch_execute_many(cursor, statement) + assert len(conn.executed) == 1 + + +def test_pymysql_dispatch_execute_uses_statement_config() -> None: + """dispatch_execute respects statement.statement_config.""" + from sqlspec.core import StatementConfig + + conn = _FakeConnection() + cursor = conn.cursor() + driver = PyMysqlDriver(cast("Any", conn)) + custom_config = StatementConfig(enable_caching=False) + statement = SQL("SELECT 1", statement_config=custom_config) + result = driver.dispatch_execute(cursor, statement) + assert result is not None + + +def test_pymysql_dispatch_execute_script_delegates_to_super() -> None: + """dispatch_execute_script emits script statement events via super().""" + conn = _FakeConnection() + cursor = conn.cursor() + driver = PyMysqlDriver(cast("Any", conn)) + statement = SQL("SELECT 1; SELECT 2").as_script() + result = driver.dispatch_execute_script(cursor, statement) + assert result.statement_count == 2 + assert result.successful_statements == 2 diff --git a/tests/unit/adapters/test_pymysql/test_litestar_store.py b/tests/unit/adapters/test_pymysql/test_litestar_store.py index 08eed3e27..a73029c8b 100644 --- a/tests/unit/adapters/test_pymysql/test_litestar_store.py +++ b/tests/unit/adapters/test_pymysql/test_litestar_store.py @@ -46,3 +46,29 @@ def test_pymysql_litestar_store_rejects_non_string_options() -> None: """Explicit option values cannot silently degrade to an empty clause.""" with pytest.raises(ImproperConfigurationError, match="table_options"): PyMysqlStore(_mock_config({"table_options": 42})) + + +def test_pymysql_litestar_store_drop_table_sql_omits_drop_index() -> None: + """_drop_table_sql must emit only DROP TABLE IF EXISTS without preceding DROP INDEX.""" + store = PyMysqlStore(_mock_config()) + assert store._drop_table_sql() == ["DROP TABLE IF EXISTS litestar_session"] + + +def test_pymysql_litestar_store_set_uses_mariadb_compatible_upsert() -> None: + """_set must use VALUES(col) in ON DUPLICATE KEY UPDATE for MySQL and MariaDB compatibility.""" + cursor = MagicMock() + conn = MagicMock() + conn.cursor.return_value = cursor + conn.__enter__.return_value = conn + config = _mock_config() + config.provide_connection = lambda *_a, **_k: conn + + store = PyMysqlStore(config) + store._set("session-1", b"payload", 300) + + assert cursor.execute.call_count == 1 + sql = cursor.execute.call_args_list[0].args[0] + assert "ON DUPLICATE KEY UPDATE" in sql + assert "data = VALUES(data)" in sql + assert "expires_at = VALUES(expires_at)" in sql + assert "AS new" not in sql diff --git a/tests/unit/adapters/test_pymysql/test_load_from_arrow.py b/tests/unit/adapters/test_pymysql/test_load_from_arrow.py index 1a93b8a3d..dc711f707 100644 --- a/tests/unit/adapters/test_pymysql/test_load_from_arrow.py +++ b/tests/unit/adapters/test_pymysql/test_load_from_arrow.py @@ -2,6 +2,7 @@ from pathlib import Path from typing import Any, cast +from unittest.mock import patch import pyarrow as pa import pyarrow.parquet as pq @@ -123,3 +124,15 @@ def test_load_from_arrow_binds_the_infile_path_as_a_parameter() -> None: assert "LOAD DATA LOCAL INFILE %s INTO TABLE" in statement assert parameters[0].endswith(".tsv") assert parameters[0] not in statement + + +def test_load_from_arrow_needs_preparation_bypasses_local_infile() -> None: + """When arrow rows need preparation, bulk load falls back to executemany instead of local infile.""" + conn = _FakeConnection() + driver = _make_driver(conn, enable_local_infile=True) + + with patch.object(type(driver), "_arrow_rows_need_preparation", lambda _self, _table: True): + driver.load_from_arrow("orders", pa.table({"id": [1], "name": ["a"]})) + + assert conn._cursor.execute_calls == [] + assert len(conn._cursor.executemany_calls) == 1 diff --git a/tests/unit/adapters/test_spanner/test_config.py b/tests/unit/adapters/test_spanner/test_config.py index c545be198..87a708cba 100644 --- a/tests/unit/adapters/test_spanner/test_config.py +++ b/tests/unit/adapters/test_spanner/test_config.py @@ -1,7 +1,10 @@ +from asyncio import CancelledError +from collections.abc import AsyncIterator, Callable +from contextlib import asynccontextmanager from datetime import timedelta from types import SimpleNamespace from typing import TYPE_CHECKING, Any, Self, cast -from unittest.mock import AsyncMock +from unittest.mock import ANY, AsyncMock, MagicMock, Mock import pytest from google.api_core import exceptions as api_exceptions @@ -43,6 +46,119 @@ pytestmark = requires_interpreted +def _assert_sdk_call(mock: Mock, *args: object) -> None: + mock.assert_called_once_with(*args) + if isinstance(mock, AsyncMock): + mock.assert_awaited_once_with(*args) + + +def _hook_checkout(is_async: bool, hook: Callable[[Any], Any]) -> SimpleNamespace: + connection = MagicMock(committed=None, rolled_back=False, _transaction_id=b"txn", _mutations=[]) + snapshot = MagicMock() + if is_async: + connection.__aenter__ = AsyncMock(return_value=connection) + connection.commit = AsyncMock() + connection.rollback = AsyncMock() + snapshot.__aenter__ = AsyncMock(return_value=connection) + snapshot.__aexit__ = AsyncMock(return_value=False) + config_type = SpannerAsyncConfig + mock_type = AsyncMock + else: + connection.__enter__.return_value = connection + snapshot.__enter__.return_value = connection + snapshot.__exit__.return_value = False + config_type = SpannerSyncConfig + mock_type = Mock + session = SimpleNamespace(transaction=Mock(return_value=connection)) + manager = SimpleNamespace(get_session=mock_type(return_value=session), put_session=mock_type()) + database = SimpleNamespace(snapshot=Mock(return_value=snapshot), sessions_manager=manager) + config = config_type( + connection_config={"project": "p", "instance_id": "i", "database_id": "d"}, + connection_instance=cast("Any", object()), + on_connection_create=hook, + ) + config._database = cast("Any", database) + return SimpleNamespace(config=config, connection=connection, snapshot=snapshot, manager=manager, session=session) + + +@asynccontextmanager +async def _hook_session(config: Any, **kwargs: Any) -> AsyncIterator[Any]: + if isinstance(config, SpannerAsyncConfig): + async with config.provide_session(**kwargs) as session: + yield session + else: + with config.provide_session(**kwargs) as session: + yield session + + +@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize("transaction", [False, True], ids=["snapshot", "default-transaction"]) +async def test_connection_hook_runs_for_session_checkout(is_async: bool, transaction: bool) -> None: + hook = Mock(return_value=None) + checkout = _hook_checkout(is_async, hook) + kwargs = {} if transaction else {"transaction": False} + + async with _hook_session(checkout.config, **kwargs) as session: + assert session.connection is checkout.connection + hook.assert_called_once_with(checkout.connection) + + if transaction: + _assert_sdk_call(checkout.connection.commit) + _assert_sdk_call(checkout.manager.put_session, checkout.session) + else: + snapshot_exit = checkout.snapshot.__aexit__ if is_async else checkout.snapshot.__exit__ + _assert_sdk_call(snapshot_exit, None, None, None) + + +@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize("transaction", [False, True], ids=["snapshot", "transaction"]) +@pytest.mark.parametrize("error_type", [ValueError, CancelledError]) +@pytest.mark.parametrize("cleanup_fails", [False, True], ids=["cleanup-ok", "cleanup-fails"]) +async def test_connection_hook_failure_releases_checkout_and_preserves_error( + is_async: bool, transaction: bool, error_type: type[BaseException], cleanup_fails: bool +) -> None: + error = error_type("hook failed") + hook = Mock(side_effect=error) + checkout = _hook_checkout(is_async, hook) + snapshot_exit = checkout.snapshot.__aexit__ if is_async else checkout.snapshot.__exit__ + if cleanup_fails: + cleanup = checkout.connection.rollback if transaction else snapshot_exit + cleanup.side_effect = RuntimeError("cleanup failed") + + with pytest.raises(error_type) as caught: + async with _hook_session(checkout.config, transaction=transaction): + pytest.fail("A failing connection hook must prevent the session from opening") + + assert caught.value is error + hook.assert_called_once_with(checkout.connection) + if transaction: + _assert_sdk_call(checkout.connection.rollback) + checkout.connection.commit.assert_not_called() + _assert_sdk_call(checkout.manager.put_session, checkout.session) + else: + _assert_sdk_call(snapshot_exit, error_type, error, ANY) + + +@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize("cleanup_fails", [False, True], ids=["cleanup-ok", "cleanup-fails"]) +async def test_transaction_enter_cancellation_returns_checkout(is_async: bool, cleanup_fails: bool) -> None: + hook = Mock(return_value=None) + checkout = _hook_checkout(is_async, hook) + error = CancelledError("checkout cancelled") + enter = checkout.connection.__aenter__ if is_async else checkout.connection.__enter__ + enter.side_effect = error + if cleanup_fails: + checkout.manager.put_session.side_effect = RuntimeError("release failed") + + with pytest.raises(CancelledError) as caught: + async with _hook_session(checkout.config): + pytest.fail("A cancelled checkout must prevent the session from opening") + + assert caught.value is error + hook.assert_not_called() + _assert_sdk_call(checkout.manager.put_session, checkout.session) + + class _DummyDriver(SyncDriverAdapterBase): dialect = "spanner" @@ -75,6 +191,32 @@ def test_config_defaults() -> None: assert config.connection_config["size"] == 10 +@pytest.mark.parametrize("config_type", [SpannerSyncConfig, SpannerAsyncConfig]) +def test_arrow_conversion_storage_capabilities(config_type: type[SpannerSyncConfig] | type[SpannerAsyncConfig]) -> None: + config = config_type(connection_config={"project": "p", "instance_id": "i", "database_id": "d"}) + + capabilities = config.storage_capabilities() + + assert capabilities["arrow_import_enabled"] is True + assert capabilities["arrow_export_enabled"] is True + assert config.supports_native_arrow_import is False + assert config.supports_native_arrow_export is False + assert config.driver_features["storage_capabilities"] == capabilities + + +@pytest.mark.parametrize("config_type", [SpannerSyncConfig, SpannerAsyncConfig]) +def test_arrow_conversion_storage_preserves_explicit_disable( + config_type: type[SpannerSyncConfig] | type[SpannerAsyncConfig], +) -> None: + capabilities = {"arrow_import_enabled": False, "arrow_export_enabled": False} + config = config_type( + connection_config={"project": "p", "instance_id": "i", "database_id": "d"}, + driver_features={"storage_capabilities": capabilities}, + ) + + assert config.driver_features["storage_capabilities"] == capabilities + + def test_min_sessions_is_rejected() -> None: """Spanner's current session pool classes do not support min_sessions.""" with pytest.raises(ImproperConfigurationError, match="min_sessions"): diff --git a/tests/unit/adapters/test_spanner/test_core.py b/tests/unit/adapters/test_spanner/test_core.py index e63b9600b..34b91e9ca 100644 --- a/tests/unit/adapters/test_spanner/test_core.py +++ b/tests/unit/adapters/test_spanner/test_core.py @@ -18,7 +18,7 @@ from sqlspec.adapters.spanner._typing import SpannerSyncSessionContext from sqlspec.adapters.spanner.adk.artifact_store import SpannerAsyncADKArtifactStore, SpannerSyncADKArtifactStore from sqlspec.adapters.spanner.adk.store import SpannerSyncADKMemoryStore, SpannerSyncADKStore -from sqlspec.adapters.spanner.config import SpannerSyncConfig, SpannerSyncConnectionContext +from sqlspec.adapters.spanner.config import SpannerSyncConfig from sqlspec.adapters.spanner.core import ( build_param_type_signature, build_statement_config, @@ -39,8 +39,10 @@ from sqlspec.adapters.spanner.type_converter import spanner_json, spanner_to_uuid from sqlspec.core import SQL, TypedParameter from sqlspec.exceptions import ( + DataError, DeadlockError, ImproperConfigurationError, + IntegrityError, OperationalError, PermissionDeniedError, QueryTimeoutError, @@ -283,6 +285,22 @@ def test_exception_mapping_permission_denied_unauthenticated_and_timeout() -> No assert isinstance(create_mapped_exception(too_many), OperationalError) assert isinstance(create_mapped_exception(already_exists), UniqueViolationError) + failed_precondition = cast("Any", api_exceptions.FailedPrecondition)("fk violation") + out_of_range = cast("Any", api_exceptions.OutOfRange)("out of range") + resource_exhausted = cast("Any", api_exceptions.ResourceExhausted)("quota exceeded") + + assert isinstance(create_mapped_exception(failed_precondition), IntegrityError) + assert isinstance(create_mapped_exception(out_of_range), DataError) + assert isinstance(create_mapped_exception(resource_exhausted), OperationalError) + + class _MockStatus: + def __init__(self, code: int, message: str) -> None: + self.code = code + self.message = message + + status_err = create_mapped_exception(_MockStatus(9, "constraint failed")) + assert isinstance(status_err, IntegrityError) + def test_connection_context_maps_implicit_commit_and_suppresses_rollback_errors() -> None: """Verify SpannerSyncConnectionContext maps commit GoogleAPICallError and suppresses rollback errors.""" @@ -310,16 +328,16 @@ def rollback(self) -> None: sessions_manager.get_session.return_value = session_obj database = SimpleNamespace(sessions_manager=sessions_manager) - config = MagicMock() - config.get_database.return_value = database - config.driver_features = {} + config = SpannerSyncConfig( + connection_config={"project": "p", "instance_id": "i", "database_id": "d"}, + connection_instance=cast("Any", object()), + ) + config._database = cast("Any", database) - ctx = SpannerSyncConnectionContext(config, transaction=True) - with pytest.raises(DeadlockError), ctx: + with pytest.raises(DeadlockError), config.provide_connection(transaction=True): pass - ctx_rollback = SpannerSyncConnectionContext(config, transaction=True) - with pytest.raises(RuntimeError, match="user boom"), ctx_rollback: + with pytest.raises(RuntimeError, match="user boom"), config.provide_connection(transaction=True): raise RuntimeError("user boom") diff --git a/tests/unit/adapters/test_spanner/test_events_store.py b/tests/unit/adapters/test_spanner/test_events_store.py index 2cd682d1a..9b6ce922d 100644 --- a/tests/unit/adapters/test_spanner/test_events_store.py +++ b/tests/unit/adapters/test_spanner/test_events_store.py @@ -3,8 +3,8 @@ import inspect from types import SimpleNamespace -from typing import Any -from unittest.mock import AsyncMock, MagicMock +from typing import TYPE_CHECKING, Any, cast +from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -12,6 +12,9 @@ from sqlspec.adapters.spanner.events import SpannerAsyncEventQueueStore, SpannerSyncEventQueueStore from sqlspec.exceptions import ImproperConfigurationError +if TYPE_CHECKING: + from collections.abc import Callable + def _mock_spanner_config(events_settings: dict[str, Any] | None = None, *, dialect: str = "spanner") -> MagicMock: """Create a mock SpannerSyncConfig.""" @@ -318,3 +321,45 @@ async def test_ddl_operations_reject_other_config_types( result = getattr(store, operation)() if inspect.isawaitable(result): await result + + +def test_events_store_spangres_dialect() -> None: + """_is_spangres recognizes postgres and postgresql dialect names.""" + for dialect_name in ("spangres", "postgres", "postgresql", "spanner_postgresql"): + config = MagicMock() + config.statement_config = MagicMock(dialect=dialect_name) + config.connection_config = {} + config.extension_config = {"events": {}} + store = SpannerSyncEventQueueStore(config) + assert store._is_spangres() is True + + +def test_events_store_drop_table_if_exists() -> None: + """drop_statements generates IF EXISTS clauses for index and table.""" + config = _mock_spanner_config() + store = SpannerSyncEventQueueStore(config) + statements = store.drop_statements() + assert "DROP INDEX IF EXISTS" in statements[0] + assert "DROP TABLE IF EXISTS" in statements[1] + + +def test_events_store_ddl_error_mapping() -> None: + """create_table maps GoogleAPICallError to SQLSpecError and passes configured ddl_timeout.""" + from sqlspec.adapters.spanner._typing import SpannerGoogleAPICallError + from sqlspec.exceptions import SQLSpecError + + mock_db = MagicMock() + mock_db.update_ddl.side_effect = cast("Callable[[str], SpannerGoogleAPICallError]", SpannerGoogleAPICallError)( + "DDL failed" + ) + + config = MagicMock(spec=SpannerSyncConfig) + config.extension_config = {"events": {"queue_table": "test_queue", "ddl_timeout": 300.0}} + config.get_database.return_value = mock_db + + store = SpannerSyncEventQueueStore(config) + with patch("sqlspec.adapters.spanner.events.store.list_existing_table_names_sync", return_value=set()): + with pytest.raises(SQLSpecError): + store.create_table() + + mock_db.update_ddl.assert_called_once() diff --git a/tests/unit/adapters/test_spanner/test_litestar_store.py b/tests/unit/adapters/test_spanner/test_litestar_store.py index 6b6954d0d..c3a431879 100644 --- a/tests/unit/adapters/test_spanner/test_litestar_store.py +++ b/tests/unit/adapters/test_spanner/test_litestar_store.py @@ -42,7 +42,7 @@ def test_delete_uses_run_in_transaction() -> None: def test_delete_all_uses_run_in_transaction() -> None: - """Verify _delete_all uses database.run_in_transaction for write operations.""" + """Verify _delete_all uses execute_partitioned_dml.""" mock_db = _mock_database() config = MagicMock() @@ -52,11 +52,11 @@ def test_delete_all_uses_run_in_transaction() -> None: store = SpannerSyncStore(config) store._delete_all() # pyright: ignore - mock_db.run_in_transaction.assert_called_once() + mock_db.execute_partitioned_dml.assert_called_once() def test_delete_expired_uses_run_in_transaction() -> None: - """Verify _delete_expired uses database.run_in_transaction for write operations.""" + """Verify _delete_expired uses execute_partitioned_dml.""" mock_db = _mock_database() config = MagicMock() @@ -66,7 +66,49 @@ def test_delete_expired_uses_run_in_transaction() -> None: store = SpannerSyncStore(config) store._delete_expired() # pyright: ignore - mock_db.run_in_transaction.assert_called_once() + mock_db.execute_partitioned_dml.assert_called_once() + + +def test_set_single_dml_upsert() -> None: + """Verify _set issues a single INSERT OR UPDATE DML statement.""" + executed_statements: list[str] = [] + + class _MockTxn: + def execute_update(self, sql: str, params: Any = None, param_types: Any = None) -> int: + executed_statements.append(sql) + return 1 + + mock_db = MagicMock() + mock_db.run_in_transaction = MagicMock(side_effect=lambda func: func(_MockTxn())) + + config = MagicMock() + config.extension_config = {"litestar": {"session_table": "sess"}} + config.get_database.return_value = mock_db + + store = SpannerSyncStore(config) + store._set("s1", b"data", None) # pyright: ignore + + assert len(executed_statements) == 1 + assert "INSERT OR UPDATE INTO" in executed_statements[0] + + +def test_row_deletion_policy_ddl() -> None: + """Verify _table_ddl includes ROW DELETION POLICY when enabled.""" + config = MagicMock() + config.extension_config = {"litestar": {"session_table": "sess", "enable_row_deletion_policy": True}} + store = SpannerSyncStore(config) + ddl = store._table_ddl() + assert "ROW DELETION POLICY (OLDER_THAN(expires_at, INTERVAL 0 DAY))" in ddl + + +def test_litestar_store_param_types_cached() -> None: + """_get_param_types returns module-level cached constants.""" + config = MagicMock() + config.extension_config = {"litestar": {"session_table": "sess"}} + store = SpannerSyncStore(config) + t1 = store._get_param_types(session_id=True, data=True, expires_at=True) + t2 = store._get_param_types(session_id=True, data=True, expires_at=True) + assert t1 is t2 def _context_manager_yielding(value: Any) -> Any: @@ -134,6 +176,7 @@ async def _run_in_txn(func: Any) -> Any: return await func(txn) db.run_in_transaction = AsyncMock(side_effect=_run_in_txn) + db.execute_partitioned_dml = AsyncMock(return_value=1) db._txn = txn return db @@ -181,11 +224,11 @@ async def test_async_store_set_delete_delete_all_delete_expired() -> None: assert db.run_in_transaction.await_count == 2 await store.delete_all() - assert db.run_in_transaction.await_count == 3 + assert db.execute_partitioned_dml.await_count == 1 deleted = await store.delete_expired() assert deleted == 1 - assert db.run_in_transaction.await_count == 4 + assert db.execute_partitioned_dml.await_count == 2 async def test_async_store_get_exists_expires_in() -> None: diff --git a/tests/unit/adapters/test_spanner/test_load_from_arrow_mutations.py b/tests/unit/adapters/test_spanner/test_load_from_arrow_mutations.py index c562f0102..26e1bee0a 100644 --- a/tests/unit/adapters/test_spanner/test_load_from_arrow_mutations.py +++ b/tests/unit/adapters/test_spanner/test_load_from_arrow_mutations.py @@ -6,16 +6,9 @@ import pytest import sqlspec.adapters.spanner.driver as spanner_driver +from sqlspec.adapters.spanner.config import SpannerSyncConfig from sqlspec.adapters.spanner.driver import SpannerSyncDriver -CAPABILITIES = { - "arrow_export_enabled": True, - "arrow_import_enabled": True, - "parquet_export_enabled": True, - "parquet_import_enabled": True, - "partition_strategies": ["fixed"], -} - class _FakeTransaction: def __init__(self) -> None: @@ -34,7 +27,8 @@ def execute_update(self, sql: str, params: Any = None, param_types: Any = None, @pytest.fixture def mutations_driver(monkeypatch: pytest.MonkeyPatch) -> SpannerSyncDriver: monkeypatch.setattr(spanner_driver, "SpannerTransaction", _FakeTransaction) - return SpannerSyncDriver(cast("Any", _FakeTransaction()), driver_features={"storage_capabilities": CAPABILITIES}) + config = SpannerSyncConfig(connection_config={"project": "p", "instance_id": "i", "database_id": "d"}) + return SpannerSyncDriver(cast("Any", _FakeTransaction()), driver_features=config.driver_features) def test_load_from_arrow_uses_single_insert_or_update(mutations_driver: SpannerSyncDriver) -> None: diff --git a/tests/unit/adapters/test_spanner/test_session_controls.py b/tests/unit/adapters/test_spanner/test_session_controls.py index 26c915e6b..491bb3533 100644 --- a/tests/unit/adapters/test_spanner/test_session_controls.py +++ b/tests/unit/adapters/test_spanner/test_session_controls.py @@ -146,3 +146,19 @@ def __init__(self) -> None: assert database.sessions_manager.returned == 2 assert database.sessions_manager.session.create_calls == 0 assert database.sessions_manager.session.delete_calls == 0 + + +def test_driver_nested_in_transaction_guard() -> None: + """_connection_in_transaction returns True when _owns_transaction is False, preventing commit.""" + from unittest.mock import Mock + + from sqlspec.adapters.spanner.driver import SpannerSyncDriver + + mock_conn = Mock() + driver = SpannerSyncDriver(connection=mock_conn) + driver.bind_external_transaction() + assert driver._connection_in_transaction() is True + driver.commit() + assert mock_conn.commit.call_count == 0 + driver.rollback() + assert mock_conn.rollback.call_count == 0 diff --git a/tests/unit/adapters/test_spanner/test_spanner_metadata_extensions.py b/tests/unit/adapters/test_spanner/test_spanner_metadata_extensions.py index f47032e64..abd8bb6cd 100644 --- a/tests/unit/adapters/test_spanner/test_spanner_metadata_extensions.py +++ b/tests/unit/adapters/test_spanner/test_spanner_metadata_extensions.py @@ -27,7 +27,8 @@ def test_get_constraints_by_table() -> None: results = dictionary.get_constraints(cast(Any, driver), table="Albums") assert len(results) == 1 - assert driver.last_params.get("table_name") == "Albums" + table_param = driver.last_params.get("table_name") + assert getattr(table_param, "value", table_param) == "Albums" assert "TABLE_CONSTRAINTS" in str(driver.last_query) @@ -38,8 +39,10 @@ def test_get_constraints_by_schema() -> None: results = dictionary.get_constraints(cast(Any, driver), schema="catalog") assert len(results) == 1 - assert driver.last_params.get("schema_name") == "catalog" - assert driver.last_params.get("table_name") is None + schema_param = driver.last_params.get("schema_name") + assert getattr(schema_param, "value", schema_param) == "catalog" + table_param = driver.last_params.get("table_name") + assert getattr(table_param, "value", table_param) is None def test_get_sequences() -> None: @@ -133,3 +136,26 @@ def test_postgresql_mode_data_dictionary_routing() -> None: dictionary.get_change_streams(cast(Any, driver)) assert "information_schema.change_streams" in str(driver.last_query) + + +def test_untyped_none_filter_parameters_wrapped() -> None: + """get_constraints wraps None filter parameters in TypedParameter with str type.""" + from sqlspec.core import TypedParameter + + dictionary = SpannerSyncDataDictionary() + driver = MockSpannerDriver() + + dictionary.get_constraints(cast(Any, driver), schema=None, table=None) + table_param = driver.last_params.get("table_name") + assert isinstance(table_param, TypedParameter) + assert table_param.value is None + assert table_param.original_type is str + + +def test_spanner_unsupported_domains_reported() -> None: + """get_metadata_capabilities reports UNSUPPORTED for schemas and dependencies.""" + dictionary = SpannerSyncDataDictionary() + driver = MockSpannerDriver() + profile = dictionary.get_metadata_capabilities(driver) + assert profile.get("schemas").support == MetadataSupport.UNSUPPORTED + assert profile.get("dependencies").support == MetadataSupport.UNSUPPORTED diff --git a/tests/unit/adapters/test_spanner/test_spanner_query_options.py b/tests/unit/adapters/test_spanner/test_spanner_query_options.py index dcf9aa0c1..9c1ffbe15 100644 --- a/tests/unit/adapters/test_spanner/test_spanner_query_options.py +++ b/tests/unit/adapters/test_spanner/test_spanner_query_options.py @@ -1,10 +1,120 @@ """Unit tests for Spanner QueryOptions forwarding.""" -from unittest.mock import MagicMock +from collections.abc import AsyncIterator, Iterator +from types import SimpleNamespace +from typing import TYPE_CHECKING, Any, cast +from unittest.mock import AsyncMock, MagicMock + +import pytest +from google.api_core import exceptions as api_exceptions +from google.cloud.spanner_v1.types.type import TypeCode from sqlspec.adapters.spanner.config import SpannerSyncConfig from sqlspec.adapters.spanner.core import default_statement_config -from sqlspec.adapters.spanner.driver import SpannerSyncDriver +from sqlspec.adapters.spanner.driver import SpannerAsyncDriver, SpannerSyncDriver +from sqlspec.exceptions import OperationalError + +if TYPE_CHECKING: + from collections.abc import Callable + + from sqlspec.adapters.spanner._typing import SpannerAsyncConnection, SpannerSyncConnection + + +class _StreamResult: + def __init__(self, error: BaseException | None) -> None: + self.error = error + self.closed = False + self.metadata = SimpleNamespace( + row_type=SimpleNamespace(fields=[SimpleNamespace(name="id", type_=SimpleNamespace(code=TypeCode.INT64))]) + ) + + def __iter__(self) -> Iterator[list[int]]: + try: + yield [1] + if self.error is not None: + raise self.error + yield [2] + finally: + self.closed = True + + async def __aiter__(self) -> AsyncIterator[list[int]]: + try: + yield [1] + if self.error is not None: + raise self.error + yield [2] + finally: + self.closed = True + + +async def _consume_stream(driver: Any, is_async: bool, *, fails_fetch: bool = False) -> list[dict[str, Any]]: + stream = driver.select_stream("SELECT id FROM users", chunk_size=1) + if is_async: + async with stream: + if fails_fetch: + assert await anext(stream) == {"id": 1} + return [row async for row in stream] + with stream: + if fails_fetch: + assert next(stream) == {"id": 1} + return list(stream) + + +@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize("phase", ["start", "fetch"]) +@pytest.mark.parametrize("ambient", [False, True], ids=["normal", "application-except"]) +async def test_stream_maps_sdk_errors_with_ambient_application_exception( + is_async: bool, phase: str, ambient: bool +) -> None: + error = cast("Callable[[str], api_exceptions.ServiceUnavailable]", api_exceptions.ServiceUnavailable)( + "Spanner unavailable" + ) + result = _StreamResult(error) + execute = AsyncMock() if is_async else MagicMock() + execute.side_effect = error if phase == "start" else None + execute.return_value = result + connection = SimpleNamespace(execute_sql=execute) + driver = ( + SpannerAsyncDriver(connection=cast("SpannerAsyncConnection", connection)) + if is_async + else SpannerSyncDriver(connection=cast("SpannerSyncConnection", connection)) + ) + + with pytest.raises(OperationalError, match="Spanner unavailable") as caught: + try: + if ambient: + raise LookupError("application error") + except LookupError: + await _consume_stream(driver, is_async, fails_fetch=phase == "fetch") + else: + await _consume_stream(driver, is_async, fails_fetch=phase == "fetch") + + assert caught.value.__cause__ is error + execute.assert_called_once() + if phase == "fetch": + assert result.closed is True + + +@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"]) +async def test_stream_succeeds_with_ambient_application_exception(is_async: bool) -> None: + result = _StreamResult(None) + execute = AsyncMock(return_value=result) if is_async else MagicMock(return_value=result) + connection = SimpleNamespace(execute_sql=execute) + driver = ( + SpannerAsyncDriver(connection=cast("SpannerAsyncConnection", connection)) + if is_async + else SpannerSyncDriver(connection=cast("SpannerSyncConnection", connection)) + ) + error = LookupError("application error") + + try: + raise error + except LookupError as caught: + rows = await _consume_stream(driver, is_async) + assert caught is error + + assert rows == [{"id": 1}, {"id": 2}] + assert result.closed is True def test_driver_execute_with_statement_query_options() -> None: diff --git a/tests/unit/adapters/test_sqlanywhere.py b/tests/unit/adapters/test_sqlanywhere.py index 06da062cb..c2093cb38 100644 --- a/tests/unit/adapters/test_sqlanywhere.py +++ b/tests/unit/adapters/test_sqlanywhere.py @@ -1,5 +1,8 @@ """Unit tests for the SAP SQL Anywhere (sqlanydb) adapter.""" +import threading +from contextlib import nullcontext +from datetime import UTC, datetime from typing import Any, cast from unittest.mock import MagicMock @@ -20,25 +23,29 @@ default_statement_config, driver_profile, ) -from sqlspec.adapters.sqlanywhere.adk import SqlanywhereADKConfig, SqlanywhereADKMemoryStore, SqlanywhereADKStore -from sqlspec.adapters.sqlanywhere.events import SqlanywhereEventQueueStore, SqlanywhereSyncEventQueueStore -from sqlspec.adapters.sqlanywhere.litestar import SqlanywhereStore +from sqlspec.adapters.sqlanywhere._typing import SqlanywhereError, SqlanywhereSessionContext +from sqlspec.adapters.sqlanywhere.adk import SqlanywhereADKMemoryStore, SqlanywhereADKStore +from sqlspec.adapters.sqlanywhere.core import build_insert_statement, format_identifier +from sqlspec.adapters.sqlanywhere.data_dictionary import build_column_metadata, build_foreign_key_metadata +from sqlspec.adapters.sqlanywhere.events import SqlanywhereEventQueueStore +from sqlspec.adapters.sqlanywhere.litestar import SqlanywhereLitestarConfig, SqlanywhereStore from sqlspec.adapters.sqlanywhere.migrations import SqlanywhereSyncMigrationTracker from sqlspec.core import SQL, ParameterStyle from sqlspec.data_dictionary.dialects.sqlanywhere import SqlanywhereVersionInfo +from sqlspec.data_dictionary.dialects.sqlanywhere.config import SQLANYWHERE_FEATURE_FLAGS from sqlspec.exceptions import ( CheckViolationError, DatabaseConnectionError, DataError, DeadlockError, ForeignKeyViolationError, - MissingDependencyError, NotFoundError, NotNullViolationError, OperationalError, PermissionDeniedError, SerializationConflictError, SQLParsingError, + TransactionError, UniqueViolationError, ) @@ -168,7 +175,7 @@ def connect(self, **kwargs: Any) -> FakeConnection: captured_kwargs.append(kwargs) return fake_conn - monkeypatch.setattr(pool_module, "_load_sqlanydb_module", lambda: FakeSqlanydbModule()) + monkeypatch.setattr(pool_module, "SQLANYDB_MODULE", FakeSqlanydbModule()) cfg = build_connection_config({"user": "dba", "password": "sql", "database": "demo", "host": "localhost"}) pool = SqlanywhereConnectionPool(cfg, min_size=1, max_size=2) @@ -185,21 +192,6 @@ def connect(self, **kwargs: Any) -> FakeConnection: assert fake_conn.closed is True -def test_sqlanywhere_pool_raises_missing_dependency_when_sqlanydb_absent(monkeypatch: pytest.MonkeyPatch) -> None: - """SqlanywhereConnectionPool raises MissingDependencyError when sqlanydb is not installed.""" - import sqlspec.adapters.sqlanywhere.pool as pool_module - - def _raise_import_error() -> Any: - msg = "No module named 'sqlanydb'" - raise ImportError(msg) - - monkeypatch.setattr(pool_module, "import_module", lambda name: _raise_import_error()) - - pool = SqlanywhereConnectionPool({"uid": "dba", "pwd": "sql"}) - with pytest.raises(MissingDependencyError, match="sqlanydb"): - pool.acquire() - - @pytest.mark.parametrize( ("args", "expected_type"), [ @@ -376,7 +368,7 @@ def _select_side_effect(query: Any, *args: Any, **kwargs: Any) -> list[dict[str, def test_sqlanywhere_migration_tracker_and_extension_stores() -> None: - """SqlanywhereSyncMigrationTracker, SqlanywhereStore, SqlanywhereSyncEventQueueStore, and SqlanywhereADKStore emit SQL Anywhere DDL.""" + """SqlanywhereSyncMigrationTracker, SqlanywhereStore, SqlanywhereEventQueueStore, and SqlanywhereADKStore emit SQL Anywhere DDL.""" tracker = SqlanywhereSyncMigrationTracker(version_table_name="DBA.ddl_migrations") ddl_text = tracker._idempotent_tracking_table_ddl_text() assert "CREATE TABLE IF NOT EXISTS" in ddl_text @@ -408,8 +400,8 @@ def test_sqlanywhere_migration_tracker_and_extension_stores() -> None: drop_litestar_sql = litestar_store._drop_table_sql() assert "DROP TABLE IF EXISTS" in drop_litestar_sql[0] - event_store = SqlanywhereSyncEventQueueStore(config) - assert SqlanywhereEventQueueStore is SqlanywhereSyncEventQueueStore + event_store = SqlanywhereEventQueueStore(config) + assert SqlanywhereEventQueueStore is SqlanywhereEventQueueStore event_statements = event_store.create_statements() assert "LONG VARCHAR" in event_statements[0] assert "TIMESTAMP" in event_statements[0] @@ -431,7 +423,6 @@ def test_sqlanywhere_migration_tracker_and_extension_stores() -> None: assert "CREATE TABLE IF NOT EXISTS" in memory_ddl assert "LONG VARCHAR" in memory_ddl assert "TIMESTAMP" in memory_ddl - assert isinstance(SqlanywhereADKConfig(), dict) @pytest.mark.anyio @@ -487,3 +478,234 @@ def cursor(self) -> CursorWithoutClose: with SqlanywhereCursor(cast("SqlanywhereConnection", ConnWithBareCursor())) as cursor: assert isinstance(cursor, CursorWithoutClose) + + +def test_sqlanywhere_adk_store_handles_native_missing_table(monkeypatch: pytest.MonkeyPatch) -> None: + """SqlanywhereADKStore.get_session returns None when native driver signals table missing.""" + config = SqlanywhereConfig(connection_config={"user": "dba", "password": "sql", "database": "demo"}) + store = SqlanywhereADKStore(config) + + def _table_missing(_self: Any, *_args: Any, **_kwargs: Any) -> None: + raise SqlanywhereError("Table 'adk_session' not found", -141) + + monkeypatch.setattr(SqlanywhereADKStore, "_execute_fetchone", _table_missing) + assert store.get_session("app", "user", "sess") is None + + +def test_sqlanywhere_adk_memory_store_insert_sql_contains_from_sys_dummy(monkeypatch: pytest.MonkeyPatch) -> None: + """SqlanywhereADKMemoryStore.insert_memory_entries emits SELECT ... FROM SYS.DUMMY WHERE NOT EXISTS.""" + config = SqlanywhereConfig(connection_config={"user": "dba", "password": "sql", "database": "demo"}) + store = SqlanywhereADKMemoryStore(config) + executed_statements: list[str] = [] + + class DummyCursor: + rowcount = 1 + + def execute(self, sql: str, _params: Any = ()) -> None: + executed_statements.append(sql) + + def close(self) -> None: + pass + + class DummyConnection: + def cursor(self) -> DummyCursor: + return DummyCursor() + + def commit(self) -> None: + pass + + def rollback(self) -> None: + pass + + monkeypatch.setattr(SqlanywhereConfig, "provide_connection", lambda _self: nullcontext(DummyConnection())) + + entry = { + "id": "mem-1", + "session_id": "sess-1", + "app_name": "app", + "user_id": "u1", + "scope": "user", + "event_id": "evt-1", + "author": "tester", + "timestamp": datetime.now(UTC), + "content_json": {"msg": "hi"}, + "content_text": "hi", + "metadata_json": None, + } + inserted = store.insert_memory_entries([cast(Any, entry)]) + assert inserted == 1 + assert len(executed_statements) == 1 + assert "FROM SYS.DUMMY" in executed_statements[0] + assert "WHERE NOT EXISTS" in executed_statements[0] + + +def test_sqlanywhere_adk_memory_store_handles_native_missing_table(monkeypatch: pytest.MonkeyPatch) -> None: + """SqlanywhereADKMemoryStore.search_entries returns empty list when table is missing.""" + config = SqlanywhereConfig(connection_config={"user": "dba", "password": "sql", "database": "demo"}) + store = SqlanywhereADKMemoryStore(config) + + def _table_missing(_self: Any, *_args: Any, **_kwargs: Any) -> None: + raise SqlanywhereError("Table 'adk_memory' not found", -141) + + monkeypatch.setattr(SqlanywhereADKMemoryStore, "_execute_fetchall", _table_missing) + assert store.search_entries("query", "app", "user") == [] + + +def test_sqlanywhere_litestar_config_and_delete_expired(monkeypatch: pytest.MonkeyPatch) -> None: + """SqlanywhereLitestarConfig is usable and _delete_expired uses resolve_rowcount.""" + config = SqlanywhereLitestarConfig() + assert isinstance(config, dict) + + adapter_config = SqlanywhereConfig(connection_config={"user": "dba", "password": "sql", "database": "demo"}) + litestar_store = SqlanywhereStore(adapter_config) + + class DummyCursor: + rowcount = 4 + + def execute(self, _sql: str) -> None: + pass + + class DummyConnection: + def cursor(self) -> DummyCursor: + return DummyCursor() + + def commit(self) -> None: + pass + + monkeypatch.setattr(SqlanywhereConfig, "provide_connection", lambda _self: nullcontext(DummyConnection())) + assert litestar_store._delete_expired() == 4 + + +def test_sqlanywhere_feature_flags_supports_transactional_ddl() -> None: + """SQLANYWHERE_FEATURE_FLAGS supports_transactional_ddl is False.""" + assert SQLANYWHERE_FEATURE_FLAGS["supports_transactional_ddl"] is False + + +def test_sqlanywhere_build_column_metadata_is_nullable() -> None: + """build_column_metadata parses NO and YES correctly.""" + cols = build_column_metadata([ + {"column_name": "id", "is_nullable": "NO"}, + {"column_name": "val", "is_nullable": "YES"}, + ]) + assert cols[0]["is_nullable"] is False + assert cols[1]["is_nullable"] is True + + +def test_sqlanywhere_build_foreign_key_metadata_names() -> None: + """build_foreign_key_metadata supports referenced_table and referenced_table_name keys.""" + fks = build_foreign_key_metadata([ + {"table_name": "t1", "column_name": "c1", "referenced_table": "t2", "referenced_column": "c2"}, + {"table_name": "t1", "column_name": "c3", "referenced_table_name": "t3", "referenced_column_name": "c4"}, + ]) + assert fks[0].referenced_table == "t2" + assert fks[0].referenced_column == "c2" + assert fks[1].referenced_table == "t3" + assert fks[1].referenced_column == "c4" + + +def test_sqlanywhere_sync_data_dictionary_instantiation() -> None: + """SqlanywhereSyncDataDictionary instantiates cleanly without redundant __init__.""" + dd = SqlanywhereSyncDataDictionary() + assert dd.dialect == "sqlanywhere" + + +def test_sqlanywhere_has_schema_under_active_exception(monkeypatch: pytest.MonkeyPatch) -> None: + """has_schema runs cleanly inside an active exception block without isolation leak.""" + driver = SqlanywhereDriver(cast("SqlanywhereConnection", FakeConnection())) + + class DummyCursor: + def execute(self, _sql: str, _params: Any = ()) -> None: + pass + + def fetchone(self) -> Any: + return (1,) + + def close(self) -> None: + pass + + class DummyConnection: + def cursor(self) -> DummyCursor: + return DummyCursor() + + monkeypatch.setattr(driver, "connection", DummyConnection()) + + try: + raise RuntimeError("active error") + except RuntimeError: + has_user = driver.has_schema("DBA") + assert has_user is True + + +def test_sqlanywhere_format_identifier_and_insert_statement() -> None: + """format_identifier and build_insert_statement format identifiers correctly.""" + assert format_identifier("schema.my_table") == '"schema"."my_table"' + sql = build_insert_statement("my_table", ["col1", "col2"]) + assert sql == 'INSERT INTO "my_table" ("col1", "col2") VALUES (?, ?)' + + +def test_sqlanywhere_commit_and_rollback_raise_transaction_error() -> None: + """commit() and rollback() wrap caught exceptions in TransactionError.""" + driver = SqlanywhereDriver(cast("SqlanywhereConnection", FakeConnection())) + + class FailingConnection: + def commit(self) -> None: + raise RuntimeError("commit failed") + + def rollback(self) -> None: + raise RuntimeError("rollback failed") + + driver.connection = cast("SqlanywhereConnection", FailingConnection()) + with pytest.raises(TransactionError, match="Failed to commit SQL Anywhere transaction"): + driver.commit() + with pytest.raises(TransactionError, match="Failed to rollback SQL Anywhere transaction"): + driver.rollback() + + +def test_sqlanywhere_session_context_rolls_back_implicit_work() -> None: + """DML without begin() must be rolled back before a pooled connection is reused.""" + conn = FakeConnection() + ctx = SqlanywhereSessionContext( + acquire_connection=lambda: cast("SqlanywhereConnection", conn), + release_connection=lambda _c, **_kw: None, + statement_config=default_statement_config, + driver_features={}, + prepare_driver=lambda d: d, + ) + with ctx as driver: + driver.execute("INSERT INTO items (id) VALUES (1)") + + assert conn.rollbacks == 1 + + +def test_sqlanywhere_connection_pool_bounds_max_size(monkeypatch: pytest.MonkeyPatch) -> None: + """SqlanywhereConnectionPool enforces max_size in multi-threaded acquisition.""" + pool = SqlanywhereConnectionPool({"user": "dba", "password": "sql", "database": "demo"}, min_size=1, max_size=2) + + class DummyConn: + def close(self) -> None: + pass + + monkeypatch.setattr( + SqlanywhereConnectionPool, "new_connection", lambda _self: cast("SqlanywhereConnection", DummyConn()) + ) + + errors: list[Exception] = [] + conns: list[Any] = [] + barrier = threading.Barrier(3) + + def worker() -> None: + try: + conn = pool._get_thread_connection() + conns.append(conn) + barrier.wait(timeout=2.0) + except Exception as exc: + errors.append(exc) + barrier.wait(timeout=2.0) + + threads = [threading.Thread(target=worker) for _ in range(3)] + for t in threads: + t.start() + for t in threads: + t.join(timeout=3.0) + + assert any(isinstance(err, OperationalError) and "pool exhausted" in str(err) for err in errors) diff --git a/tests/unit/adapters/test_sqlite/test_adk_config.py b/tests/unit/adapters/test_sqlite/test_adk_config.py index 22c846449..06b47814f 100644 --- a/tests/unit/adapters/test_sqlite/test_adk_config.py +++ b/tests/unit/adapters/test_sqlite/test_adk_config.py @@ -9,6 +9,7 @@ from sqlspec.adapters.sqlite.config import SqliteConfig from sqlspec.config import ADKConfig from sqlspec.exceptions import ImproperConfigurationError +from sqlspec.extensions.adk import StoredMemory class RecordingConnection: @@ -189,3 +190,76 @@ def test_sqlite_list_sessions_rejects_invalid_options_before_connecting(options: store.list_sessions("app", **options) assert conn.calls == [] + + +def test_sqlite_adk_get_events_with_limit_binds_parameter() -> None: + """Verify get_events parameterizes LIMIT instead of string interpolation.""" + config = SqliteConfig(connection_config={"database": ":memory:"}) + store = SqliteADKStore(config) + calls: list[tuple[str, tuple[Any, ...]]] = [] + + class _MockEventsCursor: + def fetchall(self) -> list[Any]: + return [] + + class _MockEventsConnection: + def execute(self, sql: str, params: Any = None) -> _MockEventsCursor: + calls.append((sql, params)) + return _MockEventsCursor() + + class _SyncContextManager: + def __enter__(self) -> Any: + return _MockEventsConnection() + + def __exit__(self, *_: Any) -> None: + return None + + config.provide_connection = MagicMock(return_value=_SyncContextManager()) # type: ignore[method-assign] + + events = store.get_events("app", "user", "session", limit=10) + assert events == [] + assert len(calls) >= 1 + sql, params = calls[-1] + assert "LIMIT ?" in sql + assert 10 in params + + +def test_sqlite_adk_memory_store_persists_and_retrieves_float32_embeddings() -> None: + """Verify memory entries persist and retrieve float32 embeddings.""" + from datetime import UTC, datetime + from uuid import uuid4 + + config = SqliteConfig(connection_config={"database": ":memory:"}) + store = SqliteADKMemoryStore(config) + try: + store.create_tables() + + now = datetime.now(UTC) + embedding = [0.1, 0.2, 0.3, 0.4] + record = StoredMemory( + id=str(uuid4()), + session_id="s1", + app_name="app", + user_id="u1", + scope="user", + event_id="evt-1", + author="user", + timestamp=now, + content_json={"text": "vector search test"}, + content_text="vector search test", + metadata_json=None, + inserted_at=now, + embedding=embedding, + ) + inserted = store.insert_memory_entries([record]) + assert inserted == 1 + + results = store.search_entries(query="vector", app_name="app", user_id="u1", embedding=embedding) + assert len(results) == 1 + assert results[0]["event_id"] == "evt-1" + assert results[0]["embedding"] is not None + assert len(results[0]["embedding"]) == 4 + for orig, rec in zip(embedding, results[0]["embedding"], strict=True): + assert abs(orig - rec) < 1e-6 + finally: + config.close_pool() diff --git a/tests/unit/adapters/test_sqlite/test_config.py b/tests/unit/adapters/test_sqlite/test_config.py index 7d607481c..f6b3b8d2a 100644 --- a/tests/unit/adapters/test_sqlite/test_config.py +++ b/tests/unit/adapters/test_sqlite/test_config.py @@ -1,5 +1,6 @@ """SQLite configuration normalization tests.""" +import threading from os import PathLike from pathlib import Path from typing import Literal, get_args, get_origin, get_type_hints @@ -275,3 +276,36 @@ def test_create_connection_returns_a_caller_owned_connection(tmp_path: Path) -> pooled.execute("SELECT 1") finally: config.close_pool() + + +def test_default_memory_database_uses_shared_cache() -> None: + config = SqliteConfig() + assert "cache=shared" in config.connection_config["database"] + pool = config.provide_pool() + try: + with pool.get_connection() as conn: + conn.execute("CREATE TABLE shared_probe (id INT)") + conn.commit() + + results: list[int] = [] + + def worker() -> None: + with pool.get_connection() as conn: + row = conn.execute("SELECT count(*) FROM shared_probe").fetchone() + if row: + results.append(row[0]) + + t = threading.Thread(target=worker) + t.start() + t.join() + assert results == [0] + finally: + pool.close() + + +def test_filter_connection_params_with_hoisted_helper() -> None: + from sqlspec.adapters.sqlite.core import _filter_connection_params + + raw = {"database": ":memory:", "extra": {}, "pool_min_size": 10, "timeout": 30.0, "isolation_level": None} + filtered = _filter_connection_params(raw) + assert filtered == {"database": ":memory:", "timeout": 30.0, "isolation_level": None} diff --git a/tests/unit/adapters/test_sqlite/test_driver.py b/tests/unit/adapters/test_sqlite/test_driver.py index 7d0b46fc9..c5b36e877 100644 --- a/tests/unit/adapters/test_sqlite/test_driver.py +++ b/tests/unit/adapters/test_sqlite/test_driver.py @@ -5,7 +5,7 @@ import pytest from sqlspec.adapters.sqlite import SqliteConfig, SqliteDriver -from sqlspec.adapters.sqlite.core import SQLITE_CONNECT_SUPPORTS_AUTOCOMMIT +from sqlspec.adapters.sqlite.core import SQLITE_CONNECT_SUPPORTS_AUTOCOMMIT, default_statement_config from sqlspec.adapters.sqlite.pool import SqliteConnectionPool from sqlspec.adapters.sqlite.type_converter import json_converter, register_type_handlers, unregister_type_handlers from sqlspec.exceptions import ImproperConfigurationError @@ -271,7 +271,7 @@ def test_sqlite_begin_mode_and_arrow_batch_validation() -> None: driver.rollback() driver.begin(mode="DEFERRED") driver.rollback() - with pytest.raises(ValueError, match="Transaction mode"): + with pytest.raises((ValueError, ImproperConfigurationError), match="Transaction mode"): driver.begin(mode=cast("Any", "INVALID")) with pytest.raises(ValueError, match="batch_size"): driver.load_from_arrow("sample", cast("Any", None), batch_size=0) @@ -342,3 +342,87 @@ def test_sqlite_json_converter_and_unregister() -> None: register_type_handlers() unregister_type_handlers() assert (dict, sqlite3.PrepareProtocol) not in sqlite3.adapters + + +def test_sqlite_dispatch_execute_script_emits_events() -> None: + connection = sqlite3.connect(":memory:") + try: + driver = SqliteDriver(connection=connection) + result = driver.execute_script("CREATE TABLE scr_sync (x INT); INSERT INTO scr_sync VALUES (1);") + assert result.total_statements == 2 + assert result.successful_statements == 2 + finally: + connection.close() + + +def test_sqlite_savepoint_does_not_clear_rowid_target_cache() -> None: + connection = sqlite3.connect(":memory:") + try: + driver = SqliteDriver(connection=connection) + driver._rowid_target_cache[(None, "test_target")] = True + driver.create_savepoint("sp1") + assert (None, "test_target") in driver._rowid_target_cache + driver.release_savepoint("sp1") + assert (None, "test_target") in driver._rowid_target_cache + finally: + connection.close() + + +def test_sqlite_cached_query_normalizes_parameters() -> None: + connection = sqlite3.connect(":memory:") + try: + driver = SqliteDriver(connection=connection) + driver.execute("CREATE TABLE norm_test (val INT)") + driver.execute("SELECT * FROM norm_test WHERE val = ?", [1]) + res = driver.execute("SELECT * FROM norm_test WHERE val = ?", [1]) + assert res.data == [] + finally: + connection.close() + + +def test_sqlite_begin_defaults_to_immediate_and_rejects_invalid_mode() -> None: + connection = sqlite3.connect(":memory:") + try: + driver = SqliteDriver(connection=connection) + with pytest.raises(ImproperConfigurationError, match="Transaction mode"): + driver.begin(mode=cast("Any", "INVALID")) + driver.begin() + assert connection.in_transaction + driver.rollback() + finally: + connection.close() + + +def test_sqlite_commit_and_rollback_delegate_to_end_transaction() -> None: + connection = sqlite3.connect(":memory:") + try: + driver = SqliteDriver(connection=connection) + driver.execute("CREATE TABLE cr_test (x INT)") + driver.begin() + driver.execute("INSERT INTO cr_test VALUES (1)") + driver.commit() + assert not connection.in_transaction + + driver.begin() + driver.execute("INSERT INTO cr_test VALUES (2)") + driver.rollback() + assert not connection.in_transaction + + res = driver.execute("SELECT count(*) FROM cr_test") + assert res.data == [(1,)] + finally: + connection.close() + + +def test_sqlite_session_context_initializes_with_driver_class() -> None: + from sqlspec.adapters.sqlite._typing import SqliteSessionContext + + ctx = SqliteSessionContext( + acquire_connection=lambda: None, + release_connection=lambda *_: None, + statement_config=default_statement_config, + driver_features={}, + prepare_driver=lambda d: d, + driver_class=SqliteDriver, + ) + assert ctx._driver_class is SqliteDriver diff --git a/tests/unit/adapters/test_sqlite/test_exception_mapping.py b/tests/unit/adapters/test_sqlite/test_exception_mapping.py index c77a81b57..b2eb60a9e 100644 --- a/tests/unit/adapters/test_sqlite/test_exception_mapping.py +++ b/tests/unit/adapters/test_sqlite/test_exception_mapping.py @@ -5,7 +5,17 @@ import pytest from sqlspec.adapters.sqlite.core import create_mapped_exception -from sqlspec.exceptions import DeadlockError, OperationCancelledError, PermissionDeniedError, UniqueViolationError +from sqlspec.exceptions import ( + CheckViolationError, + DataError, + DeadlockError, + IntegrityError, + NotNullViolationError, + OperationalError, + OperationCancelledError, + PermissionDeniedError, + UniqueViolationError, +) class _SqliteError(sqlite3.OperationalError): @@ -88,3 +98,54 @@ def test_primary_key_identity_takes_precedence_over_message(table: str, code: in mapped = create_mapped_exception(error) assert isinstance(mapped, UniqueViolationError) assert mapped.__cause__ is error + + +def test_not_null_error_code_maps_to_not_null_violation() -> None: + err = _SqliteIntegrityError("NOT NULL constraint failed: t.val", 1299, None) + result = create_mapped_exception(err) + assert isinstance(result, NotNullViolationError) + assert result.__cause__ is err + + +def test_check_error_code_maps_to_check_violation() -> None: + err = _SqliteIntegrityError("CHECK constraint failed: val > 0", 275, None) + result = create_mapped_exception(err) + assert isinstance(result, CheckViolationError) + assert result.__cause__ is err + + +def test_trigger_abort_does_not_map_to_not_null_violation() -> None: + err = _SqliteIntegrityError("trigger aborted", 1811, None) + result = create_mapped_exception(err) + assert not isinstance(result, NotNullViolationError) + assert isinstance(result, IntegrityError) + assert result.__cause__ is err + + +def test_no_such_column_with_locked_name_maps_to_operational_error() -> None: + err = sqlite3.OperationalError("no such column: is_locked") + result = create_mapped_exception(err) + assert isinstance(result, OperationalError) + assert not isinstance(result, DeadlockError) + assert result.__cause__ is err + + +def test_programming_error_maps_to_data_error() -> None: + err = sqlite3.ProgrammingError("cannot operate on a closed database") + result = create_mapped_exception(err) + assert isinstance(result, DataError) + assert result.__cause__ is err + + +def test_interface_error_maps_to_operational_error() -> None: + err = sqlite3.InterfaceError("error binding parameter") + result = create_mapped_exception(err) + assert isinstance(result, OperationalError) + assert result.__cause__ is err + + +def test_misuse_maps_to_data_error() -> None: + err = _SqliteError("bad call", 21, "SQLITE_MISUSE") + result = create_mapped_exception(err) + assert isinstance(result, DataError) + assert result.__cause__ is err diff --git a/tests/unit/adapters/test_sqlite/test_pool.py b/tests/unit/adapters/test_sqlite/test_pool.py index 8b55d89f8..735966e76 100644 --- a/tests/unit/adapters/test_sqlite/test_pool.py +++ b/tests/unit/adapters/test_sqlite/test_pool.py @@ -46,6 +46,25 @@ def test_default_pool_uses_native_like_shared_pragma_profile() -> None: pool.close() +def test_default_memory_connections_do_not_read_uncommitted_rows() -> None: + from sqlspec.adapters.sqlite import SqliteConfig + + config = SqliteConfig() + other = config.provide_pool().new_connection() + try: + with config.provide_connection() as connection: + connection.execute("CREATE TABLE isolated_rows (value INTEGER)") + connection.execute("BEGIN") + connection.execute("INSERT INTO isolated_rows VALUES (1)") + with pytest.raises(sqlite3.OperationalError, match="locked"): + other.execute("SELECT value FROM isolated_rows") + connection.rollback() + assert other.execute("SELECT COUNT(*) FROM isolated_rows").fetchone() == (0,) + finally: + other.close() + config.close_pool() + + def test_disable_optimizations_preserves_native_memory_profile() -> None: native_connection = sqlite3.connect(":memory:", timeout=30.0) pool = SqliteConnectionPool({"database": ":memory:", "timeout": 30.0}, enable_optimizations=False) @@ -172,6 +191,31 @@ def test_get_connection_commits_open_transaction_on_clean_exit() -> None: pool.close() +def test_get_connection_exits_cleanly_on_another_thread() -> None: + """A scope entered on one thread can be exited on another without losing the commit.""" + pool = SqliteConnectionPool(connection_parameters={"database": ":memory:"}, enable_optimizations=False) + scope = pool.get_connection() + connection = scope.__enter__() + connection.execute("CREATE TABLE widgets (id INTEGER PRIMARY KEY)") + connection.execute("INSERT INTO widgets (id) VALUES (1)") + errors: list[BaseException] = [] + + def _exit() -> None: + try: + scope.__exit__(None, None, None) + except BaseException as exc: + errors.append(exc) + + worker = threading.Thread(target=_exit) + worker.start() + worker.join(timeout=10) + + assert errors == [] + assert pool.checked_out() == 0 + assert connection.execute("SELECT COUNT(*) FROM widgets").fetchone()[0] == 1 + pool.close() + + def test_close_closes_connections_opened_on_other_threads(tmp_path: Path) -> None: """Connections opened by worker threads must not survive pool shutdown.""" pool = SqliteConnectionPool({"database": str(tmp_path / "threads.sqlite")}) @@ -389,3 +433,45 @@ def test_memory_pool_replaces_a_closed_connection() -> None: assert replacement.execute("SELECT 1").fetchone() == (1,) finally: pool.close() + + +def test_nested_get_connection_reentrancy_depth() -> None: + pool = SqliteConnectionPool({"database": ":memory:"}) + try: + with pool.get_connection() as conn1: + conn1.execute("CREATE TABLE test_depth (x INT)") + conn1.commit() + conn1.execute("INSERT INTO test_depth VALUES (1)") + with pool.get_connection() as conn2: + assert conn1 is conn2 + conn2.execute("INSERT INTO test_depth VALUES (2)") + assert getattr(conn1, "in_transaction", False) + conn1.commit() + + with pool.get_connection() as conn: + rows = conn.execute("SELECT * FROM test_depth").fetchall() + assert rows == [(1,), (2,)] + finally: + pool.close() + + +def test_sqlite_pool_checked_out_reflects_active_sessions_and_prunes_dead_threads() -> None: + pool = SqliteConnectionPool({"database": ":memory:"}) + try: + assert pool.checked_out() == 0 + with pool.get_connection(): + assert pool.checked_out() == 1 + assert pool.checked_out() == 0 + + def worker() -> None: + with pool.get_connection() as conn: + conn.execute("SELECT 1") + + t = threading.Thread(target=worker) + t.start() + t.join() + + assert not t.is_alive() + assert pool.size() == 1 + finally: + pool.close() diff --git a/tests/unit/adapters/test_sqlite/test_type_handlers.py b/tests/unit/adapters/test_sqlite/test_type_handlers.py index 6e6fbe78b..7a962b829 100644 --- a/tests/unit/adapters/test_sqlite/test_type_handlers.py +++ b/tests/unit/adapters/test_sqlite/test_type_handlers.py @@ -53,3 +53,18 @@ def test_sqlite_config_does_not_register_type_handlers_by_default() -> None: config._create_pool() # pyright: ignore[reportPrivateUsage] mock_register.assert_not_called() + + +def test_julian_conversions_match_hoisted_epoch() -> None: + """Verify datetime_to_julian and julian_to_datetime use hoisted epoch.""" + from datetime import UTC, datetime + + from sqlspec.adapters.sqlite.core import _UNIX_EPOCH, datetime_to_julian, julian_to_datetime + + now = datetime(2026, 1, 1, 12, 0, 0, tzinfo=UTC) + julian = datetime_to_julian(now) + assert julian is not None + recovered = julian_to_datetime(julian) + assert recovered is not None + assert recovered == now + assert datetime_to_julian(_UNIX_EPOCH) == 2440587.5 diff --git a/tests/unit/adapters/test_sybase.py b/tests/unit/adapters/test_sybase.py index ef23069b8..c7e0eba43 100644 --- a/tests/unit/adapters/test_sybase.py +++ b/tests/unit/adapters/test_sybase.py @@ -3,10 +3,12 @@ import os import socket import threading +from contextlib import nullcontext from datetime import UTC, datetime from typing import Any, cast from unittest.mock import AsyncMock, MagicMock +import pymssql import pytest from sqlspec.adapters.sybase import ( @@ -26,8 +28,14 @@ driver_profile, ) from sqlspec.adapters.sybase._typing import SybaseRawCursor -from sqlspec.adapters.sybase.adk import SybaseADKConfig, SybaseADKMemoryStore, SybaseADKStore -from sqlspec.adapters.sybase.events import SybaseEventQueueStore, SybaseSyncEventQueueStore +from sqlspec.adapters.sybase.adk import SybaseADKMemoryStore, SybaseADKStore +from sqlspec.adapters.sybase.core import ( + _ERROR_CODE_MAPPING, + build_insert_statement, + escape_literal_percent, + format_identifier, +) +from sqlspec.adapters.sybase.events import SybaseEventQueueStore from sqlspec.adapters.sybase.litestar import SybaseStore from sqlspec.adapters.sybase.migrations import SybaseSyncMigrationTracker from sqlspec.core import SQL, ParameterStyle @@ -45,6 +53,7 @@ PermissionDeniedError, QueryTimeoutError, SQLParsingError, + SQLSpecError, TransactionError, UniqueViolationError, ) @@ -158,7 +167,7 @@ class FakeModule: def connect(self, **kwargs: Any) -> FakeConnection: return fake_conn - monkeypatch.setattr(pool_module, "pymssql", FakeModule()) + monkeypatch.setattr(pool_module, "PYMSSQL_MODULE", FakeModule()) monkeypatch.delenv("TDSVER", raising=False) cfg_invalid = build_connection_config({"server": "ase.example.test", "user": "sa", "database": "master"}) @@ -411,7 +420,7 @@ def select(self, _sql: Any, **kwargs: Any) -> list[dict[str, Any]]: def test_sybase_migration_tracker_and_extension_stores() -> None: - """SybaseSyncMigrationTracker, SybaseStore, SybaseSyncEventQueueStore, and SybaseADKStore emit ASE-safe SQL.""" + """SybaseSyncMigrationTracker, SybaseStore, SybaseEventQueueStore, and SybaseADKStore emit ASE-safe SQL.""" tracker = SybaseSyncMigrationTracker(version_table_name="dbo.ddl_migrations") ddl_text = tracker._idempotent_tracking_table_ddl_text() assert "IF OBJECT_ID(" in ddl_text @@ -437,8 +446,8 @@ def test_sybase_migration_tracker_and_extension_stores() -> None: assert "BIGDATETIME" in litestar_ddl assert "GETUTCDATE()" in litestar_ddl - assert SybaseEventQueueStore is SybaseSyncEventQueueStore - event_store = SybaseSyncEventQueueStore(config) + assert SybaseEventQueueStore is SybaseEventQueueStore + event_store = SybaseEventQueueStore(config) event_stmts = event_store.create_statements() assert "TEXT" in event_stmts[0] assert "BIGDATETIME" in event_stmts[0] @@ -463,7 +472,6 @@ def test_sybase_migration_tracker_and_extension_stores() -> None: assert SybaseCursor.__slots__ == ("_as_dict", "connection", "cursor") assert SybaseExceptionHandler.__slots__ == () - assert isinstance(SybaseADKConfig(), dict) @pytest.mark.anyio @@ -519,3 +527,145 @@ def _capture(_self: SybaseADKStore, sql: str, params: "tuple[Any, ...]" = ()) -> ) inserted = mem_store.insert_memory_entries([entry]) assert inserted == 1 + + +def test_sybase_native_error_handling_and_missing_table(monkeypatch: pytest.MonkeyPatch) -> None: + """Native pymssql errors in commit() are caught and ADK get_session() returns None on table missing.""" + cursor = FakeCursor() + conn = FakeConnection(cursor) + driver = SybaseDriver(cast("SybaseConnection", conn)) + + def _fail_commit() -> None: + raise pymssql.OperationalError("Native pymssql commit failure") + + conn.commit = _fail_commit + with pytest.raises(SQLSpecError, match="Failed to commit Sybase ASE transaction"): + driver.commit() + + config = MagicMock() + config.extension_config = {"adk": {}} + store = SybaseADKStore(cast("Any", config)) + + def _table_missing(_self: Any, *_args: Any, **_kwargs: Any) -> None: + raise pymssql.DatabaseError(208, "Invalid object name 'adk_session'") + + monkeypatch.setattr(SybaseADKStore, "_execute_fetchone", _table_missing) + assert store.get_session("app", "user", "sess") is None + + +def test_sybase_commit_maps_deadlock_error_1205() -> None: + """Simulate error 1205 during driver.commit(); assert DeadlockError is raised.""" + cursor = FakeCursor() + conn = FakeConnection(cursor) + driver = SybaseDriver(cast("SybaseConnection", conn)) + + def _fail_commit() -> None: + raise pymssql.OperationalError(1205, "Your transaction was deadlock victim (1205)") + + conn.commit = _fail_commit + with pytest.raises(DeadlockError, match="Failed to commit Sybase ASE transaction"): + driver.commit() + + +def test_sybase_adk_store_maps_unique_violation(monkeypatch: pytest.MonkeyPatch) -> None: + """SybaseADKStore raises UniqueViolationError on duplicate key violation.""" + config = MagicMock() + config.extension_config = {"adk": {}} + store = SybaseADKStore(cast("Any", config)) + + class DummyCursor: + def execute(self, _sql: str, _params: Any = ()) -> None: + raise pymssql.IntegrityError(2601, "Attempt to insert duplicate key row with unique index") + + def close(self) -> None: + pass + + class DummyConnection: + def cursor(self, as_dict: bool = False) -> DummyCursor: + return DummyCursor() + + def rollback(self) -> None: + pass + + monkeypatch.setattr(config, "provide_connection", lambda: nullcontext(DummyConnection())) + with pytest.raises(UniqueViolationError): + store._execute("INSERT INTO dummy VALUES (1)") + + +def test_sybase_litestar_store_get_raises_sqlspec_error(monkeypatch: pytest.MonkeyPatch) -> None: + """SybaseStore._get() maps raw database errors to SQLSpecError.""" + config = MagicMock() + config.extension_config = {"litestar": {}} + store = SybaseStore(cast("Any", config)) + + class FailingCursor: + def execute(self, _sql: str, _params: Any = ()) -> None: + raise pymssql.OperationalError("Connection lost") + + def close(self) -> None: + pass + + class FailingConnection: + def cursor(self) -> FailingCursor: + return FailingCursor() + + monkeypatch.setattr(config, "provide_connection", lambda: nullcontext(FailingConnection())) + with pytest.raises(SQLSpecError): + store._get("session-1") + + +def test_sybase_has_schema_under_active_exception(monkeypatch: pytest.MonkeyPatch) -> None: + """has_schema runs cleanly inside an active exception block without isolation leak.""" + config = MagicMock() + driver = SybaseDriver(config) + + class DummyCursor: + def execute(self, _sql: str, _params: Any = ()) -> None: + pass + + def fetchone(self) -> Any: + return (1,) + + def close(self) -> None: + pass + + class DummyConnection: + def cursor(self) -> DummyCursor: + return DummyCursor() + + monkeypatch.setattr(driver, "connection", DummyConnection()) + + try: + raise RuntimeError("active error") + except RuntimeError: + has_user = driver.has_schema("dbo") + assert has_user is True + + +def test_sybase_format_identifier_and_insert_statements() -> None: + """format_identifier, build_insert_statement, and escape_literal_percent format queries correctly.""" + assert format_identifier("dbo.my_table") == "[dbo].[my_table]" + sql = build_insert_statement("my_table", ["col1", "col2"]) + assert sql == "INSERT INTO [my_table] ([col1], [col2]) VALUES (%s, %s)" + escaped = escape_literal_percent( + "SELECT * FROM t WHERE val LIKE '100%' AND id = %s", ("1",), default_statement_config.parameter_validator + ) + assert "100%%" in escaped + + +def test_sybase_migration_tracker_existing_columns_query() -> None: + """_existing_columns_query resolves USER_NAME() when schema is unqualified.""" + tracker_unqualified = SybaseSyncMigrationTracker(version_table_name="my_versions") + query_unqualified = tracker_unqualified._existing_columns_query() + assert "u.name = USER_NAME()" in query_unqualified + assert "o.name = 'my_versions'" in query_unqualified + + tracker_qualified = SybaseSyncMigrationTracker(version_table_name="custom_schema.my_versions") + query_qualified = tracker_qualified._existing_columns_query() + assert "u.name = 'custom_schema'" in query_qualified + assert "o.name = 'my_versions'" in query_qualified + + +def test_sybase_error_code_mapping_count() -> None: + """_ERROR_CODE_MAPPING contains 38 error codes.""" + assert len(_ERROR_CODE_MAPPING) == 38 diff --git a/tests/unit/builder/test_ddl_metadata_compiler.py b/tests/unit/builder/test_ddl_metadata_compiler.py index c6693994e..204b71ec7 100644 --- a/tests/unit/builder/test_ddl_metadata_compiler.py +++ b/tests/unit/builder/test_ddl_metadata_compiler.py @@ -24,6 +24,37 @@ from sqlspec.exceptions import SQLBuilderError +def test_create_table_from_metadata_foreign_key_preserves_schemas_and_actions() -> None: + """Owned foreign key metadata preserves reference qualification and both actions.""" + foreign_key = ForeignKeyMetadata( + table_name="orders", + column_name="customer_id", + referenced_table="customers", + referenced_column="id", + constraint_name="fk_orders_customer", + schema="sales", + referenced_schema="crm", + on_delete="CASCADE", + on_update="SET NULL", + ) + + built = CreateTable.from_metadata( + "orders", + [{"column_name": "customer_id", "data_type": "INTEGER"}], + foreign_keys=[foreign_key], + schema="sales", + dialect="postgres", + ).build() + + assert built.sql == ( + 'CREATE TABLE "sales"."orders" (\n' + ' "customer_id" INT,\n' + ' CONSTRAINT "fk_orders_customer" FOREIGN KEY ("customer_id") ' + 'REFERENCES "crm"."customers"("id") ON DELETE CASCADE ON UPDATE SET NULL\n' + ")" + ) + + def test_create_table_from_metadata_composite_pk_fk_check_unique_defaults_and_computed() -> None: """Test CreateTable.from_metadata with composite PK, FK, CHECK, UNIQUE, defaults, identity, and computed columns.""" columns: list[dict[str, Any]] = [ diff --git a/tests/unit/core/test_cache.py b/tests/unit/core/test_cache.py index 31caeba6d..12fa9cf0a 100644 --- a/tests/unit/core/test_cache.py +++ b/tests/unit/core/test_cache.py @@ -16,7 +16,6 @@ across the entire SQLSpec system. """ -import inspect import logging import threading import time @@ -451,21 +450,6 @@ def test_filters_view_external_api_smoke() -> None: assert len(view.to_canonical()) == 3 -def test_cache_external_api_docstring_markers() -> None: - """Dead-code policy markers identify externally callable cache helpers.""" - for member in ( - NamespacedCache.delete_statement, - NamespacedCache.delete_expression, - NamespacedCache.delete_optimized, - NamespacedCache.delete_builder, - NamespacedCache.delete_file, - cache_module.FiltersView.get_by_field, - cache_module.FiltersView.has_field, - cache_module.FiltersView.to_canonical, - ): - assert "External/extension API" in (inspect.getdoc(member) or "") - - def test_namespaced_cache_respects_config_flags() -> None: """Test that cache config flags disable namespaces.""" original_config = get_cache_config() diff --git a/tests/unit/core/test_cursor_filter.py b/tests/unit/core/test_cursor_filter.py index 42bfcdc8b..8fde202c3 100644 --- a/tests/unit/core/test_cursor_filter.py +++ b/tests/unit/core/test_cursor_filter.py @@ -90,14 +90,19 @@ def test_composite_cursor_key_shorthand(keys: Any, expected: list[CursorKey]) -> assert shorthand.encode({"v": 10, "id": 3}, backward=True) == explicit.encode({"v": 10, "id": 3}, backward=True) -@pytest.mark.parametrize( - "keys", ["", [""], [("id", "bad")], [("id",)], [("id", "asc", "last")], [3], ["id", CursorKey("id")]] -) +@pytest.mark.parametrize("keys", ["", [""], [("id", "bad")], ["id", CursorKey("id")]]) def test_invalid_cursor_key_shorthand(keys: Any) -> None: with pytest.raises(ValueError): CursorFilter(keys, 2) +@pytest.mark.parametrize("keys", [[("id",)], [("id", "asc", "last")], [3]]) +def test_cursor_key_shorthand_rejects_unsupported_input_shapes(keys: Any) -> None: + """Inputs outside CursorKeys are rejected by validation or compiled type guards.""" + with pytest.raises((TypeError, ValueError)): + CursorFilter(keys, 2) + + def test_eager_cursor_validation() -> None: with pytest.raises(InvalidCursorError): CursorFilter([CursorKey("id")], 10, "invalid") diff --git a/tests/unit/core/test_data_dictionary.py b/tests/unit/core/test_data_dictionary.py index c968aa666..929d37411 100644 --- a/tests/unit/core/test_data_dictionary.py +++ b/tests/unit/core/test_data_dictionary.py @@ -1,8 +1,10 @@ """Unit tests for the data dictionary registry and loader.""" import re +import subprocess +import sys from importlib import resources -from typing import Any, cast +from typing import Any import pytest @@ -10,7 +12,6 @@ from sqlspec.core import SQL from sqlspec.data_dictionary import ( DataDictionaryLoader, - DialectConfig, VersionInfo, get_data_dictionary_loader, get_dialect_config, @@ -151,11 +152,6 @@ def test_row_locking_capability_flags_are_introspectable() -> None: assert get_dialect_config("mysql").get_feature_version("supports_skip_locked") == VersionInfo(8, 0, 1) -def test_registry_dialects_loaded_annotation_is_bool() -> None: - """_DIALECTS_LOADED should be annotated to avoid Literal[False] narrowing.""" - assert registry.__annotations__["_DIALECTS_LOADED"] is bool - - def test_registry_load_default_dialects_sets_loaded_flag(monkeypatch) -> None: """_load_default_dialects sets the loaded flag after importing dialects.""" calls: list[str] = [] @@ -177,16 +173,18 @@ def fail_import(name: str) -> None: registry._load_default_dialects() -def test_registry_get_dialect_config_triggers_load(monkeypatch) -> None: - """get_dialect_config calls the default dialect loader before lookup.""" - calls = 0 - config = cast("DialectConfig", object()) - monkeypatch.setitem(registry._DIALECT_CONFIGS, "example", config) - - def fake_load() -> None: - nonlocal calls - calls += 1 - - monkeypatch.setattr(registry, "_load_default_dialects", fake_load) - assert registry.get_dialect_config("example") is config - assert calls == 1 +def test_registry_get_dialect_config_triggers_load() -> None: + """The first public lookup loads built-in dialect configurations.""" + script = """ +import sys +from sqlspec.data_dictionary import DialectConfig, get_dialect_config + +assert "sqlspec.data_dictionary.dialects" not in sys.modules +config = get_dialect_config("PostgreSQL") +assert isinstance(config, DialectConfig) +assert config.name == "postgres" +assert config.get_feature_flag("supports_transactions") is True +assert "sqlspec.data_dictionary.dialects" in sys.modules +assert get_dialect_config("pg") is config +""" + subprocess.run([sys.executable, "-c", script], capture_output=True, text=True, check=True) diff --git a/tests/unit/core/test_declared_params_carriage.py b/tests/unit/core/test_declared_params_carriage.py index cb53336be..31faf6adb 100644 --- a/tests/unit/core/test_declared_params_carriage.py +++ b/tests/unit/core/test_declared_params_carriage.py @@ -4,9 +4,15 @@ plus loader population and driver-derivation preservation. """ +import copy +import pickle + +import pytest + from sqlspec.core import ParameterDeclaration from sqlspec.core._pool import get_sql_pool from sqlspec.core.statement import SQL +from sqlspec.loader import SQLFileLoader _SENTINEL = (ParameterDeclaration("a", "int", required=False),) @@ -60,8 +66,6 @@ def test_pool_recycle_does_not_leak_declared_parameters() -> None: def test_get_sql_populates_declared_parameters() -> None: - from sqlspec.loader import SQLFileLoader - loader = SQLFileLoader() loader.add_named_sql("q", "select :a", parameters=[ParameterDeclaration("a", "int", description="id")]) sql = loader.get_sql("q") @@ -69,8 +73,6 @@ def test_get_sql_populates_declared_parameters() -> None: def test_undeclared_get_sql_has_empty_declarations() -> None: - from sqlspec.loader import SQLFileLoader - loader = SQLFileLoader() loader.add_named_sql("plain", "select 1") assert loader.get_sql("plain").declared_parameters == () @@ -81,20 +83,25 @@ def test_constructor_kwarg_sets_declarations() -> None: assert sql.declared_parameters == _SENTINEL -def test_declarations_survive_deepcopy() -> None: - import copy - - sql = SQL("select :a", {"a": 1}, declared_parameters=_SENTINEL) - new = copy.deepcopy(sql) - assert new.declared_parameters == _SENTINEL - - -def test_declarations_survive_pickle_roundtrip() -> None: - import pickle - - sql = SQL("select :a", {"a": 1}, declared_parameters=_SENTINEL) - new = pickle.loads(pickle.dumps(sql)) - assert new.declared_parameters == _SENTINEL +@pytest.mark.parametrize("operation", ["pickle", "deepcopy"]) +@pytest.mark.parametrize("source", ["constructor", "loader"]) +def test_declarations_survive_serialization(operation: str, source: str) -> None: + """Public SQL round trips preserve declaration types and all metadata fields.""" + declarations = ( + ParameterDeclaration("a", "int", "Account identifier", required=False), + ParameterDeclaration("b", "str", required=True), + ) + if source == "loader": + loader = SQLFileLoader() + loader.add_named_sql("declared", "SELECT :a, :b", parameters=list(declarations)) + sql = loader.get_sql("declared").copy(parameters={"a": 1, "b": "group"}) + else: + sql = SQL("SELECT :a, :b", a=1, b="group", declared_parameters=declarations) + + restored = pickle.loads(pickle.dumps(sql)) if operation == "pickle" else copy.deepcopy(sql) + + assert restored.declared_parameters == declarations + assert all(type(declaration) is ParameterDeclaration for declaration in restored.declared_parameters) def test_declarations_survive_driver_prepare_with_filter() -> None: diff --git a/tests/unit/core/test_explain_classification.py b/tests/unit/core/test_explain_classification.py index 8988b47c5..d3d3f29d4 100644 --- a/tests/unit/core/test_explain_classification.py +++ b/tests/unit/core/test_explain_classification.py @@ -95,3 +95,51 @@ def test_sql_explain_oracle_stays_non_row_returning() -> None: assert statement.sql == "EXPLAIN PLAN FOR SELECT 1" assert statement.returns_rows() is False + + +@pytest.mark.parametrize( + "sql", + [ + "INSERT INTO t (a) OUTPUT inserted.id INTO @ids VALUES (1)", + "UPDATE t SET a = 1 OUTPUT inserted.id INTO @ids WHERE id = 1", + "DELETE t OUTPUT deleted.id INTO @ids WHERE id = 1", + ], +) +def test_tsql_output_into_does_not_return_client_rows(sql: str) -> None: + """OUTPUT ... INTO writes into a table variable, so no rows reach the client.""" + statement = SQL(sql, statement_config=StatementConfig(dialect="tsql")) + + assert statement.returns_rows() is False + + +def test_tsql_output_without_into_returns_client_rows() -> None: + """A bare OUTPUT clause streams the affected rows to the client.""" + statement = SQL("INSERT INTO t (a) OUTPUT inserted.id VALUES (1)", statement_config=StatementConfig(dialect="tsql")) + + assert statement.returns_rows() is True + + +def test_oracle_returning_into_does_not_return_client_rows() -> None: + """RETURNING INTO fills output variables rather than a client result set.""" + statement = SQL( + "INSERT INTO t (a) VALUES (1) RETURNING id INTO :result", + result=0, + statement_config=StatementConfig(dialect="oracle"), + ) + + assert statement.returns_rows() is False + + +@pytest.mark.parametrize( + "sql", + [ + "INSERT INTO t (a) VALUES (1) RETURNING id", + "UPDATE t SET a = 1 WHERE id = 1 RETURNING id", + "DELETE FROM t WHERE id = 1 RETURNING id", + ], +) +def test_postgres_returning_keeps_client_rows(sql: str) -> None: + """Ordinary RETURNING statements still expose the affected rows to callers.""" + statement = SQL(sql, statement_config=StatementConfig(dialect="postgres")) + + assert statement.returns_rows() is True diff --git a/tests/unit/core/test_filters.py b/tests/unit/core/test_filters.py index cb68201b2..f52b5cbe6 100644 --- a/tests/unit/core/test_filters.py +++ b/tests/unit/core/test_filters.py @@ -7,11 +7,10 @@ import tempfile from dataclasses import dataclass from datetime import datetime -from pathlib import Path from typing import TYPE_CHECKING, Any, assert_type, cast import pytest -from sqlglot import exp +from sqlglot import exp, parse_one import sqlspec.core.filters as filters_module from sqlspec import sql as sql_builder @@ -61,31 +60,15 @@ def test_public_filters_alias_points_to_core_filters_module() -> None: assert sqlspec.filters is filters_module -def test_filters_docs_render_inherited_members() -> None: - """Docs include public filters whose methods are inherited.""" - docs = Path("docs/reference/core/filters.md").read_text() - for class_name in ("BeforeAfterFilter", "SearchFilter", "InCollectionFilter", "NotInCollectionFilter"): - assert f"::: sqlspec.core.filters.{class_name}" in docs +def test_search_filter_preserves_qualified_column_and_binds_value() -> None: + """A qualified field remains qualified in the rendered search predicate.""" + result = apply_filter(SQL("SELECT * FROM users"), SearchFilter("users.name", "needle")) + condition = parse_one(result.raw_sql).find(exp.Like) - -def test_statement_filter_get_column_expression_delegates_to_parse_column_for_condition( - monkeypatch: pytest.MonkeyPatch, -) -> None: - """StatementFilter should use the shared fast column parser.""" - calls: list[str | exp.Expression] = [] - original = StatementFilter._get_column_expression - - def fake_parse(field_name: str | exp.Expression) -> exp.Expr: - calls.append(field_name) - return exp.column("parsed") - - monkeypatch.setattr("sqlspec.core.filters.parse_column_for_condition", fake_parse) - - result = original(SearchFilter("name", "needle"), "users.name") - - assert isinstance(result, exp.Column) - assert result.name == "parsed" - assert calls == ["users.name"] + assert isinstance(condition, exp.Like) + assert condition.this == exp.column("name", table="users") + assert condition.expression == exp.Placeholder(this="users_name_search") + assert result.named_parameters == {"users_name_search": "%needle%"} def test_canonicalize_filters_deduplicates_and_sorts_by_class_and_field() -> None: @@ -945,12 +928,6 @@ def test_order_by_filter_rejects_invalid_sort_order_at_construction() -> None: OrderByFilter("name", "BAD") # type: ignore[arg-type] -def test_not_in_search_filter_declares_empty_slots() -> None: - """NotInSearchFilter should not allocate an instance dict.""" - - assert NotInSearchFilter.__slots__ == () - - def test_not_in_collection_parameter_suffix_is_literal() -> None: """The not-in suffix should be the direct literal used in parameter names.""" diff --git a/tests/unit/core/test_type_conversion.py b/tests/unit/core/test_type_conversion.py index 0a5f69d36..ce7e17bf9 100644 --- a/tests/unit/core/test_type_conversion.py +++ b/tests/unit/core/test_type_conversion.py @@ -1,6 +1,6 @@ """Tests for centralized scalar conversion helpers.""" -from datetime import UTC, date, datetime, time +from datetime import UTC, date, datetime, time, timedelta, timezone from decimal import Decimal from uuid import UUID @@ -114,20 +114,24 @@ def test_parse_datetime_rfc3339_z_suffix() -> None: assert parsed.year == 2023 -def test_parse_datetime_rfc3339_delegates_to_iso_converter(monkeypatch: pytest.MonkeyPatch) -> None: - """RFC3339 parsing should reuse the ISO datetime converter.""" - - import sqlspec.core.type_converter as type_converter - - expected = datetime(2023, 12, 25, 10, 30, 0, tzinfo=UTC) - - def fake_convert(value: str) -> datetime: - assert value == "2023-12-25T10:30:00+00:00" - return expected - - monkeypatch.setattr(type_converter, "convert_iso_datetime", fake_convert) +@pytest.mark.parametrize( + ("value", "expected"), + [ + ("2023-12-25T10:30:00Z", datetime(2023, 12, 25, 10, 30, tzinfo=UTC)), + ( + "2023-12-25T10:30:00.123456+05:30", + datetime(2023, 12, 25, 10, 30, 0, 123456, tzinfo=timezone(timedelta(hours=5, minutes=30))), + ), + ("2023-12-25T10:30:00-04:00", datetime(2023, 12, 25, 10, 30, tzinfo=timezone(timedelta(hours=-4)))), + ], +) +def test_parse_datetime_rfc3339_preserves_timezone_and_precision(value: str, expected: datetime) -> None: + """RFC3339 parsing preserves the represented instant, offset, and microseconds.""" + parsed = parse_datetime_rfc3339(value) - assert type_converter.parse_datetime_rfc3339("2023-12-25T10:30:00Z") is expected + assert parsed == expected + assert parsed.utcoffset() == expected.utcoffset() + assert parsed.microsecond == expected.microsecond def test_datetime_round_trip() -> None: diff --git a/tests/unit/data_dictionary/test_mssql_metadata.py b/tests/unit/data_dictionary/test_mssql_metadata.py index 7f6679551..470422978 100644 --- a/tests/unit/data_dictionary/test_mssql_metadata.py +++ b/tests/unit/data_dictionary/test_mssql_metadata.py @@ -1,11 +1,15 @@ """SQL Server data-dictionary metadata tests.""" -from typing import Any, cast +import sqlite3 +from collections.abc import Iterator +from typing import TYPE_CHECKING, Any, TypeVar, cast, overload import pytest from sqlspec.adapters.mssql_python.data_dictionary import MssqlPythonSyncDataDictionary +from sqlspec.adapters.mssql_python.driver import MssqlPythonDriver from sqlspec.adapters.pymssql.data_dictionary import PymssqlSyncDataDictionary +from sqlspec.adapters.pymssql.driver import PymssqlDriver from sqlspec.data_dictionary import ( DataDictionaryLoader, DDLResult, @@ -16,6 +20,12 @@ SystemMetadataRequest, SystemMetadataResult, ) +from sqlspec.driver import SyncDataDictionaryBase + +if TYPE_CHECKING: + from collections.abc import Callable + +_CursorT = TypeVar("_CursorT", bound=sqlite3.Cursor) REQUIRED_MSSQL_DOMAIN_QUERIES = { "schemas": ("list",), @@ -61,6 +71,88 @@ def select(self, statement: Any, **kwargs: Any) -> list[dict[str, Any]]: return [] +class CatalogCursor(sqlite3.Cursor): + """Execute SDK-bound SQL against a local substitute SQL Server catalog.""" + + def execute(self, sql: str, parameters: Any = ()) -> "CatalogCursor": + super().execute(sql.replace("%s", "?"), parameters) + return self + + +class CatalogConnection(sqlite3.Connection): + """Provide the DB-API cursor surface shared by the SQL Server SDKs.""" + + @overload + def cursor(self, factory: None = None) -> CatalogCursor: ... + + @overload + def cursor(self, factory: "Callable[[sqlite3.Connection], _CursorT]") -> _CursorT: ... + + def cursor(self, factory: "Callable[[sqlite3.Connection], sqlite3.Cursor] | None" = None) -> sqlite3.Cursor: + return super().cursor(factory=CatalogCursor if factory is None else factory) + + +@pytest.fixture(params=[MssqlPythonDriver, PymssqlDriver], ids=["mssql_python", "pymssql"]) +def catalog_driver(request: pytest.FixtureRequest) -> Iterator[MssqlPythonDriver | PymssqlDriver]: + """Run real drivers and dictionary SQL with only the external SDK substituted.""" + connection = sqlite3.connect(":memory:", factory=CatalogConnection) + connection.create_function("DB_NAME", 0, lambda: "app") + connection.create_function("SCHEMA_NAME", 0, lambda: "dbo") + connection.executescript(""" + ATTACH DATABASE ':memory:' AS sys; + CREATE TABLE sys.schemas (schema_id INTEGER, name TEXT); + CREATE TABLE sys.tables ( + object_id INTEGER, schema_id INTEGER, name TEXT, create_date TEXT, modify_date TEXT, + is_memory_optimized INTEGER, durability INTEGER, durability_desc TEXT, + temporal_type INTEGER, temporal_type_desc TEXT, history_table_id INTEGER + ); + INSERT INTO sys.schemas VALUES (1, 'dbo'); + INSERT INTO sys.tables VALUES + (10, 1, 'accounts', '2026-01-01', '2026-02-01', 0, 0, 'SCHEMA_AND_DATA', + 2, 'SYSTEM_VERSIONED_TEMPORAL_TABLE', 11), + (11, 1, 'accounts_history', '2026-01-01', '2026-01-01', 0, 0, 'SCHEMA_AND_DATA', + 1, 'HISTORY_TABLE', NULL); + """) + try: + yield request.param(cast(Any, connection)) + finally: + connection.close() + + +@pytest.mark.parametrize("schema", [None, "dbo"]) +def test_table_details_selects_requested_table_with_rich_metadata( + catalog_driver: MssqlPythonDriver | PymssqlDriver, schema: str | None +) -> None: + """A detail lookup preserves table identity, rich fields, and missing-table behavior.""" + dictionary = cast("SyncDataDictionaryBase", catalog_driver.data_dictionary) + result = dictionary.get_table_details(catalog_driver, "accounts", schema) + + assert result.items == ( + { + "catalog_name": "app", + "schema_name": "dbo", + "table_name": "accounts", + "table_type": "BASE TABLE", + "create_date": "2026-01-01", + "modify_date": "2026-02-01", + "is_memory_optimized": 0, + "durability": 0, + "durability_desc": "SCHEMA_AND_DATA", + "temporal_type": 2, + "temporal_type_desc": "SYSTEM_VERSIONED_TEMPORAL_TABLE", + "history_schema_name": "dbo", + "history_table_name": "accounts_history", + }, + ) + history = dictionary.get_table_details(catalog_driver, "accounts_history", schema) + assert len(history.items) == 1 + history_row = cast("dict[str, object]", history.items[0]) + assert history_row["table_name"] == "accounts_history" + assert history_row["history_table_name"] is None + missing = dictionary.get_table_details(catalog_driver, "missing", schema) + assert missing.items == () + + def test_mssql_domain_query_pack_supports_required_catalog_queries() -> None: """SQL Server query packs should load from direct domain paths.""" loader = DataDictionaryLoader() diff --git a/tests/unit/data_dictionary/test_mysql_metadata.py b/tests/unit/data_dictionary/test_mysql_metadata.py index 0ec27bf81..6c4bb6405 100644 --- a/tests/unit/data_dictionary/test_mysql_metadata.py +++ b/tests/unit/data_dictionary/test_mysql_metadata.py @@ -1,7 +1,9 @@ """Unit tests for MySQL-family data-dictionary metadata.""" from importlib import import_module +from inspect import isawaitable from typing import Any, cast +from unittest.mock import AsyncMock, Mock import pytest import sqlglot @@ -11,6 +13,7 @@ from sqlspec.adapters.asyncmy.data_dictionary import AsyncmyDataDictionary from sqlspec.adapters.mysqlconnector.data_dictionary import MysqlConnectorSyncDataDictionary from sqlspec.adapters.pymysql.data_dictionary import PyMysqlDataDictionary +from sqlspec.core import SQL from sqlspec.data_dictionary import ( DataDictionaryLoader, DDLResult, @@ -28,9 +31,18 @@ class _SyncDDLDriver: def __init__(self) -> None: self.statements: list[str] = [] + self.sql_mode = "STRICT_TRANS_TABLES,NO_ENGINE_SUBSTITUTION" + self.sql_quote_show_create = "1" def select_one(self, statement: str, *args: Any, **kwargs: Any) -> dict[str, str]: self.statements.append(str(statement)) + query = str(statement) + if "SELECT VERSION()" in query: + return { + "raw_version": "8.4.0", + "sql_mode": self.sql_mode, + "sql_quote_show_create": self.sql_quote_show_create, + } return {"Table": "orders", "Create Table": "CREATE TABLE `orders` (`id` int)"} def select_value_or_none(self, statement: object, *args: Any, **kwargs: Any) -> str | None: @@ -259,3 +271,47 @@ def select(self, statement: str, *args: Any, **kwargs: Any) -> list[dict[str, ob synthesized = dictionary.get_object_ddl(cast(Any, synth_driver), "orders", schema="shop", prefer_native=False) assert synthesized.status == MetadataSupport.UNSUPPORTED assert not any(statement.startswith("SHOW CREATE") for statement in synth_driver.statements) + + +@pytest.mark.parametrize("dictionary_type", [PyMysqlDataDictionary, AiomysqlDataDictionary]) +async def test_mysql_native_ddl_refreshes_session_context(dictionary_type: type[Any]) -> None: + """DDL captures mutable SQL mode and quote settings for every statement.""" + session = _SyncDDLDriver() + driver = ( + Mock(select_one=AsyncMock(side_effect=session.select_one)) + if dictionary_type is AiomysqlDataDictionary + else session + ) + dictionary = dictionary_type() + for sql_mode, quotes in [(session.sql_mode, "1"), ("NO_BACKSLASH_ESCAPES", "0")]: + session.sql_mode = sql_mode + session.sql_quote_show_create = quotes + result = dictionary.get_ddl(driver, "orders", schema="shop") + if isawaitable(result): + result = await result + assert result.context["sql_mode"] == sql_mode + assert result.context["sql_quote_show_create"] == quotes + + +@pytest.mark.parametrize("dictionary_type", [PyMysqlDataDictionary, AiomysqlDataDictionary]) +@pytest.mark.parametrize("object_name", [None, "orders"]) +async def test_mysql_get_privileges_binds_optional_object_name( + dictionary_type: type[Any], object_name: str | None +) -> None: + from sqlspec.adapters.pymysql.core import default_statement_config + + def select_rows(statement: str, **parameters: Any) -> list[dict[str, str]]: + compiled_sql, bound = SQL(str(statement), statement_config=default_statement_config, **parameters).compile() + assert "information_schema.table_privileges" in compiled_sql + assert bound == {"schema_name": "shop", "object_name": object_name} + return [{"table_name": "orders"}] + + is_async = dictionary_type is AiomysqlDataDictionary + driver = Mock( + select=AsyncMock(side_effect=select_rows) if is_async else Mock(side_effect=select_rows), + select_value_or_none=AsyncMock(return_value="8.4.0") if is_async else Mock(return_value="8.4.0"), + ) + result = dictionary_type().get_privileges(driver, object_name=object_name, schema="shop") + if isawaitable(result): + result = await result + assert result.items == ({"table_name": "orders"},) diff --git a/tests/unit/data_dictionary/test_postgres_cockroach_query_packs.py b/tests/unit/data_dictionary/test_postgres_cockroach_query_packs.py index c53aa7311..290d4b534 100644 --- a/tests/unit/data_dictionary/test_postgres_cockroach_query_packs.py +++ b/tests/unit/data_dictionary/test_postgres_cockroach_query_packs.py @@ -150,3 +150,33 @@ def test_postgres_and_cockroach_foreign_keys_and_indexes_contracts() -> None: assert "AS is_primary" in crdb_columns_by_table assert "AS generation_expression" in crdb_columns_by_table assert "COALESCE(" in crdb_columns_by_table + + +def test_postgres_system_and_constraints_query_contracts() -> None: + """PostgreSQL pg_stat_statements omits pg_catalog prefix and constraints pre-aggregate columns.""" + loader = DataDictionaryLoader() + sys_query = loader.get_domain_query_text("postgres", "system", "pg_stat_statements") + assert sys_query is not None + assert "FROM pg_catalog.pg_stat_statements" not in sys_query + assert "FROM pg_stat_statements" in sys_query + + con_query = loader.get_domain_query_text("postgres", "constraints", "by_schema") + assert con_query is not None + assert "WITH con_cols AS (" in con_query + assert "CROSS JOIN LATERAL unnest(con.conkey) WITH ORDINALITY" in con_query + assert "ARRAY(SELECT att.attname" not in con_query + + +def test_cockroach_constraints_and_columns_query_contracts() -> None: + """CockroachDB constraints and columns queries pre-aggregate and factor scalar subqueries.""" + loader = DataDictionaryLoader() + con_query = loader.get_domain_query_text("cockroachdb", "constraints", "by_schema") + assert con_query is not None + assert "WITH con_cols AS (" in con_query + assert "array_agg(kcu.column_name::text ORDER BY kcu.ordinal_position)" in con_query + assert "ARRAY(SELECT kcu.column_name" not in con_query + + cols_query = loader.get_domain_query_text("cockroachdb", "columns", "by_schema") + assert cols_query is not None + assert "WITH resolved_table_schema AS (" in cols_query + assert "FROM resolved_table_schema" in cols_query diff --git a/tests/unit/data_dictionary/test_sqlite_duckdb_metadata.py b/tests/unit/data_dictionary/test_sqlite_duckdb_metadata.py index 933f10c19..39a7af6b7 100644 --- a/tests/unit/data_dictionary/test_sqlite_duckdb_metadata.py +++ b/tests/unit/data_dictionary/test_sqlite_duckdb_metadata.py @@ -12,6 +12,7 @@ from sqlspec.adapters.sqlite import SqliteConfig, SqliteDriver from sqlspec.adapters.sqlite.data_dictionary import SqliteDataDictionary from sqlspec.data_dictionary import DataDictionaryLoader, DependencyMetadata, MetadataSupport +from sqlspec.exceptions import SQLFileNotFoundError from tests.conftest import requires_interpreted pytestmark = requires_interpreted @@ -396,3 +397,168 @@ def test_sqlite_constraints_include_rowid_alias_primary_key() -> None: types = {item["constraint_type"]: item["columns"] for item in constraint_items} assert types["PRIMARY KEY"] == ["id"] assert types["UNIQUE"] == ["code"] + + +def test_sqlite_get_columns_and_indexes_with_none_schema() -> None: + """Introspection with schema=None should coalesce to main.""" + connection = sqlite3.connect(":memory:") + connection.row_factory = sqlite3.Row + driver = SqliteDriver(connection) + try: + driver.execute_script(""" + CREATE TABLE users ( + id INTEGER PRIMARY KEY, + name TEXT NOT NULL, + email TEXT NOT NULL + ); + CREATE INDEX idx_users_name ON users(name); + """) + + cols = driver.data_dictionary.get_columns(driver, table="users", schema=None) + assert len(cols) == 3 + assert all(c["schema_name"] == "main" for c in cols) + + indexes = driver.data_dictionary.get_indexes(driver, table="users", schema=None) + assert len(indexes) == 1 + assert indexes[0]["schema_name"] == "main" + assert indexes[0]["index_name"] == "idx_users_name" + finally: + connection.close() + + +def test_sqlite_data_dictionary_missing_query_raises_sql_file_not_found() -> None: + """Missing query in SQLite data dictionary raises SQLFileNotFoundError.""" + dd = SqliteDataDictionary() + with pytest.raises(SQLFileNotFoundError, match="Missing SQLite data-dictionary query"): + dd._get_domain_query_text("nonexistent_domain", "nonexistent_query") + + +def test_sqlite_generated_column_memoization(monkeypatch: pytest.MonkeyPatch) -> None: + """Verify generated column expressions are parsed once per table.""" + from sqlspec.adapters.sqlite import data_dictionary as dd_module + + connection = sqlite3.connect(":memory:") + connection.row_factory = sqlite3.Row + driver = SqliteDriver(connection) + try: + driver.execute_script(""" + CREATE TABLE gen_test ( + id INTEGER PRIMARY KEY, + a INT, + b INT, + c INT GENERATED ALWAYS AS (a + b) VIRTUAL, + d INT GENERATED ALWAYS AS (a * b) STORED + ); + """) + + call_count = 0 + original_fn = dd_module.sqlite_generated_expressions + + def counting_parser(sql: str | None) -> dict[str, str]: + nonlocal call_count + call_count += 1 + return original_fn(sql) + + monkeypatch.setattr(dd_module, "sqlite_generated_expressions", counting_parser) + + cols = driver.data_dictionary.get_columns(driver, table="gen_test") + assert call_count == 1 + gen_cols = {c["column_name"]: c["generation_expression"] for c in cols if c.get("is_generated")} + assert gen_cols == {"c": "a + b", "d": "a * b"} + finally: + connection.close() + + +def test_duckdb_metadata_capabilities_profile() -> None: + """DuckDBDataDictionary reports supported capabilities across metadata domains.""" + config = DuckDBConfig(connection_config={"database": ":memory:"}) + driver = DuckDBDriver(config.create_connection()) + try: + profile = driver.data_dictionary.get_metadata_capabilities(driver) + assert profile.dialect == "duckdb" + assert profile.get("tables").support == MetadataSupport.SUPPORTED + assert profile.get("columns").support == MetadataSupport.SUPPORTED + assert profile.get("schemas").support == MetadataSupport.SUPPORTED + assert profile.get("objects").support == MetadataSupport.SUPPORTED + assert profile.get("routines").support == MetadataSupport.SUPPORTED + assert profile.get("system").support == MetadataSupport.SUPPORTED + assert profile.get("databases").support == MetadataSupport.SUPPORTED + assert profile.get("extensions").support == MetadataSupport.SUPPORTED + assert profile.get("privileges").support == MetadataSupport.UNSUPPORTED + finally: + driver.connection.close() + + +@pytest.mark.parametrize(("table", "expected_names"), [("target", ["target"]), ("missing", [])]) +def test_duckdb_table_details_filters_requested_table(table: str, expected_names: list[str]) -> None: + config = DuckDBConfig(connection_config={"database": ":memory:"}) + driver = DuckDBDriver(config.create_connection()) + try: + driver.execute_script(""" + CREATE SCHEMA app; + CREATE TABLE app.target (id INTEGER); + CREATE TABLE app.other (id INTEGER); + CREATE TABLE target (id INTEGER); + """) + + details = driver.data_dictionary.get_table_details(driver, table=table, schema="app") + + assert {row["table_name"] for row in driver.data_dictionary.get_tables(driver, schema="app")} == { + "target", + "other", + } + assert details.capability.support == MetadataSupport.SUPPORTED + assert [cast("dict[str, object]", item)["table_name"] for item in details.items] == expected_names + finally: + driver.connection.close() + + +def test_duckdb_unwired_sql_domains() -> None: + """DuckDBDataDictionary wires schemas, objects, routines, databases, extensions, system, and table_details.""" + config = DuckDBConfig(connection_config={"database": ":memory:"}) + driver = DuckDBDriver(config.create_connection()) + try: + driver.execute("CREATE SCHEMA custom_schema") + driver.execute("CREATE TABLE custom_schema.test_unwired (id INT PRIMARY KEY, name VARCHAR)") + + schemas = driver.data_dictionary.get_schemas(driver) + assert schemas.domain == "schemas" + assert len(schemas.items) > 0 + + objects = driver.data_dictionary.get_objects(driver, schema="custom_schema") + assert objects.domain == "objects" + assert any(cast("dict[str, object]", obj).get("object_name") == "test_unwired" for obj in objects.items) + + routines = driver.data_dictionary.get_routines(driver, schema="custom_schema") + assert routines.domain == "routines" + + databases = driver.data_dictionary.get_databases(driver) + assert databases.domain == "databases" + assert len(databases.items) > 0 + + extensions = driver.data_dictionary.get_extensions(driver) + assert extensions.domain == "extensions" + assert len(extensions.items) > 0 + + table_details = driver.data_dictionary.get_table_details(driver, table="test_unwired", schema="custom_schema") + assert table_details.domain == "tables" + assert len(table_details.items) == 1 + assert cast("dict[str, object]", table_details.items[0])["table_name"] == "test_unwired" + + system_settings = driver.data_dictionary.get_system_metadata(driver, "settings", include_system=True) + assert system_settings.request.domain == "settings" + assert len(system_settings.rows) > 0 + finally: + driver.connection.close() + + +def test_references_identifier() -> None: + """DuckDBDataDictionary._references_identifier matches whole identifiers inherited from base.""" + config = DuckDBConfig(connection_config={"database": ":memory:"}) + driver = DuckDBDriver(config.create_connection()) + try: + dd = driver.data_dictionary + assert dd._references_identifier("SELECT * FROM my_table", "my_table") is True + assert dd._references_identifier("SELECT * FROM my_table_extra", "my_table") is False + finally: + driver.connection.close() diff --git a/tests/unit/data_dictionary/test_sybase_sqlanywhere_dialects.py b/tests/unit/data_dictionary/test_sybase_sqlanywhere_dialects.py index a2a72c5dd..9a22cd808 100644 --- a/tests/unit/data_dictionary/test_sybase_sqlanywhere_dialects.py +++ b/tests/unit/data_dictionary/test_sybase_sqlanywhere_dialects.py @@ -50,7 +50,7 @@ def test_sybase_and_sqlanywhere_dialect_configs_registered() -> None: assert SQLANYWHERE_CONFIG.get_feature_flag("supports_merge") is True assert SQLANYWHERE_CONFIG.get_feature_flag("supports_for_update") is True assert SQLANYWHERE_CONFIG.get_feature_flag("supports_skip_locked") is True - assert SQLANYWHERE_CONFIG.get_feature_flag("supports_transactional_ddl") is True + assert SQLANYWHERE_CONFIG.get_feature_flag("supports_transactional_ddl") is False assert SQLANYWHERE_CONFIG.get_optimal_type("uuid") == "UNIQUEIDENTIFIER" assert SQLANYWHERE_CONFIG.get_optimal_type("text") == "LONG VARCHAR" assert SQLANYWHERE_CONFIG.get_optimal_type("blob") == "LONG BINARY" diff --git a/tests/unit/dialects/test_generators.py b/tests/unit/dialects/test_generators.py index 21ed2ad28..54a3b4204 100644 --- a/tests/unit/dialects/test_generators.py +++ b/tests/unit/dialects/test_generators.py @@ -4,8 +4,6 @@ produces correct SQL output for all extension operators and Spanner DDL. """ -from pathlib import Path - from sqlglot import parse_one from sqlspec.dialects.spanner._generators import _get_dialect_name @@ -95,12 +93,3 @@ def test_spangres_normalizes_row_deletion_policy_to_ttl() -> None: sql = "\n CREATE TABLE events (\n id VARCHAR PRIMARY KEY,\n created_at TIMESTAMP\n ) ROW DELETION POLICY (OLDER_THAN(created_at, INTERVAL 30 DAY))\n " rendered = parse_one(sql, dialect="spangres").sql(dialect="spangres") assert "TTL INTERVAL '30 days' ON created_at" in rendered - - -def test_compat_removed_dialects_compat_module_removed() -> None: - assert not Path("sqlspec/dialects/_compat.py").exists() - - -def test_compat_removed_dialects_compat_not_in_mypyc_include() -> None: - assert "sqlspec/dialects/_compat.py" not in Path("pyproject.toml").read_text() - assert "sqlspec/dialects/_compat.py" not in Path("Makefile").read_text() diff --git a/tests/unit/dialects/test_registration.py b/tests/unit/dialects/test_registration.py index dc97e86ed..d037f594b 100644 --- a/tests/unit/dialects/test_registration.py +++ b/tests/unit/dialects/test_registration.py @@ -4,6 +4,10 @@ import sys from importlib.metadata import entry_points +import pytest + +from tests.conftest import is_compiled + DIALECT_ENTRY_POINTS = { "db2": "sqlspec.dialects.db2", "paradedb": "sqlspec.dialects.postgres", @@ -36,8 +40,9 @@ def test_sqlglot_resolves_dialects_without_sqlspec_dialect_import() -> None: assert "<=>" in result.stdout +@pytest.mark.skipif(is_compiled(), reason="Compiled root facade eagerly binds dialect exports for mypyc compatibility") def test_importing_sqlspec_does_not_eagerly_load_dialect_machinery() -> None: - """``import sqlspec`` must not pay the sqlglot dialect registration cost.""" + """The interpreted root facade defers dialect registration until first use.""" code = ( "import sys\n" "import sqlspec\n" diff --git a/tests/unit/dialects/test_spanner.py b/tests/unit/dialects/test_spanner.py index 38047436c..f713aa1cc 100644 --- a/tests/unit/dialects/test_spanner.py +++ b/tests/unit/dialects/test_spanner.py @@ -5,8 +5,6 @@ clauses, and DAY-based row deletion intervals. """ -from pathlib import Path - import pytest from sqlglot import Dialect, exp, parse_one @@ -112,8 +110,6 @@ def test_interleave_create_repairs_command_fallback() -> None: expression = parse_one(OFFICIAL_INTERLEAVE_DDL, dialect="spanner") assert isinstance(expression, exp.Create) assert "INTERLEAVE IN PARENT Singers ON DELETE CASCADE" in expression.sql(dialect="spanner") - assert not hasattr(_parsers, "_INTERLEAVE_PATTERN") - assert not hasattr(_parsers, "extract_interleave_property") ordered_pk_ddl = ( "CREATE TABLE Albums (SingerId INT64 NOT NULL, AlbumId INT64 NOT NULL) " @@ -200,19 +196,6 @@ def test_spanner_property_moves_after_non_spanner_property() -> None: assert rendered.index("enable_change_stream_capture") < rendered.index("INTERLEAVE IN PARENT") -def test_dedup_normalize_interval_expression_has_single_definition() -> None: - package_root = Path("sqlspec/dialects/spanner") - matches = [path for path in package_root.glob("_*.py") if "def _normalize_interval_expression" in path.read_text()] - assert matches == [package_root / "_generators.py"] - - -def test_dedup_spanner_constants_have_single_definition() -> None: - package_root = Path("sqlspec/dialects/spanner") - for name in ("_TTL_MIN_COMPONENTS", "_ROW_DELETION_NAME", "_INTERLEAVE_NAME", "_INTERLEAVE_IN_NAME"): - matches = [path for path in package_root.glob("_*.py") if f"{name} =" in path.read_text()] - assert matches == [package_root / "_generators.py"] - - def test_dedup_parsers_import_canonical_helpers() -> None: assert _parsers._normalize_interval_expression is _generators._normalize_interval_expression assert _parsers._ROW_DELETION_NAME is _generators._ROW_DELETION_NAME diff --git a/tests/unit/driver/test_query_cache.py b/tests/unit/driver/test_query_cache.py index 5721a8dee..b8a72bb29 100644 --- a/tests/unit/driver/test_query_cache.py +++ b/tests/unit/driver/test_query_cache.py @@ -698,7 +698,7 @@ async def fetchall_async() -> Any: operation_type="SELECT" if returns_rows else "INSERT", operation_profile=OperationProfile(returns_rows=returns_rows, modifies_rows=not returns_rows), ) - result = await driver._execute_cache_hit_with_cursor(sql, ("once",), cached) + result = await AsyncDriverAdapterBase._execute_cache_hit_with_cursor(driver, sql, ("once",), cached) executed = [statement for statement in statements if statement.startswith(("SELECT", "INSERT"))] assert len(executed) == 1 @@ -747,7 +747,7 @@ async def execute_async(sql: str, params: Any) -> None: operation_profile=OperationProfile(returns_rows=returns_rows, modifies_rows=not returns_rows), ) with pytest.raises(error_type) as exc_info: - await driver._execute_cache_hit_with_cursor(sql, ("once",), cached) + await AsyncDriverAdapterBase._execute_cache_hit_with_cursor(driver, sql, ("once",), cached) assert exc_info.value is error assert len([statement for statement in statements if statement.startswith(("SELECT", "INSERT"))]) == 1 @@ -769,7 +769,9 @@ def fetch_rows() -> list[Any]: cursor.fetchall = Mock(side_effect=fetch_rows) sql = "SELECT id, name FROM users ORDER BY id" - result = await driver._execute_cache_hit_with_cursor(sql, (), _make_cached(compiled_sql=sql)) + result = await AsyncDriverAdapterBase._execute_cache_hit_with_cursor( + driver, sql, (), _make_cached(compiled_sql=sql) + ) assert result.get_data() == [{"id": 1, "name": "test"}, {"id": 2, "name": "example"}] assert result.column_names == ["id", "name"] @@ -792,7 +794,7 @@ async def test_async_cache_hit_missing_cursor_capability_dispatches_once( operation_type="SELECT" if returns_rows else "INSERT", operation_profile=OperationProfile(returns_rows=returns_rows, modifies_rows=not returns_rows), ) - result = await driver._execute_cache_hit_with_cursor(sql, ("once",), cached) + result = await AsyncDriverAdapterBase._execute_cache_hit_with_cursor(driver, sql, ("once",), cached) assert len([statement for statement in statements if statement.startswith(("SELECT", "INSERT"))]) == 1 fallback.assert_awaited_once() diff --git a/tests/unit/driver/test_storage_parallel_and_auth.py b/tests/unit/driver/test_storage_parallel_and_auth.py index 4371e28ea..c5b7e80fd 100644 --- a/tests/unit/driver/test_storage_parallel_and_auth.py +++ b/tests/unit/driver/test_storage_parallel_and_auth.py @@ -238,10 +238,15 @@ def fake_psqlpy_pool(**pool_kwargs: Any) -> MagicMock: with patch("sqlspec.adapters.psqlpy.config.ConnectionPool", side_effect=fake_psqlpy_pool): await psqlpy_cfg._create_pool() assert captured_psqlpy["password"] == "async-iam-token" - assert ( - build_psqlpy_connection_config({"username": "u", "auth_token_provider": lambda: "sync-tok"})["password"] - == "sync-tok" - ) + sync_calls: list[str] = [] + + def sync_provider() -> str: + sync_calls.append("called") + return "sync-tok" + + deferred = build_psqlpy_connection_config({"username": "u", "auth_token_provider": sync_provider}) + assert "password" not in deferred + assert sync_calls == [] def test_adbc_auth_token_provider_resolves_fresh_credentials() -> None: diff --git a/tests/unit/exceptions/test_exception_handler.py b/tests/unit/exceptions/test_exception_handler.py index 66d6d1294..ea92b8347 100644 --- a/tests/unit/exceptions/test_exception_handler.py +++ b/tests/unit/exceptions/test_exception_handler.py @@ -150,39 +150,6 @@ def fake_create_mapped_exception(exc_val: BaseException) -> Exception: assert seen == {"exc_val": error} -def test_mysqlconnector_sync_exception_handler_preserves_suppression(monkeypatch: pytest.MonkeyPatch) -> None: - """mysql-connector sync handler should preserve migration-suppression sentinel values.""" - pytest.importorskip("mysql.connector") - import mysql.connector - - from sqlspec.adapters.mysqlconnector import driver as mysqlconnector_driver - - monkeypatch.setattr(mysqlconnector_driver, "create_mapped_exception", lambda *args, **kwargs: True) - - error = mysql.connector.Error("skip mapping") - handler = mysqlconnector_driver.MysqlConnectorSyncExceptionHandler() - - assert handler.__exit__(type(error), error, None) is True - assert handler.pending_exception is None - - -@pytest.mark.anyio -async def test_mysqlconnector_async_exception_handler_preserves_suppression(monkeypatch: pytest.MonkeyPatch) -> None: - """mysql-connector async handler should preserve migration-suppression sentinel values.""" - pytest.importorskip("mysql.connector") - import mysql.connector - - from sqlspec.adapters.mysqlconnector import driver as mysqlconnector_driver - - monkeypatch.setattr(mysqlconnector_driver, "create_mapped_exception", lambda *args, **kwargs: True) - - error = mysql.connector.Error("skip mapping") - handler = mysqlconnector_driver.MysqlConnectorAsyncExceptionHandler() - - assert await handler.__aexit__(type(error), error, None) is True - assert handler.pending_exception is None - - @pytest.mark.anyio async def test_cockroach_asyncpg_exception_handler_preserves_serialization_conflicts( monkeypatch: pytest.MonkeyPatch, diff --git a/tests/unit/extensions/test_adk/test_store_instantiation.py b/tests/unit/extensions/test_adk/test_store_instantiation.py index 1e6787868..d7891588e 100644 --- a/tests/unit/extensions/test_adk/test_store_instantiation.py +++ b/tests/unit/extensions/test_adk/test_store_instantiation.py @@ -15,6 +15,7 @@ import pytest +from sqlspec.core import StatementConfig from sqlspec.exceptions import SQLSpecError from sqlspec.extensions.adk.memory.store import BaseAsyncADKMemoryStore, BaseSyncADKMemoryStore from sqlspec.extensions.adk.store import BaseAsyncADKStore, BaseSyncADKStore @@ -29,7 +30,7 @@ "sqlspec.adapters.cockroach_asyncpg.adk.CockroachAsyncpgADKStore", "sqlspec.adapters.cockroach_psycopg.adk.CockroachPsycopgAsyncADKStore", "sqlspec.adapters.cockroach_psycopg.adk.CockroachPsycopgSyncADKStore", - "sqlspec.adapters.duckdb.adk.DuckdbADKStore", + "sqlspec.adapters.duckdb.adk.DuckDBADKStore", "sqlspec.adapters.mssql_python.adk.MssqlPythonADKStore", "sqlspec.adapters.mysqlconnector.adk.MysqlConnectorAsyncADKStore", "sqlspec.adapters.mysqlconnector.adk.MysqlConnectorSyncADKStore", @@ -55,7 +56,7 @@ "sqlspec.adapters.cockroach_asyncpg.adk.CockroachAsyncpgADKMemoryStore", "sqlspec.adapters.cockroach_psycopg.adk.CockroachPsycopgAsyncADKMemoryStore", "sqlspec.adapters.cockroach_psycopg.adk.CockroachPsycopgSyncADKMemoryStore", - "sqlspec.adapters.duckdb.adk.DuckdbADKMemoryStore", + "sqlspec.adapters.duckdb.adk.DuckDBADKMemoryStore", "sqlspec.adapters.mssql_python.adk.MssqlPythonADKMemoryStore", "sqlspec.adapters.mysqlconnector.adk.MysqlConnectorAsyncADKMemoryStore", "sqlspec.adapters.mysqlconnector.adk.MysqlConnectorSyncADKMemoryStore", @@ -77,7 +78,7 @@ "sqlspec.adapters.adbc.adk.AdbcADKStore", "sqlspec.adapters.arrow_odbc.adk.ArrowOdbcADKStore", "sqlspec.adapters.cockroach_psycopg.adk.CockroachPsycopgSyncADKStore", - "sqlspec.adapters.duckdb.adk.DuckdbADKStore", + "sqlspec.adapters.duckdb.adk.DuckDBADKStore", "sqlspec.adapters.mssql_python.adk.MssqlPythonADKStore", "sqlspec.adapters.mysqlconnector.adk.MysqlConnectorSyncADKStore", "sqlspec.adapters.oracledb.adk.OracleSyncADKStore", @@ -92,7 +93,7 @@ "sqlspec.adapters.adbc.adk.AdbcADKMemoryStore", "sqlspec.adapters.arrow_odbc.adk.ArrowOdbcADKMemoryStore", "sqlspec.adapters.cockroach_psycopg.adk.CockroachPsycopgSyncADKMemoryStore", - "sqlspec.adapters.duckdb.adk.DuckdbADKMemoryStore", + "sqlspec.adapters.duckdb.adk.DuckDBADKMemoryStore", "sqlspec.adapters.mssql_python.adk.MssqlPythonADKMemoryStore", "sqlspec.adapters.mysqlconnector.adk.MysqlConnectorSyncADKMemoryStore", "sqlspec.adapters.oracledb.adk.OracleSyncADKMemoryStore", @@ -192,6 +193,56 @@ def test_adk_store_registration_validator_resolves_duckdb_store_classes() -> Non _ensure_adk_store_registration(config) +@pytest.mark.parametrize( + ("adapter", "prefix", "dialect"), + [ + ("db2", "Db2Sync", "db2"), + ("db2", "Db2Async", "db2"), + ("arrow_odbc", "ArrowOdbc", "tsql"), + ("arrow_odbc", "ArrowOdbc", "db2"), + ("cockroach_asyncpg", "CockroachAsyncpg", "postgres"), + ("cockroach_psycopg", "CockroachPsycopgSync", "postgres"), + ("cockroach_psycopg", "CockroachPsycopgAsync", "postgres"), + ("duckdb", "DuckDB", "duckdb"), + ("mssql_python", "MssqlPython", "tsql"), + ("pymssql", "Pymssql", "tsql"), + ("sybase", "Sybase", "sybase"), + ("sqlanywhere", "Sqlanywhere", "sqlanywhere"), + ], +) +def test_qualified_table_names_reach_supporting_adk_stores(adapter: str, prefix: str, dialect: str) -> None: + """Supported stores accept qualified tables through their public configuration.""" + session_tables = { + "session_table": "app.sessions", + "events_table": "app.events", + "app_state_table": "app.app_state", + "user_state_table": "app.user_state", + "metadata_table": "app.metadata", + } + config = _load_class(f"sqlspec.adapters.{adapter}.{prefix}Config")( + extension_config={"adk": {**session_tables, "memory_table": "app.memories"}}, + statement_config=StatementConfig(dialect=dialect), + ) + + session_store = _load_class(f"sqlspec.adapters.{adapter}.adk.{prefix}ADKStore")(config) + memory_store = _load_class(f"sqlspec.adapters.{adapter}.adk.{prefix}ADKMemoryStore")(config) + + assert {name: getattr(session_store, name) for name in session_tables} == session_tables + assert memory_store.memory_table == "app.memories" + + +@pytest.mark.parametrize("table", ["dataset.sessions", "project.dataset.sessions"]) +def test_bigquery_adk_preserves_qualified_table_configuration(table: str) -> None: + """BigQuery keeps dataset and project qualification supplied by the caller.""" + from sqlspec.adapters.bigquery import BigQueryConfig + from sqlspec.adapters.bigquery.adk import BigQueryADKStore + + store = BigQueryADKStore(BigQueryConfig(extension_config={"adk": {"session_table": table, "enable_memory": False}})) + + assert store.session_table == table + assert f"CREATE TABLE IF NOT EXISTS {table}" in store._sessions_table_ddl() + + def test_adk_store_registration_validator_fails_fast_for_broken_store_mapping(monkeypatch: pytest.MonkeyPatch) -> None: """Broken adapter-to-store naming fails during config registration.""" from sqlspec.adapters.sqlite import SqliteConfig diff --git a/tests/unit/extensions/test_events/test_listener_hubs.py b/tests/unit/extensions/test_events/test_listener_hubs.py index e1cf2d486..60b221f2e 100644 --- a/tests/unit/extensions/test_events/test_listener_hubs.py +++ b/tests/unit/extensions/test_events/test_listener_hubs.py @@ -1035,7 +1035,7 @@ async def test_listener_hub_refuses_transaction_pooler_without_direct_dsn() -> N await psycopg_async_hub.subscribe("alerts") psycopg_sync_cfg = _PsycopgSyncConfig(_PsycopgSyncConnection()) - psycopg_sync_cfg.driver_features = {"pgbouncer": True, "direct_dsn": None} # type: ignore[attr-defined] + psycopg_sync_cfg.driver_features = {"pooler_mode": "transaction", "direct_dsn": None} # type: ignore[attr-defined] psycopg_sync_hub = PsycopgSyncListenerHub(psycopg_sync_cfg) # type: ignore[arg-type] with pytest.raises(ImproperConfigurationError, match="direct_dsn"): psycopg_sync_hub.subscribe("alerts") diff --git a/tests/unit/extensions/test_extension_config_boundaries.py b/tests/unit/extensions/test_extension_config_boundaries.py index aadf54099..e4281c2c3 100644 --- a/tests/unit/extensions/test_extension_config_boundaries.py +++ b/tests/unit/extensions/test_extension_config_boundaries.py @@ -5,7 +5,7 @@ from sqlspec.adapters.sqlite.adk.store import SqliteADKMemoryStore, SqliteADKStore from sqlspec.adapters.sqlite.config import SqliteConfig from sqlspec.adapters.sqlite.events.store import SqliteEventQueueStore -from sqlspec.adapters.sqlite.litestar.store import SQLiteStore +from sqlspec.adapters.sqlite.litestar.store import SqliteStore from sqlspec.exceptions import ImproperConfigurationError @@ -14,7 +14,7 @@ def test_extension_store_rejects_unused_migration_switch(extension: str) -> None with pytest.raises(ImproperConfigurationError, match="run_migrations"): config = SqliteConfig(extension_config={extension: {"run_migrations": True}}) if extension == "litestar": - SQLiteStore(config) + SqliteStore(config) elif extension == "events": SqliteEventQueueStore(config) else: @@ -35,3 +35,15 @@ def test_sqlite_adk_rejects_postgres_memory_tuning(option: str, value: object) - with pytest.raises(ImproperConfigurationError, match=option): config = SqliteConfig(extension_config={"adk": {option: value}}) SqliteADKMemoryStore(config) + + +@pytest.mark.parametrize( + "option", ["session_table", "events_table", "app_state_table", "user_state_table", "metadata_table", "memory_table"] +) +def test_sqlite_adk_rejects_qualified_table_names(option: str) -> None: + """Unsupported qualified names fail before any schema creation can partially succeed.""" + config = SqliteConfig(extension_config={"adk": {option: "main.custom_table"}}) + store_type = SqliteADKMemoryStore if option == "memory_table" else SqliteADKStore + + with pytest.raises(ValueError, match="Invalid table name"): + store_type(config) diff --git a/tests/unit/extensions/test_flask/test_correlation.py b/tests/unit/extensions/test_flask/test_correlation.py index 60b09913b..3806070f8 100644 --- a/tests/unit/extensions/test_flask/test_correlation.py +++ b/tests/unit/extensions/test_flask/test_correlation.py @@ -163,8 +163,8 @@ def test_endpoint() -> str: assert seen_correlation_id[0] is False -def test_flask_correlation_disabled_default_is_disabled() -> None: - """Correlation middleware should be disabled by default.""" +def test_flask_correlation_is_enabled_by_default() -> None: + """Correlation middleware is enabled by default, matching the other framework extensions.""" sqlspec = SQLSpec() config = SqliteConfig(connection_config={"database": ":memory:"}) sqlspec.add_config(config) @@ -180,7 +180,7 @@ def test_endpoint() -> str: with app.test_client() as client: response = client.get("/test", headers={"x-request-id": "default-test-123"}) assert response.status_code == 200 - assert seen_correlation_id[0] is False + assert seen_correlation_id[0] is True def test_flask_correlation_trace_headers_extracts_traceparent() -> None: diff --git a/tests/unit/extensions/test_sanic/test_middleware.py b/tests/unit/extensions/test_sanic/test_middleware.py index 785f182f2..7171becc9 100644 --- a/tests/unit/extensions/test_sanic/test_middleware.py +++ b/tests/unit/extensions/test_sanic/test_middleware.py @@ -162,7 +162,9 @@ def _make_config( ) -> MagicMock: config = MagicMock() config.supports_connection_pooling = supports_connection_pooling - config.extension_config = {"sanic": {"commit_mode": commit_mode, "disable_di": disable_di}} + config.extension_config = { + "sanic": {"commit_mode": commit_mode, "disable_di": disable_di, "enable_correlation_middleware": False} + } config.provide_connection = MagicMock(return_value=_AsyncConnectionManager(connection)) config.create_connection = AsyncMock(return_value=connection) config.statement_config = {} @@ -174,7 +176,9 @@ def _make_request(*, pool: object | None) -> SimpleNamespace: app_ctx = SimpleNamespace() if pool is not None: app_ctx.db_pool = pool - return SimpleNamespace(app=SimpleNamespace(ctx=app_ctx), ctx=SimpleNamespace(db_session_instance=object())) + return SimpleNamespace( + app=SimpleNamespace(ctx=app_ctx), ctx=SimpleNamespace(db_session_instance=object()), headers={} + ) class _AsyncConnectionManager: diff --git a/tests/unit/extensions/test_schema_reconciliation.py b/tests/unit/extensions/test_schema_reconciliation.py index 4b9e9e991..90b0f1f29 100644 --- a/tests/unit/extensions/test_schema_reconciliation.py +++ b/tests/unit/extensions/test_schema_reconciliation.py @@ -5,11 +5,11 @@ from sqlspec.adapters.sqlite.adk.store import SqliteADKMemoryStore, SqliteADKStore from sqlspec.adapters.sqlite.config import SqliteConfig from sqlspec.adapters.sqlite.events.store import SqliteEventQueueStore -from sqlspec.adapters.sqlite.litestar.store import SQLiteStore +from sqlspec.adapters.sqlite.litestar.store import SqliteStore from sqlspec.migrations.schema import SchemaTarget, ensure_schema_sync -class AdditiveSQLiteStore(SQLiteStore): +class AdditiveSQLiteStore(SqliteStore): """Session store target with one additive column.""" def _table_ddl(self) -> str: @@ -31,11 +31,7 @@ class AdditiveSqliteADKMemoryStore(SqliteADKMemoryStore): """ADK memory target with one additive column.""" def _memory_table_ddl(self) -> str: - return ( - super() - ._memory_table_ddl() - .replace("inserted_at REAL NOT NULL\n", "inserted_at REAL NOT NULL,\n schema_tag TEXT\n") - ) + return super()._memory_table_ddl().replace("embedding BLOB\n", "embedding BLOB,\n schema_tag TEXT\n") class AdditiveSqliteEventQueueStore(SqliteEventQueueStore): @@ -50,7 +46,7 @@ async def test_litestar_session_store_reconcile_uses_target_ddl(tmp_path: Path) connection_config={"database": str(tmp_path / "litestar.db")}, extension_config={"litestar": {"session_table": "sessions"}}, ) - await SQLiteStore(config).create_table() + await SqliteStore(config).create_table() store = AdditiveSQLiteStore(config) await store.reconcile_schema(assume_existing=True) diff --git a/tests/unit/extensions/test_storage_options.py b/tests/unit/extensions/test_storage_options.py index 9195b31aa..0f55215e4 100644 --- a/tests/unit/extensions/test_storage_options.py +++ b/tests/unit/extensions/test_storage_options.py @@ -20,7 +20,7 @@ from sqlspec.adapters.cockroach_psycopg.litestar import CockroachPsycopgAsyncStore, CockroachPsycopgSyncStore from sqlspec.adapters.sqlite.config import SqliteConfig from sqlspec.adapters.sqlite.events import SqliteEventQueueStore -from sqlspec.adapters.sqlite.litestar import SQLiteStore +from sqlspec.adapters.sqlite.litestar import SqliteStore from sqlspec.exceptions import ImproperConfigurationError @@ -138,7 +138,7 @@ def test_sqlite_litestar_pragma_profile_applied() -> None: "litestar": {"pragma_profile": True, "pragma_overrides": {"cache_size": -32000, "journal_mode": "WAL"}} } ) - store = SQLiteStore(config) + store = SqliteStore(config) connection = RecordingConnection() store.prepare_schema_sync(SimpleNamespace(connection=connection)) diff --git a/tests/unit/test_import_time.py b/tests/unit/test_import_time.py index 4b3a1b3df..4ba4446f6 100644 --- a/tests/unit/test_import_time.py +++ b/tests/unit/test_import_time.py @@ -200,3 +200,137 @@ def test_pytest_collect_only_succeeds_without_plugin_registration_errors() -> No [sys.executable, "-m", "pytest", "--collect-only", "-q", "tests"], capture_output=True, text=True, check=False ) assert result.returncode == 0, f"pytest --collect-only failed:\n{result.stderr}\n{result.stdout}" + + +ADAPTER_DRIVERS: dict[str, tuple[str, ...]] = { + "adbc": ("adbc_driver_manager",), + "aiomysql": ("aiomysql",), + "aiosqlite": ("aiosqlite",), + "arrow_odbc": ("arrow_odbc",), + "asyncmy": ("asyncmy",), + "asyncpg": ("asyncpg",), + "bigquery": ("google.cloud.bigquery",), + "cockroach_asyncpg": ("asyncpg",), + "cockroach_psycopg": ("psycopg",), + "db2": ("ibm_db", "ibm_db_dbi"), + "duckdb": ("duckdb",), + "mssql_python": ("mssql_python",), + "mysqlconnector": ("mysql.connector",), + "oracledb": ("oracledb",), + "psqlpy": ("psqlpy",), + "psycopg": ("psycopg",), + "pymssql": ("pymssql",), + "pymysql": ("pymysql",), + "spanner": ("google.cloud.spanner_v1",), + "sqlanywhere": ("sqlanydb",), + "sqlite": (), + "sybase": ("pymssql",), +} +"""Driver module(s) each adapter imports; every package under ``sqlspec/adapters/`` must be registered.""" + +DRIVER_PACKAGE_DEPENDENCIES: dict[str, tuple[str, ...]] = {"aiomysql": ("pymysql",)} +"""Other registered driver modules that a driver package itself depends on (aiomysql requires PyMySQL).""" + +SECONDARY_OPTIONAL_MODULES = ( + "pyarrow", + "pandas", + "polars", + "numpy", + "pgvector", + "rich", + "opentelemetry", + "litestar", + "pydantic", + "fsspec", + "obstore", + "google.cloud.storage", + "google.cloud.sql.connector", + "google.cloud.alloydb.connector", +) + +DRIVER_OWNED_OPTIONAL_IMPORTS: dict[str, frozenset[str]] = { + "adbc": frozenset({"pyarrow", "pandas", "polars", "numpy"}), + "arrow_odbc": frozenset({"pyarrow", "numpy"}), + "bigquery": frozenset({"pyarrow", "pandas", "numpy", "opentelemetry"}), + "mysqlconnector": frozenset({"opentelemetry"}), + "spanner": frozenset({"opentelemetry"}), +} +"""Optional modules the driver SDK imports on its own when installed; sqlspec code does not import them.""" + + +def test_every_adapter_package_is_registered() -> None: + """Every adapter package declares its driver so the import-contract gates cover it.""" + packages = { + path.name + for path in Path("sqlspec/adapters").iterdir() + if path.is_dir() and (path / "__init__.py").exists() and not path.name.startswith(("_", ".")) + } + assert packages == set(ADAPTER_DRIVERS) + + +@pytest.mark.parametrize("adapter", sorted(name for name, drivers in ADAPTER_DRIVERS.items() if drivers)) +def test_adapter_import_raises_driver_module_not_found(adapter: str) -> None: + """Importing an adapter without its driver raises the driver's own ModuleNotFoundError.""" + drivers = ADAPTER_DRIVERS[adapter] + script = f""" +import importlib.abc +import sys + +blocked = {drivers!r} + +class BlockDriver(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path, target=None): + if any(fullname == name or fullname.startswith(name + ".") for name in blocked): + raise ModuleNotFoundError(f"No module named {{fullname!r}}", name=fullname) + return None + +sys.meta_path.insert(0, BlockDriver()) +try: + import sqlspec.adapters.{adapter} +except ModuleNotFoundError as exc: + assert exc.name and any(exc.name == name or exc.name.startswith(name + ".") or name.startswith(exc.name + ".") for name in blocked), exc.name +else: + raise AssertionError("adapter imported without its driver") +""" + result = subprocess.run([sys.executable, "-c", script], capture_output=True, text=True, check=False) + assert result.returncode == 0, result.stderr + + +@pytest.mark.parametrize("adapter", sorted(ADAPTER_DRIVERS)) +def test_adapter_imports_with_every_other_driver_blocked(adapter: str) -> None: + """Each adapter works without any other adapter's driver installed.""" + own = {*ADAPTER_DRIVERS[adapter], *DRIVER_PACKAGE_DEPENDENCIES.get(adapter, ())} + blocked = sorted({name for drivers in ADAPTER_DRIVERS.values() for name in drivers} - own) + script = f""" +import importlib +import importlib.abc +import sys + +blocked = {blocked!r} + +class BlockOtherDrivers(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path, target=None): + if any(fullname == name or fullname.startswith(name + ".") for name in blocked): + raise ModuleNotFoundError(f"No module named {{fullname!r}}", name=fullname) + return None + +sys.meta_path.insert(0, BlockOtherDrivers()) +package = importlib.import_module("sqlspec.adapters.{adapter}") +for name in getattr(package, "__all__", ()): + getattr(package, name) +""" + result = subprocess.run([sys.executable, "-c", script], capture_output=True, text=True, check=False) + assert result.returncode == 0, result.stderr + + +@pytest.mark.parametrize("adapter", sorted(ADAPTER_DRIVERS)) +def test_adapter_import_does_not_load_secondary_optional_modules(adapter: str) -> None: + """Importing an adapter loads only its driver, plus optional modules the driver SDK imports itself.""" + allowed = sorted(DRIVER_OWNED_OPTIONAL_IMPORTS.get(adapter, frozenset())) + script = ( + f"import sys, importlib; importlib.import_module('sqlspec.adapters.{adapter}'); " + f"print(','.join(m for m in {SECONDARY_OPTIONAL_MODULES!r} if m in sys.modules and m not in {allowed!r}))" + ) + result = subprocess.run([sys.executable, "-c", script], capture_output=True, text=True, check=True) + leaked = [name for name in result.stdout.strip().split(",") if name] + assert not leaked, f"import sqlspec.adapters.{adapter} eagerly imported: {leaked}" diff --git a/tests/unit/test_mypyc_sybase_surface.py b/tests/unit/test_mypyc_sybase_surface.py index 2575c641f..37ce9573a 100644 --- a/tests/unit/test_mypyc_sybase_surface.py +++ b/tests/unit/test_mypyc_sybase_surface.py @@ -92,13 +92,13 @@ def test_mypyc_smoke_imports_include_sybase_and_sqlanywhere() -> None: assert smoke_by_name["sqlanywhere_pool"].require_compiled is True -def test_importing_sybase_and_sqlanywhere_adapters_does_not_eagerly_import_driver_packages() -> None: - """Importing sqlspec.adapters.sybase and sqlspec.adapters.sqlanywhere defers pymssql and sqlanydb.""" +def test_importing_sybase_and_sqlanywhere_adapters_does_not_import_rich() -> None: + """Importing sqlspec.adapters.sybase and sqlspec.adapters.sqlanywhere does not load rich.""" script = ( "import sys\n" "import sqlspec.adapters.sybase\n" "import sqlspec.adapters.sqlanywhere\n" - "leaked = [name for name in ('pymssql', 'sqlanydb', 'rich') if name in sys.modules]\n" + "leaked = [name for name in ('rich',) if name in sys.modules]\n" "assert not leaked, f'Eagerly imported: {leaked}'\n" ) subprocess.run([sys.executable, "-c", script], capture_output=True, text=True, check=True) diff --git a/tools/scripts/bench_adbc_pool.py b/tools/scripts/bench_adbc_pool.py new file mode 100644 index 000000000..3a57edd5b --- /dev/null +++ b/tools/scripts/bench_adbc_pool.py @@ -0,0 +1,276 @@ +"""Measure complete PostgreSQL ADBC sessions with unpooled, cloned and pooled connections.""" + +import hashlib +import json +import math +import os +import platform +import re +import shutil +import statistics +import subprocess +import sys +from concurrent.futures import ThreadPoolExecutor +from contextlib import ExitStack +from datetime import UTC, datetime +from importlib import metadata +from pathlib import Path +from threading import Barrier, BrokenBarrierError, Event, Lock +from time import perf_counter_ns +from typing import TYPE_CHECKING +from uuid import uuid4 + +import click + +from sqlspec.adapters.adbc import AdbcConfig + +if TYPE_CHECKING: + from adbc_driver_manager.dbapi import Connection + + +class CloneConfig(AdbcConfig): + """Close each session clone while retaining its original caller-owned seed.""" + + __slots__ = ("_clone_lock", "_seed") + + def __init__(self, uri: str, seed: "Connection") -> None: + super().__init__( + connection_config={"driver_name": "postgresql", "uri": uri, "autocommit": False}, + driver_features=_features(), + ) + self._seed = seed + self._clone_lock = Lock() + + def create_connection(self) -> "Connection": + with self._clone_lock: + return self._seed.adbc_clone() + + +@click.command() +@click.option("--uri-env", default="ADBC_POSTGRESQL_TEST_URI", show_default=True) +@click.option("--mode", type=click.Choice(["all", "unpooled", "clone", "pooled"]), default="all", show_default=True) +@click.option("--workers", type=click.Choice(["1", "4"]), default="1", show_default=True) +@click.option("--warmup", type=click.IntRange(min=1), default=100, show_default=True) +@click.option("--samples", type=click.IntRange(min=1), default=1000, show_default=True) +@click.option("--repeats", type=click.IntRange(min=1), default=7, show_default=True) +@click.option("--output", type=click.Path(path_type=Path, dir_okay=False)) +def main(uri_env: str, mode: str, workers: str, warmup: int, samples: int, repeats: int, output: Path | None) -> None: + """Compare per-session latency on independent worker rows in a disposable schema. + + Copy this script outside the checkout to run installed source/baseline wheels. + Set SQLSPEC_BENCH_REVISION to the verified source SHA for an installed wheel. + URI values and credentials are never included in the JSON report. + """ + uri = os.environ.get(uri_env) + if not uri: + raise click.ClickException(f"Set the connection environment variable {uri_env} before running") + provenance = _provenance() + worker_count = int(workers) + schema = "sqlspec_adbc_pool_bench_" + uuid4().hex + if re.fullmatch(r"sqlspec_adbc_pool_bench_[0-9a-f]{32}", schema) is None: + raise RuntimeError("Invalid generated benchmark schema") + table = schema + ".sessions" + unpooled = AdbcConfig( + connection_config={"driver_name": "postgresql", "uri": uri, "autocommit": False}, driver_features=_features() + ) + modes = ["unpooled", "clone", "pooled"] if mode == "all" else [mode] + results: list[dict[str, object]] = [] + combined: dict[str, list[int]] = {name: [] for name in modes} + with ExitStack() as cleanup: + configs = {"unpooled": unpooled} + if "clone" in modes: + seed = unpooled.create_connection() + cleanup.callback(seed.close) + configs["clone"] = CloneConfig(uri, seed) + if "pooled" in modes: + pooled = AdbcConfig( + connection_config={"driver_name": "postgresql", "uri": uri, "autocommit": False}, + driver_features=_features(), + pool_config={"max_size": worker_count, "timeout": 30.0}, + ) + cleanup.callback(pooled.close_pool) + pooled.create_pool() + configs["pooled"] = pooled + with unpooled.provide_session() as session: + server_version = str(session.select_value("SELECT current_setting('server_version')")) + server_identity = str( + session.select_value("SELECT inet_server_addr()::text || ':' || inet_server_port()::text") + ) + session.commit() + cleanup.callback(_drop_schema, unpooled, schema) + with unpooled.provide_session() as session: + session.execute(f"CREATE SCHEMA {schema}") + session.execute( + f"CREATE TABLE {table} (worker_id INTEGER PRIMARY KEY, value BIGINT NOT NULL, payload VARCHAR(32) NOT NULL)" + ) + for worker_id in range(1, worker_count + 1): + session.execute( + f"INSERT INTO {table} (worker_id, value, payload) VALUES (?, ?, ?)", worker_id, 0, "sqlspec" + ) + session.commit() + with ThreadPoolExecutor(max_workers=worker_count) as executor: + for repeat in range(repeats): + ordered = modes[repeat % len(modes) :] + modes[: repeat % len(modes)] + for name in ordered: + with unpooled.provide_session() as session: + session.execute(f"UPDATE {table} SET value = 0") + session.commit() + timings = _measure(executor, configs[name], table, worker_count, warmup, samples) + with unpooled.provide_session() as session: + rows = session.select(f"SELECT worker_id, value, payload FROM {table} ORDER BY worker_id") + expected = [ + {"worker_id": worker_id, "value": warmup + samples, "payload": "sqlspec"} + for worker_id in range(1, worker_count + 1) + ] + if rows != expected: + raise RuntimeError("Benchmark row increment or payload verification failed") + session.commit() + combined[name].extend(timings) + results.append({ + "mode": name, + "repeat": repeat + 1, + **_summary(timings), + "verified_rows": worker_count, + }) + report = { + "metadata": { + **provenance, + "timestamp": datetime.now(UTC).isoformat(), + "server_version": server_version, + "server_identifier_sha256": hashlib.sha256(server_identity.encode()).hexdigest(), + "workers": worker_count, + "warmup_per_worker": warmup, + "samples_per_worker": samples, + "repeats": repeats, + "timing_unit": "nanoseconds", + "workload": "begin/read/update/read/commit/session-release", + }, + "results": results, + "aggregate": {name: _summary(timings) for name, timings in combined.items()}, + } + encoded = json.dumps(report, indent=2, allow_nan=False) + "\n" + if output is None: + click.echo(encoded, nl=False) + else: + output.write_text(encoded, encoding="utf-8") + click.echo(f"Wrote benchmark results to {output}") + + +def _features() -> dict[str, bool]: + return { + "enable_pgvector": False, + "enable_paradedb": False, + "enable_pg_textsearch": False, + "enable_statement_profiling": False, + } + + +def _drop_schema(config: AdbcConfig, schema: str) -> None: + with config.provide_session() as session: + session.execute(f"DROP SCHEMA IF EXISTS {schema} CASCADE") + session.commit() + + +def _measure( + executor: ThreadPoolExecutor, config: AdbcConfig, table: str, workers: int, warmup: int, samples: int +) -> list[int]: + barrier = Barrier(workers, timeout=60) + stop = Event() + + def worker(worker_id: int) -> list[int]: + try: + for _ in range(warmup): + if stop.is_set(): + raise BrokenBarrierError("A benchmark peer failed") + transact(worker_id) + barrier.wait() + timings = [] + for _ in range(samples): + if stop.is_set(): + raise BrokenBarrierError("A benchmark peer failed") + start = perf_counter_ns() + transact(worker_id) + timings.append(perf_counter_ns() - start) + return timings + except BaseException: + stop.set() + barrier.abort() + raise + + def transact(worker_id: int) -> None: + with config.provide_session() as session: + session.begin() + before = int(session.select_value(f"SELECT value FROM {table} WHERE worker_id = ?", worker_id)) + session.execute(f"UPDATE {table} SET value = ? WHERE worker_id = ?", before + 1, worker_id) + after = session.select_value(f"SELECT value FROM {table} WHERE worker_id = ?", worker_id) + if after != before + 1: + raise RuntimeError("Transactional read-after-write verification failed") + session.commit() + + futures = [executor.submit(worker, worker_id) for worker_id in range(1, workers + 1)] + timings: list[int] = [] + failures: list[BaseException] = [] + for future in futures: + try: + timings.extend(future.result()) + except BaseException as exc: + failures.append(exc) + if failures: + raise next((exc for exc in failures if not isinstance(exc, BrokenBarrierError)), failures[0]) + return timings + + +def _summary(timings: list[int]) -> dict[str, object]: + if not timings or any(value <= 0 for value in timings): + raise RuntimeError("Missing or invalid session timings") + ordered = sorted(timings) + median = statistics.median(ordered) + deviation = statistics.pstdev(ordered) + return { + "count": len(ordered), + "median_ns": median, + "p95_ns": ordered[math.ceil(0.95 * len(ordered)) - 1], + "stddev_ns": deviation, + "noise_ratio": deviation / median, + "samples_ns": timings, + } + + +def _provenance() -> dict[str, object]: + import sqlspec + from sqlspec.driver import _sync + + origin = Path(sqlspec.__file__ or "").resolve() + source_root = origin.parent.parent + revision = os.environ.get("SQLSPEC_BENCH_REVISION") + dirty: bool | None = None + if (source_root / ".git").exists(): + git = shutil.which("git") + if git is None: + raise click.ClickException("Git is required to establish source-checkout provenance") + revision = subprocess.check_output([git, "-C", str(source_root), "rev-parse", "HEAD"], text=True).strip() + dirty = bool(subprocess.check_output([git, "-C", str(source_root), "status", "--porcelain"], text=True).strip()) + if not revision or re.fullmatch(r"[0-9a-f]{7,40}", revision) is None: + raise click.ClickException("Set SQLSPEC_BENCH_REVISION to the verified source SHA for installed packages") + versions = {} + for distribution in ("sqlspec", "adbc-driver-manager", "adbc-driver-postgresql", "pyarrow", "sqlglot"): + try: + versions[distribution] = metadata.version(distribution) + except metadata.PackageNotFoundError: + versions[distribution] = "not-installed" + return { + "revision": revision, + "source_dirty": dirty, + "harness_sha256": hashlib.sha256(Path(__file__).read_bytes()).hexdigest(), + "sqlspec_origin": str(origin), + "driver_origin": _sync.__file__, + "mypyc_compiled": (_sync.__file__ or "").endswith((".so", ".pyd")), + "python": sys.version, + "executable": sys.executable, + "platform": platform.platform(), + "versions": versions, + } + + +if __name__ == "__main__": + main() diff --git a/tools/scripts/mypyc_smoke.py b/tools/scripts/mypyc_smoke.py index 81f15174f..f25bca47c 100644 --- a/tools/scripts/mypyc_smoke.py +++ b/tools/scripts/mypyc_smoke.py @@ -44,8 +44,8 @@ class SmokeImport(NamedTuple): SmokeImport("storage_backend_fsspec", "sqlspec.storage.backends.fsspec", "FSSpecBackend", True), SmokeImport("storage_backend_obstore", "sqlspec.storage.backends.obstore", "ObStoreBackend", True), SmokeImport("sqlite_pool", "sqlspec.adapters.sqlite.pool", "SqliteConnectionPool", True), - SmokeImport("sybase_pool", "sqlspec.adapters.sybase.pool", "SybaseConnectionPool", True), - SmokeImport("sqlanywhere_pool", "sqlspec.adapters.sqlanywhere.pool", "SqlanywhereConnectionPool", True), + SmokeImport("sybase_pool", "sqlspec.adapters.sybase.pool", "SybaseConnectionPool", True, "pymssql"), + SmokeImport("sqlanywhere_pool", "sqlspec.adapters.sqlanywhere.pool", "SqlanywhereConnectionPool", True, "sqlanydb"), SmokeImport("data_dictionary_registry", "sqlspec.data_dictionary._registry", "get_dialect_config", True), SmokeImport("data_dictionary_loader", "sqlspec.data_dictionary._loader", "DataDictionaryLoader", True), SmokeImport("pgvector_dialect", "sqlspec.dialects.postgres._pgvector", "PGVector"), diff --git a/uv.lock b/uv.lock index d5a2a3e77..c5201c2a5 100644 --- a/uv.lock +++ b/uv.lock @@ -3,17 +3,14 @@ revision = 5 requires-python = ">=3.11, <4.0" resolution-markers = [ "python_full_version >= '3.15' and sys_platform == 'win32'", - "python_full_version >= '3.15' and sys_platform == 'emscripten'", - "python_full_version >= '3.15' and sys_platform != 'emscripten' and sys_platform != 'win32'", "python_full_version == '3.14.*' and sys_platform == 'win32'", + "python_full_version >= '3.15' and sys_platform == 'emscripten'", "python_full_version == '3.14.*' and sys_platform == 'emscripten'", + "python_full_version >= '3.15' and sys_platform != 'emscripten' and sys_platform != 'win32'", "python_full_version == '3.14.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", - "python_full_version == '3.13.*' and sys_platform == 'win32'", - "python_full_version == '3.13.*' and sys_platform == 'emscripten'", - "python_full_version == '3.13.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", - "python_full_version == '3.12.*' and sys_platform == 'win32'", - "python_full_version == '3.12.*' and sys_platform == 'emscripten'", - "python_full_version == '3.12.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'win32'", + "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'emscripten'", + "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32'", "python_full_version < '3.12' and sys_platform == 'win32'", "python_full_version < '3.12' and sys_platform == 'emscripten'", "python_full_version < '3.12' and sys_platform != 'emscripten' and sys_platform != 'win32'", @@ -509,19 +506,19 @@ wheels = [ [[package]] name = "arrow-odbc" -version = "10.6.1" +version = "10.6.2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "cffi" }, { name = "pyarrow" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/12/f8/7d0d0b88cadd9d28050ad7bce851af086857eb1401d09b54b26045bfbe99/arrow_odbc-10.6.1.tar.gz", hash = "sha256:f7449d292632f455f72e7ad8bb6b4fe386ac5602738225526294e1ffee0374d5", size = 107639, upload-time = "2026-10-19T20:22:59.75Z" } +sdist = { url = "https://files.pythonhosted.org/packages/97/76/990a0332ab2285a870ac1b39197917413f6fbca2ff192f859a5cc6d25575/arrow_odbc-10.6.2.tar.gz", hash = "sha256:badfa0587e9d569e265082f485686b32df0434734277b41348220268aa9e4152", size = 108670, upload-time = "2026-10-25T14:34:07.167Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/7d/17/8c32bd7afa88b6cb8ee120e0051218d2bab52dd7e1e60387230ac5b00ef1/arrow_odbc-10.6.1-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:4de744a4f20ebe574030bb1943f9aa10cb242b77ac7128a1efd1539d943cee6e", size = 753105, upload-time = "2026-10-19T20:25:14.844Z" }, - { url = "https://files.pythonhosted.org/packages/88/7f/e7c884fe16844f38899053f0e2e9f5ad66a984bb79823b1aef6d17b743d5/arrow_odbc-10.6.1-py3-none-macosx_11_0_arm64.whl", hash = "sha256:eafdf71851aac18c37822a67b57d704f2b9a7b133205da253445d374ca6c90c2", size = 735430, upload-time = "2026-10-19T20:22:58.19Z" }, - { url = "https://files.pythonhosted.org/packages/67/54/0deab13f42cbdd8aaeec9b957164b49b54d805fb3c94ca27d3857a07ee09/arrow_odbc-10.6.1-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:72d97c33cc41bde4b071bad8b20932e360d11f4dad562bc0687cd59b8a11a3cc", size = 843468, upload-time = "2026-10-19T20:28:59.043Z" }, - { url = "https://files.pythonhosted.org/packages/94/bb/920f8e03a0340b8c8ad945c6fde73df07519596def2e40ff6ef70358ec8f/arrow_odbc-10.6.1-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:78d48b9fe0291b32f50e8b1fb81c5bb2cd74fe7f636b34052551070a51ae2a5a", size = 869855, upload-time = "2026-10-19T20:25:20.298Z" }, - { url = "https://files.pythonhosted.org/packages/ec/12/311cd3f26939ead9b7f276988b9bd9d7b049c47c84af6b71cd35bd3d24c0/arrow_odbc-10.6.1-py3-none-win_amd64.whl", hash = "sha256:6226df0d91af94bdb47f575933e6c704e46d525d8b37348c43f7207f4b7d5e9b", size = 654895, upload-time = "2026-10-19T20:25:03.548Z" }, + { url = "https://files.pythonhosted.org/packages/68/a4/c7e9406f07406b61d6e7c7c491f5e6cf7602a4d49727adcaf76376d03524/arrow_odbc-10.6.2-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:d1603f351058daa8fc9d7d6df9c5da5bd57a0d2368b21cc8710928f540551073", size = 759654, upload-time = "2026-10-25T14:32:14.506Z" }, + { url = "https://files.pythonhosted.org/packages/f5/33/212da57507d06cf4d7f452e77339be3c66500b76a886b26b7a106a1968d5/arrow_odbc-10.6.2-py3-none-macosx_11_0_arm64.whl", hash = "sha256:50fea0f9441b7cef2ea689cd90705ad4f9d214ec65488bb256354ca2f77ce1b4", size = 739566, upload-time = "2026-10-25T14:34:06.155Z" }, + { url = "https://files.pythonhosted.org/packages/54/3a/105b89a01a25b8bd7b8a19ccd9cac9dfe03498df8f7553956c8ca23d076a/arrow_odbc-10.6.2-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:54e25aa9c094c713298741048b5725530bbb1f5637adc3917b935045fd0c08e4", size = 846711, upload-time = "2026-10-25T14:32:23.305Z" }, + { url = "https://files.pythonhosted.org/packages/36/1c/3856b0cfeb0e84b8c9b1b1744a2a393a5b6553580039588cf6da3d8af80a/arrow_odbc-10.6.2-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:3cf90b019f4a52c771e15a7847504f7826061aa04b77d290353f608abf62ef07", size = 872316, upload-time = "2026-10-25T14:32:53.888Z" }, + { url = "https://files.pythonhosted.org/packages/f1/80/2976600c3e8f0693c9901cf544f741aa6195d0027e5a5db2174fe3822d0a/arrow_odbc-10.6.2-py3-none-win_amd64.whl", hash = "sha256:763a49240a8a38f86974789ca0692560e02d60b389ad6abb8040cd2338c63773", size = 659820, upload-time = "2026-10-25T14:32:37.285Z" }, ] [[package]] @@ -846,7 +843,7 @@ wheels = [ [[package]] name = "bump-my-version" -version = "1.6.5" +version = "1.6.6" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "click" }, @@ -859,9 +856,9 @@ dependencies = [ { name = "tomlkit" }, { name = "wcmatch" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/a9/65/338597a14c53ae2b9d004031dbcb61a0641484b20469b5b7bd5250a37a34/bump_my_version-1.6.5.tar.gz", hash = "sha256:0fd40347bca359bdee9999c93e101d2ad4612cf10319bc5e81181e57fc313e0d", size = 1158255, upload-time = "2026-10-24T14:02:28.037Z" } +sdist = { url = "https://files.pythonhosted.org/packages/2b/df/ff4a0fc20547b4bf30b91fb6300cf0bd91d40a33dbb3cdb9cc698d0e71c0/bump_my_version-1.6.6.tar.gz", hash = "sha256:3f0a140b2a275dc030b2adc4dacb86247101ede38a5032d9fc9f291a17cfeb9d", size = 1161185, upload-time = "2026-10-25T13:45:13.457Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/15/bd/a8698a2143cccd7d5a5f23e7ec8520c2133e4083ccda1799531490fad17d/bump_my_version-1.6.5-py3-none-any.whl", hash = "sha256:7fc77cfa499edfa18072aaf7cc9e6697b8a66c7ba76f467c3c4fba4c554f2521", size = 67396, upload-time = "2026-10-24T14:02:26.214Z" }, + { url = "https://files.pythonhosted.org/packages/19/21/01f6f456e466363ef970a69885403bba3337c5a442bc4a7f52d2f5a26dcd/bump_my_version-1.6.6-py3-none-any.whl", hash = "sha256:99c16933a7ca5c7ebeb622ab6bf5a964ad0bda40c8e5f62311b27d9414edf123", size = 68031, upload-time = "2026-10-25T13:45:11.376Z" }, ] [[package]] @@ -1384,11 +1381,11 @@ wheels = [ [[package]] name = "deepmerge" -version = "3.0.1" +version = "4.0.0" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/38/6e/5cb3548b4d3112fea529375e55e6f3cdc52b8054e3a66f203b1f888ba885/deepmerge-3.0.1.tar.gz", hash = "sha256:35b39a4cb92cf328d6eca61cbbf65f68a37c2ceb3085f0f853cbb2e52a59fc23", size = 22328, upload-time = "2026-09-01T14:09:44.383Z" } +sdist = { url = "https://files.pythonhosted.org/packages/ca/e9/5e99e56e03f9461858e686fdaf82e2516671bc1bd46f8b7ad7b9a8375f6f/deepmerge-4.0.0.tar.gz", hash = "sha256:948565f2bc9cdb41259a4a23d26638e5840c9e084a343ab279ffd3edcfdf74d3", size = 22792, upload-time = "2026-10-25T19:25:28.967Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/20/91/600003aaad107e27553fbc9cbfc57e96fa37e0223d2ef9d7d3a0e8d8d070/deepmerge-3.0.1-py3-none-any.whl", hash = "sha256:35c96f6a68fcf90719a5b31d9f8042ecef6c00fb56836660d33455d0f5cfda65", size = 14909, upload-time = "2026-09-01T14:09:43.364Z" }, + { url = "https://files.pythonhosted.org/packages/c2/89/362146b2ec7d0d503f824ac8496c88f0408ee0d39203b02cc10feac08146/deepmerge-4.0.0-py3-none-any.whl", hash = "sha256:a36dd5810fc3b90bd3e0b681158221361d6a31a6109411b27b211b371fed6d5b", size = 15429, upload-time = "2026-10-25T19:25:28.012Z" }, ] [[package]] @@ -3743,17 +3740,14 @@ version = "2.5.4" source = { registry = "https://pypi.org/simple" } resolution-markers = [ "python_full_version >= '3.15' and sys_platform == 'win32'", - "python_full_version >= '3.15' and sys_platform == 'emscripten'", - "python_full_version >= '3.15' and sys_platform != 'emscripten' and sys_platform != 'win32'", "python_full_version == '3.14.*' and sys_platform == 'win32'", + "python_full_version >= '3.15' and sys_platform == 'emscripten'", "python_full_version == '3.14.*' and sys_platform == 'emscripten'", + "python_full_version >= '3.15' and sys_platform != 'emscripten' and sys_platform != 'win32'", "python_full_version == '3.14.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", - "python_full_version == '3.13.*' and sys_platform == 'win32'", - "python_full_version == '3.13.*' and sys_platform == 'emscripten'", - "python_full_version == '3.13.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", - "python_full_version == '3.12.*' and sys_platform == 'win32'", - "python_full_version == '3.12.*' and sys_platform == 'emscripten'", - "python_full_version == '3.12.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'win32'", + "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'emscripten'", + "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32'", ] sdist = { url = "https://files.pythonhosted.org/packages/95/b0/c7453d0b6e2073c3264468b106ee1563750cecc910965e67357e3698c83e/numpy-2.5.4.tar.gz", hash = "sha256:9a94cf751c9ad8ebaa835bcd3d40dacf8534ad086b88c38029b65123c7999d2a", size = 20866315, upload-time = "2026-10-10T20:05:31.422Z" } wheels = [ @@ -4204,30 +4198,30 @@ wheels = [ [[package]] name = "polars" -version = "2.0.0" +version = "2.1.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "polars-runtime-32" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/8e/e9/001f371ec6a1bb54893f599ceebd56e6144fed4091f09f09fec0021a9276/polars-2.0.0.tar.gz", hash = "sha256:62da109e27a19a9d36657ee25dc035c9d3f87e7bd610526fe467dc37ea7dc115", size = 778215, upload-time = "2026-10-06T11:51:29.679Z" } +sdist = { url = "https://files.pythonhosted.org/packages/b9/da/c27daf205a824f89f71172a75ad0098c5837a5a31ddeaac7f8737ea8574b/polars-2.1.0.tar.gz", hash = "sha256:9eee488d6e10773c028753dfc8dac647fd60a708108dcb818a733e85d33df253", size = 811508, upload-time = "2026-10-25T18:00:33.43Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/ac/09/cc33bbd5463749c116b62c204d88bed6c02a6cb901eac7adab0d38651b07/polars-2.0.0-py3-none-any.whl", hash = "sha256:35d62f3541b7a6d4c360a2e2f07fccc0c2bcbd33b0ea51c83a25417a47a3f3ad", size = 876611, upload-time = "2026-10-06T11:44:04.327Z" }, + { url = "https://files.pythonhosted.org/packages/d5/d4/e696a46c74ec8cac6434764deaf72f7edcc2e91c2145a9b3f61900a7aa7a/polars-2.1.0-py3-none-any.whl", hash = "sha256:63c50ce3533f5a03d17c3a2a24f1d0e0c7eda98fbde4268a3cba900d34060e67", size = 911474, upload-time = "2026-10-25T17:58:26.863Z" }, ] [[package]] name = "polars-runtime-32" -version = "2.0.0" +version = "2.1.0" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/34/ad/dbb6f6d7070867951532bcfe5e6a648d8777b416b18cddabc07030404e8c/polars_runtime_32-2.0.0.tar.gz", hash = "sha256:b5f9afcc742b4a67eabd2c680ff0f12eb02ede9b4bf807bffabd6dbb9a58d5c7", size = 3591339, upload-time = "2026-10-06T11:51:31.076Z" } +sdist = { url = "https://files.pythonhosted.org/packages/e0/a6/46e09af693326735b1fba3a8bf5a725b57136fb1342e707e75d3113c826b/polars_runtime_32-2.1.0.tar.gz", hash = "sha256:c575f5f3cdb8a5f9be65477b9b4202b062593687fe62a0082add87fbd4571067", size = 3878004, upload-time = "2026-10-25T18:00:34.932Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/82/88/d35dec6c8928dfbaa1cccf9b626a1067da906e792c92d9f994ca825ab2b5/polars_runtime_32-2.0.0-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:ffb7ac6cf4e8c4a652df1951e3c3840c7c23a033603d5a9efd422fa8dd699d82", size = 52494314, upload-time = "2026-10-06T11:44:07.768Z" }, - { url = "https://files.pythonhosted.org/packages/5f/fd/2237bf53ffaff47cdf1edc6c10587a7a6444d4951150eeb08d84f3493ff8/polars_runtime_32-2.0.0-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:7012d8a0201bd95638545ce8f256c0efe2c5cab0f806eb043021dddde5a9498b", size = 47930083, upload-time = "2026-10-06T11:44:11.592Z" }, - { url = "https://files.pythonhosted.org/packages/0d/0d/85e3ed90417996fc09770be91b39979074fe2978fc15b431bf8a9459760d/polars_runtime_32-2.0.0-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:8b85bb42e6009acc9629afcc70a83473fd468694d6a30ffb0ab376c8dd1a0a17", size = 50417889, upload-time = "2026-10-06T11:50:20.774Z" }, - { url = "https://files.pythonhosted.org/packages/83/88/e9fecfd49159da92f54ff2445883577a0f1bc195da53ecc9535c458d55dd/polars_runtime_32-2.0.0-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0d6ac584ea2b38913784db943879412380d92e28ab9cb88e20a77ba71ba3f911", size = 54475036, upload-time = "2026-10-06T11:50:24.411Z" }, - { url = "https://files.pythonhosted.org/packages/48/ad/b2abf732697b21467aaaeaac0f3bf7eee0d89c59ce8125f1ed41b28a2d97/polars_runtime_32-2.0.0-cp310-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:a6bf5e260e0a6f00d0f9181438fe9e45776df8c66cee9cba16e3675cc3888488", size = 50579474, upload-time = "2026-10-06T11:50:28.377Z" }, - { url = "https://files.pythonhosted.org/packages/7f/05/304deee59a95865e1b5e9ec7b066069b49093b81b768f473d9d3b165c686/polars_runtime_32-2.0.0-cp310-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:55c26eef325b6840584d91aac232e9cf3ac19e1b904594b9b54131be1edeab4d", size = 54413293, upload-time = "2026-10-06T11:50:31.828Z" }, - { url = "https://files.pythonhosted.org/packages/61/59/8c9fd7199f7c4eb1b64e640306a946a2e4a46337b3bbb33b840972c7d84b/polars_runtime_32-2.0.0-cp310-abi3-win_amd64.whl", hash = "sha256:7da1caf3c7b4f397fb213c984013a0c755557619a2d511899a1ff74392484078", size = 54229989, upload-time = "2026-10-06T11:50:35.206Z" }, - { url = "https://files.pythonhosted.org/packages/e2/93/43608026f38aa6ed4d22da8597706a61682ee403caef0021ce8e6dc73227/polars_runtime_32-2.0.0-cp310-abi3-win_arm64.whl", hash = "sha256:c30ba698c8904048df4a9bc3d6c5033cc2d0a7cbb0e13f4fd2de5a1947b61994", size = 48730655, upload-time = "2026-10-06T11:50:38.756Z" }, + { url = "https://files.pythonhosted.org/packages/7a/40/b424f3946c702ff48eef14e44d35df81a25a784f20f77355f39c746528fc/polars_runtime_32-2.1.0-cp311-abi3-macosx_10_12_x86_64.whl", hash = "sha256:de243b6e5a0b9f3abeaefe542136fcdc6dd8df5be3c988a40838d075314e246c", size = 52169672, upload-time = "2026-10-25T17:58:32.437Z" }, + { url = "https://files.pythonhosted.org/packages/fe/38/07066430ac655a8cab2f9e96ea1fec6ca1f20ec47844afe117662d43ace0/polars_runtime_32-2.1.0-cp311-abi3-macosx_11_0_arm64.whl", hash = "sha256:317314d65949f1d58c4939adf83aa0d2cfaf8da50fd4a725a2de27bde2c10018", size = 47626883, upload-time = "2026-10-25T17:58:38.688Z" }, + { url = "https://files.pythonhosted.org/packages/49/a8/a773e20b451bbaab4d8267867b7c43a0c3aa67ebb6f781423a42a5a1af37/polars_runtime_32-2.1.0-cp311-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:777228cc8bf9cda36dee85ba74d13a25974ad4caffc5da93e10b1da9d810ce2d", size = 50033514, upload-time = "2026-10-25T17:58:44.261Z" }, + { url = "https://files.pythonhosted.org/packages/d2/e6/649c524252d32b24d0ba92278033823fa8e0040ee3748ce5afbeda107403/polars_runtime_32-2.1.0-cp311-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0b838056cf365d7497971d9d9825c547766923338771c48c545b86ead2ebc572", size = 53863954, upload-time = "2026-10-25T17:58:48.737Z" }, + { url = "https://files.pythonhosted.org/packages/90/f7/2eb405bf5ead6373a33fa5dd4637fe996565a1f07ccce8911ed167673740/polars_runtime_32-2.1.0-cp311-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:04cfc3db1571b74ea4145944de7feceed4bfc9b26e5b2c265e27ad3b2d3866a9", size = 50204290, upload-time = "2026-10-25T17:58:53.33Z" }, + { url = "https://files.pythonhosted.org/packages/33/4d/d76cd42042605e079d6df91cc4fea09ed5f4cb836c72c952a17f585860a9/polars_runtime_32-2.1.0-cp311-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:af58b5f6e5cc44b578f2a93d752f87c135c3a53205b4360301bc3a9b89dfbfd1", size = 53794253, upload-time = "2026-10-25T17:58:57.99Z" }, + { url = "https://files.pythonhosted.org/packages/4d/4e/2a419d0def2908255dd1b7318c48430319859dc92ca156e2f81ff7bcd78b/polars_runtime_32-2.1.0-cp311-abi3-win_amd64.whl", hash = "sha256:f2fe32f48af2395cd6ec52a433a03970a69aab1dbb66a613729fd549ce92c904", size = 53926265, upload-time = "2026-10-25T17:59:03.13Z" }, + { url = "https://files.pythonhosted.org/packages/c8/1b/15ffba3264baf21d2f04c51f54b01159cb6f4fad4e84de1404be552d601f/polars_runtime_32-2.1.0-cp311-abi3-win_arm64.whl", hash = "sha256:0af30dd53a7725f703d0d21dd421a39e0165378bfaa41431008f286608376d71", size = 48636436, upload-time = "2026-10-25T17:59:08.476Z" }, ] [[package]] @@ -5067,7 +5061,7 @@ wheels = [ [[package]] name = "pytest-databases" version = "0.19.0" -source = { git = "https://github.com/litestar-org/pytest-databases.git?branch=test%2Fclientless-all#9f6fc17c5703f618b5895fa43ce62e736e44c28b" } +source = { git = "https://github.com/litestar-org/pytest-databases.git?branch=test%2Fclientless-all#09a0f8f82b4513c90c2c9628cd762ce75de31d5d" } dependencies = [ { name = "docker" }, { name = "filelock" }, @@ -5945,6 +5939,7 @@ dev = [ { name = "slotscheck" }, { name = "sniffio" }, { name = "sqlalchemy", extra = ["asyncio"] }, + { name = "sqlanydb" }, { name = "types-cffi" }, { name = "types-colorama" }, { name = "types-protobuf" }, @@ -5978,6 +5973,7 @@ extras = [ { name = "pgvector" }, { name = "polars" }, { name = "pyarrow" }, + { name = "sqlanydb" }, ] lint = [ { name = "asyncpg-stubs" }, @@ -6137,6 +6133,7 @@ dev = [ { name = "slotscheck", specifier = ">=0.16.5" }, { name = "sniffio" }, { name = "sqlalchemy", extras = ["asyncio"] }, + { name = "sqlanydb", specifier = ">=1.0.14" }, { name = "types-cffi" }, { name = "types-colorama" }, { name = "types-protobuf" }, @@ -6170,6 +6167,7 @@ extras = [ { name = "pgvector" }, { name = "polars" }, { name = "pyarrow" }, + { name = "sqlanydb", specifier = ">=1.0.14" }, ] lint = [ { name = "asyncpg-stubs" },