sqlspec 0.25.0__py3-none-any.whl → 0.26.0__py3-none-any.whl

sqlspec/adapters/adbc/data_dictionary.py

@@ -0,0 +1,290 @@
+"""ADBC multi-dialect data dictionary for metadata queries."""
+
+import re
+from typing import TYPE_CHECKING, Optional, cast
+
+from sqlspec.driver import SyncDataDictionaryBase, SyncDriverAdapterBase, VersionInfo
+from sqlspec.utils.logging import get_logger
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from sqlspec.adapters.adbc.driver import AdbcDriver
+
+logger = get_logger("adapters.adbc.data_dictionary")
+
+POSTGRES_VERSION_PATTERN = re.compile(r"PostgreSQL (\d+)\.(\d+)(?:\.(\d+))?")
+SQLITE_VERSION_PATTERN = re.compile(r"(\d+)\.(\d+)\.(\d+)")
+DUCKDB_VERSION_PATTERN = re.compile(r"v?(\d+)\.(\d+)\.(\d+)")
+MYSQL_VERSION_PATTERN = re.compile(r"(\d+)\.(\d+)\.(\d+)")
+
+__all__ = ("AdbcDataDictionary",)
+
+
+class AdbcDataDictionary(SyncDataDictionaryBase):
+    """ADBC multi-dialect data dictionary.
+
+    Delegates to appropriate dialect-specific logic based on the driver's dialect.
+    """
+
+    def _get_dialect(self, driver: SyncDriverAdapterBase) -> str:
+        """Get dialect from ADBC driver.
+
+        Args:
+            driver: ADBC driver instance
+
+        Returns:
+            Dialect name
+        """
+        return str(cast("AdbcDriver", driver).dialect)
+
+    def get_version(self, driver: SyncDriverAdapterBase) -> "Optional[VersionInfo]":
+        """Get database version information based on detected dialect.
+
+        Args:
+            driver: ADBC driver instance
+
+        Returns:
+            Database version information or None if detection fails
+        """
+        dialect = self._get_dialect(driver)
+        adbc_driver = cast("AdbcDriver", driver)
+
+        try:
+            if dialect == "postgres":
+                version_str = adbc_driver.select_value("SELECT version()")
+                if version_str:
+                    match = POSTGRES_VERSION_PATTERN.search(str(version_str))
+                    if match:
+                        major = int(match.group(1))
+                        minor = int(match.group(2))
+                        patch = int(match.group(3)) if match.group(3) else 0
+                        return VersionInfo(major, minor, patch)
+
+            elif dialect == "sqlite":
+                version_str = adbc_driver.select_value("SELECT sqlite_version()")
+                if version_str:
+                    match = SQLITE_VERSION_PATTERN.match(str(version_str))
+                    if match:
+                        major, minor, patch = map(int, match.groups())
+                        return VersionInfo(major, minor, patch)
+
+            elif dialect == "duckdb":
+                version_str = adbc_driver.select_value("SELECT version()")
+                if version_str:
+                    match = DUCKDB_VERSION_PATTERN.search(str(version_str))
+                    if match:
+                        major, minor, patch = map(int, match.groups())
+                        return VersionInfo(major, minor, patch)
+
+            elif dialect == "mysql":
+                version_str = adbc_driver.select_value("SELECT VERSION()")
+                if version_str:
+                    match = MYSQL_VERSION_PATTERN.search(str(version_str))
+                    if match:
+                        major, minor, patch = map(int, match.groups())
+                        return VersionInfo(major, minor, patch)
+
+            elif dialect == "bigquery":
+                return VersionInfo(1, 0, 0)
+
+        except Exception:
+            logger.warning("Failed to get %s version", dialect)
+
+        return None
+
+    def get_feature_flag(self, driver: SyncDriverAdapterBase, feature: str) -> bool:
+        """Check if database supports a specific feature based on detected dialect.
+
+        Args:
+            driver: ADBC driver instance
+            feature: Feature name to check
+
+        Returns:
+            True if feature is supported, False otherwise
+        """
+        dialect = self._get_dialect(driver)
+        version_info = self.get_version(driver)
+
+        if dialect == "postgres":
+            feature_checks: dict[str, Callable[..., bool]] = {
+                "supports_json": lambda v: v and v >= VersionInfo(9, 2, 0),
+                "supports_jsonb": lambda v: v and v >= VersionInfo(9, 4, 0),
+                "supports_uuid": lambda _: True,
+                "supports_arrays": lambda _: True,
+                "supports_returning": lambda v: v and v >= VersionInfo(8, 2, 0),
+                "supports_upsert": lambda v: v and v >= VersionInfo(9, 5, 0),
+                "supports_window_functions": lambda v: v and v >= VersionInfo(8, 4, 0),
+                "supports_cte": lambda v: v and v >= VersionInfo(8, 4, 0),
+                "supports_transactions": lambda _: True,
+                "supports_prepared_statements": lambda _: True,
+                "supports_schemas": lambda _: True,
+            }
+        elif dialect == "sqlite":
+            feature_checks = {
+                "supports_json": lambda v: v and v >= VersionInfo(3, 38, 0),
+                "supports_returning": lambda v: v and v >= VersionInfo(3, 35, 0),
+                "supports_upsert": lambda v: v and v >= VersionInfo(3, 24, 0),
+                "supports_window_functions": lambda v: v and v >= VersionInfo(3, 25, 0),
+                "supports_cte": lambda v: v and v >= VersionInfo(3, 8, 3),
+                "supports_transactions": lambda _: True,
+                "supports_prepared_statements": lambda _: True,
+                "supports_schemas": lambda _: False,
+                "supports_arrays": lambda _: False,
+                "supports_uuid": lambda _: False,
+            }
+        elif dialect == "duckdb":
+            feature_checks = {
+                "supports_json": lambda _: True,
+                "supports_arrays": lambda _: True,
+                "supports_uuid": lambda _: True,
+                "supports_returning": lambda v: v and v >= VersionInfo(0, 8, 0),
+                "supports_upsert": lambda v: v and v >= VersionInfo(0, 8, 0),
+                "supports_window_functions": lambda _: True,
+                "supports_cte": lambda _: True,
+                "supports_transactions": lambda _: True,
+                "supports_prepared_statements": lambda _: True,
+                "supports_schemas": lambda _: True,
+            }
+        elif dialect == "mysql":
+            feature_checks = {
+                "supports_json": lambda v: v and v >= VersionInfo(5, 7, 8),
+                "supports_cte": lambda v: v and v >= VersionInfo(8, 0, 1),
+                "supports_returning": lambda _: False,
+                "supports_upsert": lambda _: True,
+                "supports_window_functions": lambda v: v and v >= VersionInfo(8, 0, 2),
+                "supports_transactions": lambda _: True,
+                "supports_prepared_statements": lambda _: True,
+                "supports_schemas": lambda _: True,
+                "supports_uuid": lambda _: False,
+                "supports_arrays": lambda _: False,
+            }
+        elif dialect == "bigquery":
+            feature_checks = {
+                "supports_json": lambda _: True,
+                "supports_arrays": lambda _: True,
+                "supports_structs": lambda _: True,
+                "supports_returning": lambda _: False,
+                "supports_upsert": lambda _: True,
+                "supports_window_functions": lambda _: True,
+                "supports_cte": lambda _: True,
+                "supports_transactions": lambda _: False,
+                "supports_prepared_statements": lambda _: True,
+                "supports_schemas": lambda _: True,
+                "supports_uuid": lambda _: False,
+            }
+        else:
+            feature_checks = {
+                "supports_transactions": lambda _: True,
+                "supports_prepared_statements": lambda _: True,
+                "supports_window_functions": lambda _: True,
+                "supports_cte": lambda _: True,
+            }
+
+        if feature in feature_checks:
+            return bool(feature_checks[feature](version_info))
+
+        return False
+
+    def get_optimal_type(self, driver: SyncDriverAdapterBase, type_category: str) -> str:
+        """Get optimal database type for a category based on detected dialect.
+
+        Args:
+            driver: ADBC driver instance
+            type_category: Type category
+
+        Returns:
+            Database-specific type name
+        """
+        dialect = self._get_dialect(driver)
+        version_info = self.get_version(driver)
+
+        if dialect == "postgres":
+            if type_category == "json":
+                if version_info and version_info >= VersionInfo(9, 4, 0):
+                    return "JSONB"
+                if version_info and version_info >= VersionInfo(9, 2, 0):
+                    return "JSON"
+                return "TEXT"
+            type_map = {
+                "uuid": "UUID",
+                "boolean": "BOOLEAN",
+                "timestamp": "TIMESTAMP WITH TIME ZONE",
+                "text": "TEXT",
+                "blob": "BYTEA",
+                "array": "ARRAY",
+            }
+
+        elif dialect == "sqlite":
+            if type_category == "json":
+                if version_info and version_info >= VersionInfo(3, 38, 0):
+                    return "JSON"
+                return "TEXT"
+            type_map = {"uuid": "TEXT", "boolean": "INTEGER", "timestamp": "TIMESTAMP", "text": "TEXT", "blob": "BLOB"}
+
+        elif dialect == "duckdb":
+            type_map = {
+                "json": "JSON",
+                "uuid": "UUID",
+                "boolean": "BOOLEAN",
+                "timestamp": "TIMESTAMP",
+                "text": "TEXT",
+                "blob": "BLOB",
+                "array": "LIST",
+            }
+
+        elif dialect == "mysql":
+            if type_category == "json":
+                if version_info and version_info >= VersionInfo(5, 7, 8):
+                    return "JSON"
+                return "TEXT"
+            type_map = {
+                "uuid": "VARCHAR(36)",
+                "boolean": "TINYINT(1)",
+                "timestamp": "TIMESTAMP",
+                "text": "TEXT",
+                "blob": "BLOB",
+            }
+
+        elif dialect == "bigquery":
+            type_map = {
+                "json": "JSON",
+                "uuid": "STRING",
+                "boolean": "BOOL",
+                "timestamp": "TIMESTAMP",
+                "text": "STRING",
+                "blob": "BYTES",
+                "array": "ARRAY",
+            }
+        else:
+            type_map = {
+                "json": "TEXT",
+                "uuid": "VARCHAR(36)",
+                "boolean": "INTEGER",
+                "timestamp": "TIMESTAMP",
+                "text": "TEXT",
+                "blob": "BLOB",
+            }
+
+        return type_map.get(type_category, "TEXT")
+
+    def list_available_features(self) -> "list[str]":
+        """List available feature flags across all supported dialects.
+
+        Returns:
+            List of supported feature names
+        """
+        return [
+            "supports_json",
+            "supports_jsonb",
+            "supports_uuid",
+            "supports_arrays",
+            "supports_structs",
+            "supports_returning",
+            "supports_upsert",
+            "supports_window_functions",
+            "supports_cte",
+            "supports_transactions",
+            "supports_prepared_statements",
+            "supports_schemas",
+        ]

sqlspec/adapters/adbc/driver.py

@@ -12,13 +12,15 @@ from typing import TYPE_CHECKING, Any, Optional, cast
 from adbc_driver_manager.dbapi import DatabaseError, IntegrityError, OperationalError, ProgrammingError
 from sqlglot import exp
 
+from sqlspec.adapters.adbc.data_dictionary import AdbcDataDictionary
+from sqlspec.adapters.adbc.type_converter import ADBCTypeConverter
 from sqlspec.core.cache import get_cache_config
 from sqlspec.core.parameters import ParameterStyle, ParameterStyleConfig
 from sqlspec.core.statement import SQL, StatementConfig
 from sqlspec.driver import SyncDriverAdapterBase
 from sqlspec.exceptions import MissingDependencyError, SQLParsingError, SQLSpecError
+from sqlspec.typing import Empty
 from sqlspec.utils.logging import get_logger
-from sqlspec.utils.serializers import to_json
 
 if TYPE_CHECKING:
     from contextlib import AbstractContextManager
@@ -28,6 +30,7 @@ if TYPE_CHECKING:
     from sqlspec.adapters.adbc._types import AdbcConnection
     from sqlspec.core.result import SQLResult
     from sqlspec.driver import ExecutionResult
+    from sqlspec.driver._sync import SyncDataDictionaryBase
 
 __all__ = ("AdbcCursor", "AdbcDriver", "AdbcExceptionHandler", "get_adbc_statement_config")
 
@@ -293,7 +296,7 @@ def _convert_array_for_postgres_adbc(value: Any) -> Any:
 
 
 def get_type_coercion_map(dialect: str) -> "dict[type, Any]":
-    """Get type coercion map for Arrow type handling.
+    """Get type coercion map for Arrow type handling with dialect-aware type conversion.
 
     Args:
         dialect: Database dialect name
@@ -301,7 +304,9 @@ def get_type_coercion_map(dialect: str) -> "dict[type, Any]":
     Returns:
         Mapping of Python types to conversion functions
     """
-    type_map = {
+    tc = ADBCTypeConverter(dialect)
+
+    return {
         datetime.datetime: lambda x: x,
         datetime.date: lambda x: x,
         datetime.time: lambda x: x,
@@ -309,18 +314,13 @@ def get_type_coercion_map(dialect: str) -> "dict[type, Any]":
         bool: lambda x: x,
         int: lambda x: x,
         float: lambda x: x,
-        str: lambda x: x,
+        str: tc.convert_if_detected,
         bytes: lambda x: x,
         tuple: _convert_array_for_postgres_adbc,
         list: _convert_array_for_postgres_adbc,
         dict: lambda x: x,
     }
 
-    if dialect in {"postgres", "postgresql"}:
-        type_map[dict] = lambda x: to_json(x) if x is not None else None
-
-    return type_map
-
 
 class AdbcCursor:
     """Context manager for cursor management."""
@@ -335,8 +335,7 @@ class AdbcCursor:
         self.cursor = self.connection.cursor()
         return self.cursor
 
-    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
-        _ = (exc_type, exc_val, exc_tb)
+    def __exit__(self, *_: Any) -> None:
         if self.cursor is not None:
             with contextlib.suppress(Exception):
                 self.cursor.close()  # type: ignore[no-untyped-call]
@@ -394,7 +393,7 @@ class AdbcDriver(SyncDriverAdapterBase):
     database dialects, parameter style conversion, and transaction management.
     """
 
-    __slots__ = ("_detected_dialect", "dialect")
+    __slots__ = ("_data_dictionary", "_detected_dialect", "dialect")
 
     def __init__(
         self,
@@ -413,6 +412,7 @@ class AdbcDriver(SyncDriverAdapterBase):
 
         super().__init__(connection=connection, statement_config=statement_config, driver_features=driver_features)
         self.dialect = statement_config.dialect
+        self._data_dictionary: Optional[SyncDataDictionaryBase] = None
 
     @staticmethod
     def _ensure_pyarrow_installed() -> None:
@@ -475,6 +475,87 @@ class AdbcDriver(SyncDriverAdapterBase):
             return None
         return parameters
 
+    def prepare_driver_parameters(
+        self,
+        parameters: Any,
+        statement_config: "StatementConfig",
+        is_many: bool = False,
+        prepared_statement: Optional[Any] = None,
+    ) -> Any:
+        """Prepare parameters with cast-aware type coercion for ADBC.
+
+        For PostgreSQL, applies cast-aware parameter processing using metadata from the compiled statement.
+        This allows proper handling of JSONB casts and other type conversions.
+
+        Args:
+            parameters: Parameters in any format
+            statement_config: Statement configuration
+            is_many: Whether this is for execute_many operation
+            prepared_statement: Prepared statement containing the original SQL statement
+
+        Returns:
+            Parameters with cast-aware type coercion applied
+        """
+        if prepared_statement and self.dialect in {"postgres", "postgresql"} and not is_many:
+            parameter_casts = self._get_parameter_casts(prepared_statement)
+            postgres_compatible = self._handle_postgres_empty_parameters(parameters)
+            return self._prepare_parameters_with_casts(postgres_compatible, parameter_casts, statement_config)
+
+        return super().prepare_driver_parameters(parameters, statement_config, is_many, prepared_statement)
+
+    def _get_parameter_casts(self, statement: SQL) -> "dict[int, str]":
+        """Get parameter cast metadata from compiled statement.
+
+        Args:
+            statement: SQL statement with compiled metadata
+
+        Returns:
+            Dict mapping parameter positions to cast types
+        """
+
+        processed_state = statement.get_processed_state()
+        if processed_state is not Empty:
+            return processed_state.parameter_casts or {}
+        return {}
+
+    def _prepare_parameters_with_casts(
+        self, parameters: Any, parameter_casts: "dict[int, str]", statement_config: "StatementConfig"
+    ) -> Any:
+        """Prepare parameters with cast-aware type coercion.
+
+        Uses type coercion map for non-dict types and dialect-aware dict handling.
+
+        Args:
+            parameters: Parameter values (list, tuple, or scalar)
+            parameter_casts: Mapping of parameter positions to cast types
+            statement_config: Statement configuration for type coercion
+
+        Returns:
+            Parameters with cast-aware type coercion applied
+        """
+        from sqlspec._serialization import encode_json
+
+        if isinstance(parameters, (list, tuple)):
+            result: list[Any] = []
+            for idx, param in enumerate(parameters, start=1):  # pyright: ignore
+                cast_type = parameter_casts.get(idx, "").upper()
+                if cast_type in {"JSON", "JSONB", "TYPE.JSON", "TYPE.JSONB"}:
+                    if isinstance(param, dict):
+                        result.append(encode_json(param))
+                    else:
+                        result.append(param)
+                elif isinstance(param, dict):
+                    result.append(ADBCTypeConverter(self.dialect).convert_dict(param))  # type: ignore[arg-type]
+                else:
+                    if statement_config.parameter_config.type_coercion_map:
+                        for type_check, converter in statement_config.parameter_config.type_coercion_map.items():
+                            if type_check is not dict and isinstance(param, type_check):
+                                param = converter(param)
+                                break
+                    result.append(param)
+            return tuple(result) if isinstance(parameters, tuple) else result
+        return parameters
+
     def with_cursor(self, connection: "AdbcConnection") -> "AdbcCursor":
         """Create context manager for cursor.
 
@@ -519,17 +600,26 @@ class AdbcDriver(SyncDriverAdapterBase):
         """
         sql, prepared_parameters = self._get_compiled_sql(statement, self.statement_config)
 
+        parameter_casts = self._get_parameter_casts(statement)
+
         try:
             if not prepared_parameters:
                 cursor._rowcount = 0  # pyright: ignore[reportPrivateUsage]
                 row_count = 0
-            elif isinstance(prepared_parameters, list) and prepared_parameters:
+            elif isinstance(prepared_parameters, (list, tuple)) and prepared_parameters:
                 processed_params = []
                 for param_set in prepared_parameters:
                     postgres_compatible = self._handle_postgres_empty_parameters(param_set)
-                    formatted_params = self.prepare_driver_parameters(
-                        postgres_compatible, self.statement_config, is_many=False
-                    )
+
+                    if self.dialect in {"postgres", "postgresql"}:
+                        # For postgres, always use cast-aware parameter preparation
+                        formatted_params = self._prepare_parameters_with_casts(
+                            postgres_compatible, parameter_casts, self.statement_config
+                        )
+                    else:
+                        formatted_params = self.prepare_driver_parameters(
+                            postgres_compatible, self.statement_config, is_many=False
+                        )
                     processed_params.append(formatted_params)
 
                 cursor.executemany(sql, processed_params)
@@ -540,7 +630,6 @@ class AdbcDriver(SyncDriverAdapterBase):
 
         except Exception:
             self._handle_postgres_rollback(cursor)
-            logger.exception("Executemany failed")
             raise
 
         return self.create_execution_result(cursor, rowcount_override=row_count, is_many_result=True)
@@ -557,9 +646,18 @@ class AdbcDriver(SyncDriverAdapterBase):
         """
         sql, prepared_parameters = self._get_compiled_sql(statement, self.statement_config)
 
+        parameter_casts = self._get_parameter_casts(statement)
+
         try:
             postgres_compatible_params = self._handle_postgres_empty_parameters(prepared_parameters)
-            cursor.execute(sql, parameters=postgres_compatible_params)
+
+            if self.dialect in {"postgres", "postgresql"}:
+                formatted_params = self._prepare_parameters_with_casts(
+                    postgres_compatible_params, parameter_casts, self.statement_config
+                )
+                cursor.execute(sql, parameters=formatted_params)
+            else:
+                cursor.execute(sql, parameters=postgres_compatible_params)
 
         except Exception:
             self._handle_postgres_rollback(cursor)
@@ -655,3 +753,14 @@ class AdbcDriver(SyncDriverAdapterBase):
         except Exception as e:
             msg = f"Failed to commit transaction: {e}"
             raise SQLSpecError(msg) from e
+
+    @property
+    def data_dictionary(self) -> "SyncDataDictionaryBase":
+        """Get the data dictionary for this driver.
+
+        Returns:
+            Data dictionary instance for metadata queries
+        """
+        if self._data_dictionary is None:
+            self._data_dictionary = AdbcDataDictionary()
+        return self._data_dictionary