sqlspec 0.69.0__py3-none-any.whl → 0.70.0__py3-none-any.whl

sqlspec/__main__.py

@@ -1,3 +1,5 @@
+"""Command-line entry point for executing SQLSpec via ``python -m sqlspec``."""
+
 from sqlspec.cli import add_migration_commands as build_cli_interface
 
 __all__ = ("run_cli",)

sqlspec/_typing.py

@@ -38,6 +38,7 @@ if TYPE_CHECKING:
     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 @@ __all__ = (
     "UNSET",
     "UNSET_STUB",
     "UUID_UTILS_INSTALLED",
+    "ArrowException",
     "ArrowRecordBatch",
     "ArrowRecordBatchReader",
     "ArrowRecordBatchReaderProtocol",
@@ -472,6 +474,13 @@ _ARROW_SCHEMA_SHIM = ArrowSchemaProtocol
 _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 @@ _CATTRS_UNSTRUCTURE_SHIM = cattrs_unstructure_stub
 _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),

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 @@ __all__ = (
     "AdbcDriver",
     "AdbcDriverFeatures",
     "AdbcExceptionHandler",
+    "AdbcPoolConfig",
 )

sqlspec/adapters/adbc/_typing.py

@@ -7,6 +7,7 @@ compilation to avoid ABI boundary issues.
 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 @@ __all__ = (
     "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 @@ class AdbcSessionContext:
         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
 
 

sqlspec/adapters/adbc/adk/store.py

@@ -8,11 +8,14 @@ from functools import lru_cache
 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 import BaseSyncADKStore, StoredEvent, StoredSession,
 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 @@ class AdbcADKStore(BaseSyncADKStore["AdbcConfig"]):
         """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 @@ class AdbcADKMemoryStore(BaseSyncADKMemoryStore["AdbcConfig"]):
         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 @@ class AdbcADKMemoryStore(BaseSyncADKMemoryStore["AdbcConfig"]):
                     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))