sqlspec 0.64.0__py3-none-any.whl → 0.65.0__py3-none-any.whl

sqlspec/_typing.py

@@ -50,6 +50,7 @@ __all__ = (
     "CATTRS_INSTALLED",
     "CLOUD_SQL_CONNECTOR_INSTALLED",
     "FSSPEC_INSTALLED",
+    "IBM_DB_INSTALLED",
     "LITESTAR_INSTALLED",
     "MSGSPEC_INSTALLED",
     "NANOID_INSTALLED",
@@ -666,6 +667,7 @@ ATTRS_INSTALLED = dependency_flag("attrs")
 CATTRS_INSTALLED = dependency_flag("cattrs")
 CLOUD_SQL_CONNECTOR_INSTALLED = dependency_flag("google.cloud.sql.connector")
 FSSPEC_INSTALLED = dependency_flag("fsspec")
+IBM_DB_INSTALLED = dependency_flag("ibm_db")
 LITESTAR_INSTALLED = dependency_flag("litestar")
 MSGSPEC_INSTALLED = dependency_flag("msgspec")
 NUMPY_INSTALLED = dependency_flag("numpy")

sqlspec/adapters/adbc/_typing.py

@@ -31,8 +31,7 @@ if TYPE_CHECKING:
     AdbcConnection: TypeAlias = _AdbcConnection
     AdbcRawCursor: TypeAlias = _AdbcRawCursor
     AdbcNativeError: TypeAlias = _AdbcNativeError
-
-if not TYPE_CHECKING:
+else:
     AdbcConnection = _AdbcConnection
     AdbcRawCursor = _AdbcRawCursor
     AdbcNativeError = _AdbcNativeError

sqlspec/adapters/adbc/adk/store.py

@@ -1,10 +1,16 @@
 """ADBC ADK store for Google Agent Development Kit session/event storage."""
 
 import contextlib
-import re
 from datetime import datetime, timedelta, timezone
+from functools import lru_cache
 from typing import TYPE_CHECKING, Any, Final, Literal
 
+from sqlspec.adapters.adbc.core import (
+    driver_kind_from_driver_name,
+    is_shared_object_driver,
+    resolve_dialect_from_config,
+)
+from sqlspec.core.parameters._validator import PARAMETER_REGEX
 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
@@ -36,6 +42,21 @@ ADBC_TABLE_NOT_FOUND_PATTERNS: Final = (
 )
 
 
+@lru_cache(maxsize=256)
+def _postgresql_store_sql(sql: str) -> str:
+    """Cache positional SQL conversion without changing quoted or commented text."""
+    segments: list[str] = []
+    last_end = 0
+    index = 0
+    for match in PARAMETER_REGEX.finditer(sql):
+        if match.lastgroup == "qmark":
+            index += 1
+            segments.extend((sql[last_end : match.start()], f"${index}"))
+            last_end = match.end()
+    segments.append(sql[last_end:])
+    return "".join(segments)
+
+
 class AdbcADKStore(BaseSyncADKStore["AdbcConfig"]):
     """ADBC synchronous ADK store for Arrow Database Connectivity.
 
@@ -192,15 +213,28 @@ class AdbcADKStore(BaseSyncADKStore["AdbcConfig"]):
         Returns:
             Dialect identifier for DDL generation.
         """
-        driver_name = self._config.connection_config.get("driver_name", "").lower()
-
-        if "postgres" in driver_name:
+        driver_name = self._config.connection_config.get("driver_name", "")
+        if (
+            isinstance(driver_name, str)
+            and driver_name
+            and driver_kind_from_driver_name(driver_name) is None
+            and not is_shared_object_driver(driver_name)
+        ):
+            logger.warning(
+                "Unknown ADBC driver: %s. Using generic SQL dialect. "
+                "Consider using a direct adapter for better performance.",
+                driver_name,
+            )
+            return DIALECT_GENERIC
+
+        resolved = resolve_dialect_from_config(self._config.connection_config)
+        if resolved == "postgres":
             return DIALECT_POSTGRESQL
-        if "sqlite" in driver_name:
+        if resolved == "sqlite":
             return DIALECT_SQLITE
-        if "duckdb" in driver_name:
+        if resolved == "duckdb":
             return DIALECT_DUCKDB
-        if "snowflake" in driver_name:
+        if resolved == "snowflake":
             return DIALECT_SNOWFLAKE
 
         logger.warning(
@@ -605,14 +639,7 @@ class AdbcADKStore(BaseSyncADKStore["AdbcConfig"]):
         """Return SQL with dialect-appropriate positional placeholders."""
         if self._dialect != DIALECT_POSTGRESQL:
             return sql
-        index = 0
-
-        def replace_placeholder(_match: Any) -> str:
-            nonlocal index
-            index += 1
-            return f"${index}"
-
-        return re.sub(r"\?", replace_placeholder, sql)
+        return _postgresql_store_sql(sql)
 
     def _execute(self, cursor: Any, sql: str, params: "tuple[Any, ...] | list[Any]") -> Any:
         """Execute parameterized SQL using the current ADBC dialect's placeholder style."""
@@ -961,37 +988,41 @@ class AdbcADKStore(BaseSyncADKStore["AdbcConfig"]):
         with self._config.provide_connection() as conn:
             cursor = conn.cursor()
             try:
-                self._execute(
-                    cursor,
-                    insert_sql,
-                    (
-                        event_record["id"],
-                        event_record["app_name"],
-                        event_record["user_id"],
-                        event_record["session_id"],
-                        event_record["invocation_id"],
-                        event_record["timestamp"],
-                        event_data,
-                    ),
-                )
                 self._execute(cursor, update_sql, (state_json, app_name, user_id, session_id))
-                if app_state is not None:
-                    self._execute(cursor, delete_app_state_sql, (app_name,))
-                    self._execute(
-                        cursor,
-                        insert_app_state_sql,
-                        (app_name, self._serialize_state(app_state), datetime.now(timezone.utc)),
-                    )
-                if user_state is not None:
-                    self._execute(cursor, delete_user_state_sql, (app_name, user_id))
+                self._execute(cursor, select_sql, (app_name, user_id, session_id))
+                row = cursor.fetchone()
+                if row is not None:
                     self._execute(
                         cursor,
-                        insert_user_state_sql,
-                        (app_name, user_id, self._serialize_state(user_state), datetime.now(timezone.utc)),
+                        insert_sql,
+                        (
+                            event_record["id"],
+                            event_record["app_name"],
+                            event_record["user_id"],
+                            event_record["session_id"],
+                            event_record["invocation_id"],
+                            event_record["timestamp"],
+                            event_data,
+                        ),
                     )
-                self._execute(cursor, select_sql, (app_name, user_id, session_id))
-                row = cursor.fetchone()
-                conn.commit()
+                    if app_state is not None:
+                        self._execute(cursor, delete_app_state_sql, (app_name,))
+                        self._execute(
+                            cursor,
+                            insert_app_state_sql,
+                            (app_name, self._serialize_state(app_state), datetime.now(timezone.utc)),
+                        )
+                    if user_state is not None:
+                        self._execute(cursor, delete_user_state_sql, (app_name, user_id))
+                        self._execute(
+                            cursor,
+                            insert_user_state_sql,
+                            (app_name, user_id, self._serialize_state(user_state), datetime.now(timezone.utc)),
+                        )
+                    conn.commit()
+                else:
+                    with contextlib.suppress(Exception):
+                        conn.rollback()
             except Exception:
                 with contextlib.suppress(Exception):
                     conn.rollback()
@@ -1277,6 +1308,9 @@ class AdbcADKMemoryStore(BaseSyncADKMemoryStore["AdbcConfig"]):
 
     def create_tables(self) -> None:
         """Create tables if they don't exist."""
+        if not self._enabled:
+            return
+
         if not self.create_schema_enabled:
             self.reconcile_schema()
             return
@@ -1308,18 +1342,43 @@ class AdbcADKMemoryStore(BaseSyncADKMemoryStore["AdbcConfig"]):
         return self._delete_entries_older_than(days, app_name, scope)
 
     def _detect_dialect(self) -> str:
-        driver_name = self._config.connection_config.get("driver_name", "").lower()
-        if "postgres" in driver_name:
+        driver_name = self._config.connection_config.get("driver_name", "")
+        if (
+            isinstance(driver_name, str)
+            and driver_name
+            and driver_kind_from_driver_name(driver_name) is None
+            and not is_shared_object_driver(driver_name)
+        ):
+            logger.warning("Unknown ADBC driver: %s. Using generic SQL dialect.", driver_name)
+            return DIALECT_GENERIC
+
+        resolved = resolve_dialect_from_config(self._config.connection_config)
+        if resolved == "postgres":
             return DIALECT_POSTGRESQL
-        if "sqlite" in driver_name:
+        if resolved == "sqlite":
             return DIALECT_SQLITE
-        if "duckdb" in driver_name:
+        if resolved == "duckdb":
             return DIALECT_DUCKDB
-        if "snowflake" in driver_name:
+        if resolved == "snowflake":
             return DIALECT_SNOWFLAKE
         logger.warning("Unknown ADBC driver: %s. Using generic SQL dialect.", driver_name)
         return DIALECT_GENERIC
 
+    def _format_sql(self, sql: str) -> str:
+        """Return SQL with dialect-appropriate positional placeholders."""
+        if self._dialect != DIALECT_POSTGRESQL:
+            return sql
+        return _postgresql_store_sql(sql)
+
+    def _execute(self, cursor: Any, sql: str, params: "tuple[Any, ...] | list[Any]") -> Any:
+        """Execute parameterized SQL using the current ADBC dialect's placeholder style."""
+        return cursor.execute(self._format_sql(sql), params)
+
+    def _json_placeholder(self) -> str:
+        if self._dialect == DIALECT_POSTGRESQL:
+            return "?::jsonb"
+        return "?"
+
     def _serialize_json_field(self, value: Any) -> "str | None":
         if value is None:
             return None
@@ -1484,6 +1543,7 @@ class AdbcADKMemoryStore(BaseSyncADKMemoryStore["AdbcConfig"]):
 
         inserted_count = 0
         use_returning = self._dialect in {DIALECT_SQLITE, DIALECT_POSTGRESQL, DIALECT_DUCKDB}
+        json_placeholder = self._json_placeholder()
 
         if self._owner_id_column_name:
             if use_returning:
@@ -1493,7 +1553,7 @@ class AdbcADKMemoryStore(BaseSyncADKMemoryStore["AdbcConfig"]):
                     {self._owner_id_column_name}, timestamp, content_json, content_text,
                     metadata_json, inserted_at
                 ) VALUES (
-                    ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?
+                    ?, ?, ?, ?, ?, ?, ?, ?, ?, {json_placeholder}, ?, {json_placeholder}, ?
                 ) ON CONFLICT(event_id) DO NOTHING RETURNING 1
                 """
             else:
@@ -1503,7 +1563,7 @@ class AdbcADKMemoryStore(BaseSyncADKMemoryStore["AdbcConfig"]):
                     {self._owner_id_column_name}, timestamp, content_json, content_text,
                     metadata_json, inserted_at
                 ) VALUES (
-                    ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?
+                    ?, ?, ?, ?, ?, ?, ?, ?, ?, {json_placeholder}, ?, {json_placeholder}, ?
                 )
                 """
         elif use_returning:
@@ -1512,7 +1572,7 @@ class AdbcADKMemoryStore(BaseSyncADKMemoryStore["AdbcConfig"]):
                     id, session_id, app_name, user_id, scope, event_id, author,
                     timestamp, content_json, content_text, metadata_json, inserted_at
                 ) VALUES (
-                    ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?
+                    ?, ?, ?, ?, ?, ?, ?, ?, {json_placeholder}, ?, {json_placeholder}, ?
                 ) ON CONFLICT(event_id) DO NOTHING RETURNING 1
                 """
         else:
@@ -1521,7 +1581,7 @@ class AdbcADKMemoryStore(BaseSyncADKMemoryStore["AdbcConfig"]):
                     id, session_id, app_name, user_id, scope, event_id, author,
                     timestamp, content_json, content_text, metadata_json, inserted_at
                 ) VALUES (
-                    ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?
+                    ?, ?, ?, ?, ?, ?, ?, ?, {json_placeholder}, ?, {json_placeholder}, ?
                 )
                 """
 
@@ -1530,7 +1590,7 @@ class AdbcADKMemoryStore(BaseSyncADKMemoryStore["AdbcConfig"]):
             try:
                 for entry in entries:
                     content_json = self._serialize_json_field(entry["content_json"])
-                    metadata_json = self._serialize_json_field(entry["metadata_json"])
+                    metadata_json = self._serialize_json_field(entry.get("metadata_json"))
                     params: tuple[Any, ...]
                     if self._owner_id_column_name:
                         params = (
@@ -1540,7 +1600,7 @@ class AdbcADKMemoryStore(BaseSyncADKMemoryStore["AdbcConfig"]):
                             entry["user_id"],
                             entry.get("scope", "user"),
                             entry["event_id"],
-                            entry["author"],
+                            entry.get("author"),
                             owner_id,
                             self._encode_timestamp(entry["timestamp"]),
                             content_json,
@@ -1556,7 +1616,7 @@ class AdbcADKMemoryStore(BaseSyncADKMemoryStore["AdbcConfig"]):
                             entry["user_id"],
                             entry.get("scope", "user"),
                             entry["event_id"],
-                            entry["author"],
+                            entry.get("author"),
                             self._encode_timestamp(entry["timestamp"]),
                             content_json,
                             entry["content_text"],
@@ -1564,12 +1624,12 @@ class AdbcADKMemoryStore(BaseSyncADKMemoryStore["AdbcConfig"]):
                             self._encode_timestamp(entry["inserted_at"]),
                         )
                     if use_returning:
-                        cursor.execute(sql, params)
+                        self._execute(cursor, sql, params)
                         if cursor.fetchone():
                             inserted_count += 1
                     else:
                         try:
-                            cursor.execute(sql, params)
+                            self._execute(cursor, sql, params)
                             inserted_count += 1
                         except Exception as exc:
                             exc_str = str(exc).lower()
@@ -1594,6 +1654,9 @@ class AdbcADKMemoryStore(BaseSyncADKMemoryStore["AdbcConfig"]):
             msg = "Memory store is disabled"
             raise RuntimeError(msg)
 
+        if not query or not query.strip():
+            return []
+
         if self._use_fts:
             logger.warning("ADBC memory store does not support FTS, falling back to simple search")
 
@@ -1615,7 +1678,7 @@ class AdbcADKMemoryStore(BaseSyncADKMemoryStore["AdbcConfig"]):
             with self._config.provide_connection() as conn:
                 cursor = conn.cursor()
                 try:
-                    cursor.execute(sql, (*scope_params, pattern, effective_limit))
+                    self._execute(cursor, sql, (*scope_params, pattern, effective_limit))
                     rows = cursor.fetchall()
                 finally:
                     cursor.close()
@@ -1628,6 +1691,10 @@ class AdbcADKMemoryStore(BaseSyncADKMemoryStore["AdbcConfig"]):
         return self._rows_to_records(rows)
 
     def _delete_entries_by_session(self, session_id: str) -> int:
+        if not self._enabled:
+            msg = "Memory store is disabled"
+            raise RuntimeError(msg)
+
         use_returning = self._dialect in {DIALECT_SQLITE, DIALECT_POSTGRESQL, DIALECT_DUCKDB}
         if use_returning:
             sql = f"DELETE FROM {self._memory_table} WHERE session_id = ? RETURNING 1"
@@ -1636,7 +1703,7 @@ class AdbcADKMemoryStore(BaseSyncADKMemoryStore["AdbcConfig"]):
         with self._config.provide_connection() as conn:
             cursor = conn.cursor()
             try:
-                cursor.execute(sql, (session_id,))
+                self._execute(cursor, sql, (session_id,))
                 if use_returning:
                     deleted_rows = cursor.fetchall()
                     conn.commit()
@@ -1647,6 +1714,10 @@ class AdbcADKMemoryStore(BaseSyncADKMemoryStore["AdbcConfig"]):
                 cursor.close()
 
     def _delete_entries_older_than(self, days: int, app_name: "str | None" = None, scope: "str | None" = None) -> int:
+        if not self._enabled:
+            msg = "Memory store is disabled"
+            raise RuntimeError(msg)
+
         cutoff = self._encode_timestamp(datetime.now(timezone.utc) - timedelta(days=days))
         use_returning = self._dialect in {DIALECT_SQLITE, DIALECT_POSTGRESQL, DIALECT_DUCKDB}
         clauses = ["inserted_at < ?"]
@@ -1665,7 +1736,7 @@ class AdbcADKMemoryStore(BaseSyncADKMemoryStore["AdbcConfig"]):
         with self._config.provide_connection() as conn:
             cursor = conn.cursor()
             try:
-                cursor.execute(sql, tuple(params))
+                self._execute(cursor, sql, tuple(params))
                 if use_returning:
                     deleted_rows = cursor.fetchall()
                     conn.commit()

sqlspec/adapters/adbc/config.py

@@ -76,8 +76,15 @@ class AdbcConnectionParams(TypedDict):
     role: NotRequired[str]
     authorization_header: NotRequired[str]
     grpc_options: NotRequired[dict[str, Any]]
+    grpc_max_message_size: NotRequired[int]
+    fetch_timeout: NotRequired[float]
     gizmosql_backend: NotRequired[str]
     tls_skip_verify: NotRequired[bool]
+    tls_root_certs: NotRequired[str]
+    mtls_cert_chain: NotRequired[str]
+    mtls_private_key: NotRequired[str]
+    with_cookie_middleware: NotRequired[bool]
+    headers: NotRequired[dict[str, str]]
     extra: NotRequired[dict[str, Any]]
 
 

sqlspec/adapters/adbc/core.py

@@ -11,6 +11,7 @@ import sqlglot
 from sqlglot import exp
 from sqlglot.errors import ParseError
 
+import sqlspec.adapters.adbc._typing as _adbc_typing
 from sqlspec.adapters.adbc.type_converter import get_adbc_type_converter
 from sqlspec.core import (
     DriverParameterProfile,
@@ -181,6 +182,14 @@ _DRIVER_MANAGER_CONNECT: Final[str] = "adbc_driver_manager.dbapi.connect"
 _FLIGHTSQL_DB_KWARGS_FIELDS: "tuple[str, ...]" = ("username", "password")
 _FLIGHTSQL_TLS_SKIP_VERIFY_KEY: Final[str] = "adbc.flight.sql.client_option.tls_skip_verify"
 _FLIGHTSQL_AUTHORIZATION_HEADER_KEY: Final[str] = "adbc.flight.sql.authorization_header"
+_FLIGHTSQL_WITH_MAX_MSG_SIZE_KEY: Final[str] = "adbc.flight.sql.client_option.with_max_msg_size"
+_FLIGHTSQL_TIMEOUT_QUERY_KEY: Final[str] = "adbc.flight.sql.rpc.timeout_seconds.query"
+_FLIGHTSQL_TIMEOUT_FETCH_KEY: Final[str] = "adbc.flight.sql.rpc.timeout_seconds.fetch"
+_FLIGHTSQL_TLS_ROOT_CERTS_KEY: Final[str] = "adbc.flight.sql.client_option.tls_root_certs"
+_FLIGHTSQL_MTLS_CERT_CHAIN_KEY: Final[str] = "adbc.flight.sql.client_option.mtls_cert_chain"
+_FLIGHTSQL_MTLS_PRIVATE_KEY_KEY: Final[str] = "adbc.flight.sql.client_option.mtls_private_key"
+_FLIGHTSQL_WITH_COOKIE_MIDDLEWARE_KEY: Final[str] = "adbc.flight.sql.rpc.with_cookie_middleware"
+_FLIGHTSQL_RPC_CALL_HEADER_PREFIX: Final[str] = "adbc.flight.sql.rpc.call_header."
 _SQLSTATE_CLASS_CODE_LEN = 2
 _SQLSTATE_DESCRIPTIONS: dict[str, str] = {
     "23": "integrity constraint violation",
@@ -685,7 +694,6 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpec
     sqlstate = sqlstate_attr if sqlstate_attr is not None else None
 
     if sqlstate:
-        # Use centralized SQLSTATE mapping for specific codes
         if sqlstate == "23505":
             return _create_adbc_error(error, UniqueViolationError, "unique constraint violation")
         if sqlstate == "23503":
@@ -695,7 +703,6 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpec
         if sqlstate == "23514":
             return _create_adbc_error(error, CheckViolationError, "check constraint violation")
 
-        # Deadlock and serialization errors
         if sqlstate == "40P01":
             return _create_adbc_error(error, DeadlockError, "deadlock detected")
         if sqlstate == "40001":
@@ -705,25 +712,20 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpec
             termination_class = _classify_timeout_or_cancellation(str(error)) or OperationalError
             return _create_adbc_error(error, termination_class, "query terminated")
 
-        # Permission errors
         if sqlstate == "42501":
             return _create_adbc_error(error, PermissionDeniedError, "insufficient privilege")
         if sqlstate == "28000":
             return _create_adbc_error(error, PermissionDeniedError, "invalid authorization")
 
-        # Use centralized mapping for SQLSTATE class prefixes
         exc_class = map_sqlstate_to_exception(sqlstate)
         if exc_class is not None and exc_class is not SQLSpecError:
             description = _get_sqlstate_description(sqlstate)
             return _create_adbc_error(error, exc_class, description)
 
-        # Fallback for unmapped SQLSTATE codes
         return _create_adbc_error(error, SQLSpecError, "database error")
 
-    # Message-based fallback when no SQLSTATE is available
     error_msg = str(error).lower()
 
-    # Constraint violations
     if "unique" in error_msg or "duplicate" in error_msg:
         return _create_adbc_error(error, UniqueViolationError, "unique constraint violation")
     if "foreign key" in error_msg:
@@ -735,7 +737,6 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpec
     if "constraint" in error_msg:
         return _create_adbc_error(error, IntegrityError, "integrity constraint violation")
 
-    # Deadlock/lock patterns
     if "deadlock" in error_msg:
         return _create_adbc_error(error, DeadlockError, "deadlock detected")
     if "serialization" in error_msg or "concurrent update" in error_msg:
@@ -744,15 +745,12 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpec
     if message_class := _classify_timeout_or_cancellation(error_msg):
         return _create_adbc_error(error, message_class, "query terminated")
 
-    # Permission patterns
     if "permission" in error_msg or "denied" in error_msg or "unauthorized" in error_msg:
         return _create_adbc_error(error, PermissionDeniedError, "permission denied")
 
-    # Syntax errors
     if "syntax" in error_msg:
         return _create_adbc_error(error, SQLParsingError, "SQL parsing error")
 
-    # Connection errors
     if "connection" in error_msg or "connect" in error_msg:
         return _create_adbc_error(error, DatabaseConnectionError, "connection error")
 
@@ -1154,7 +1152,7 @@ def _rewrite_postgres_uuid_placeholders(
 
     for expression in expressions:
         expression.transform(wrap_parameter, copy=False)
-    rewritten = "; ".join(expression.sql(dialect=dialect) for expression in expressions)
+    rewritten = "; ".join(expression.sql(dialect=dialect, copy=False) for expression in expressions)
     return rewritten, effective_scalars, effective_arrays
 
 
@@ -1388,10 +1386,10 @@ def _prepare_batch_with_casts(
 
 def _resolve_flightsql_db_kwargs_keys() -> "tuple[str, str]":
     try:
-        from sqlspec.adapters.adbc._typing import AdbcFlightSqlDatabaseOptions as DatabaseOptions
+        database_options = _adbc_typing.AdbcFlightSqlDatabaseOptions
     except ImportError:
         return _FLIGHTSQL_TLS_SKIP_VERIFY_KEY, _FLIGHTSQL_AUTHORIZATION_HEADER_KEY
-    return DatabaseOptions.TLS_SKIP_VERIFY.value, DatabaseOptions.AUTHORIZATION_HEADER.value
+    return database_options.TLS_SKIP_VERIFY.value, database_options.AUTHORIZATION_HEADER.value
 
 
 def _lift_flightsql_db_kwargs(config: "dict[str, Any]") -> None:
@@ -1412,6 +1410,36 @@ def _lift_flightsql_db_kwargs(config: "dict[str, Any]") -> None:
     if "authorization_header" in config:
         db_kwargs_dict.setdefault(authorization_header_key, config.pop("authorization_header"))
 
+    if "grpc_max_message_size" in config:
+        db_kwargs_dict.setdefault(_FLIGHTSQL_WITH_MAX_MSG_SIZE_KEY, str(config.pop("grpc_max_message_size")))
+
+    if "query_timeout" in config:
+        db_kwargs_dict.setdefault(_FLIGHTSQL_TIMEOUT_QUERY_KEY, str(config.pop("query_timeout")))
+
+    if "fetch_timeout" in config:
+        db_kwargs_dict.setdefault(_FLIGHTSQL_TIMEOUT_FETCH_KEY, str(config.pop("fetch_timeout")))
+
+    if "tls_root_certs" in config:
+        db_kwargs_dict.setdefault(_FLIGHTSQL_TLS_ROOT_CERTS_KEY, str(config.pop("tls_root_certs")))
+
+    if "mtls_cert_chain" in config:
+        db_kwargs_dict.setdefault(_FLIGHTSQL_MTLS_CERT_CHAIN_KEY, str(config.pop("mtls_cert_chain")))
+
+    if "mtls_private_key" in config:
+        db_kwargs_dict.setdefault(_FLIGHTSQL_MTLS_PRIVATE_KEY_KEY, str(config.pop("mtls_private_key")))
+
+    if "with_cookie_middleware" in config:
+        val = config.pop("with_cookie_middleware")
+        db_kwargs_dict.setdefault(
+            _FLIGHTSQL_WITH_COOKIE_MIDDLEWARE_KEY, str(val).lower() if isinstance(val, bool) else str(val)
+        )
+
+    if "headers" in config:
+        headers = config.pop("headers")
+        if isinstance(headers, dict):
+            for hdr_key, hdr_val in headers.items():
+                db_kwargs_dict.setdefault(f"{_FLIGHTSQL_RPC_CALL_HEADER_PREFIX}{hdr_key}", str(hdr_val))
+
     config.pop("gizmosql_backend", None)
     if db_kwargs_dict:
         config["db_kwargs"] = db_kwargs_dict

sqlspec/adapters/adbc/driver.py

@@ -257,7 +257,6 @@ class AdbcDriver(SyncDriverAdapterBase):
 
         try:
             if not prepared_parameters:
-                cursor._rowcount = 0  # pyright: ignore[reportPrivateUsage]
                 row_count = 0
             elif isinstance(prepared_parameters, (list, tuple)) and prepared_parameters:
                 parameter_count = len(prepared_parameters)
@@ -496,7 +495,8 @@ class AdbcDriver(SyncDriverAdapterBase):
                 cursor = cursor_manager.__enter__()
                 try:
                     sql, driver_params = self._compiled_sql(prepared_statement, config)
-                    cursor.execute(sql, driver_params or ())
+                    execute_parameters = normalize_postgres_empty_parameters(self._dialect_name, driver_params)
+                    cursor.execute(sql, parameters=execute_parameters)
                     fetch_record_batch = getattr(cursor, "fetch_record_batch", None)
                     if fetch_record_batch is None:
                         arrow_table = cursor.fetch_arrow_table()
@@ -512,6 +512,7 @@ class AdbcDriver(SyncDriverAdapterBase):
                     else:
                         reader = fetch_record_batch()
                 except Exception:
+                    handle_postgres_rollback(self._dialect_name, cursor, logger)
                     cursor_manager.__exit__(None, None, None)
                     cursor = None
                     raise
@@ -537,8 +538,13 @@ class AdbcDriver(SyncDriverAdapterBase):
                 raise DatabaseConnectionError(msg)
 
             sql, driver_params = self._compiled_sql(prepared_statement, config)
-            cursor.execute(sql, driver_params or ())
-            arrow_table = cursor.fetch_arrow_table()
+            execute_parameters = normalize_postgres_empty_parameters(self._dialect_name, driver_params)
+            try:
+                cursor.execute(sql, parameters=execute_parameters)
+                arrow_table = cursor.fetch_arrow_table()
+            except Exception:
+                handle_postgres_rollback(self._dialect_name, cursor, logger)
+                raise
 
             arrow_result = build_arrow_result_from_table(
                 prepared_statement,

sqlspec/adapters/aiomysql/_typing.py

@@ -39,6 +39,8 @@ if TYPE_CHECKING:
 
         def close(self) -> object: ...
 
+        def get_transaction_status(self) -> bool: ...
+
     class AiomysqlModuleProtocol(Protocol):
         async def create_pool(self, **kwargs: Any) -> "AiomysqlPool": ...
 

sqlspec/adapters/aiomysql/adk/store.py

@@ -579,12 +579,12 @@ class AiomysqlADKMemoryStore(BaseAsyncADKMemoryStore["AiomysqlConfig"]):
                             entry["user_id"],
                             entry.get("scope", "user"),
                             entry["event_id"],
-                            entry["author"],
+                            entry.get("author"),
                             owner_id,
                             entry["timestamp"],
                             to_json(entry["content_json"]),
                             entry["content_text"],
-                            to_json(entry["metadata_json"]),
+                            to_json(entry["metadata_json"]) if entry.get("metadata_json") is not None else None,
                             entry["inserted_at"],
                         )
                     else:
@@ -595,11 +595,11 @@ class AiomysqlADKMemoryStore(BaseAsyncADKMemoryStore["AiomysqlConfig"]):
                             entry["user_id"],
                             entry.get("scope", "user"),
                             entry["event_id"],
-                            entry["author"],
+                            entry.get("author"),
                             entry["timestamp"],
                             to_json(entry["content_json"]),
                             entry["content_text"],
-                            to_json(entry["metadata_json"]),
+                            to_json(entry["metadata_json"]) if entry.get("metadata_json") is not None else None,
                             entry["inserted_at"],
                         )
                     await cursor.execute(sql, params)
@@ -655,6 +655,9 @@ class AiomysqlADKMemoryStore(BaseAsyncADKMemoryStore["AiomysqlConfig"]):
         records: list[StoredMemory] = []
         for row in rows:
             rec = cast("StoredMemory", dict(zip(columns, row, strict=False)))
+            rec["content_json"] = _json_dict(rec.get("content_json"))
+            metadata_val = rec.get("metadata_json")
+            rec["metadata_json"] = _json_dict(metadata_val) if metadata_val is not None else None
             rec["embedding"] = None
             records.append(rec)
         return records