sqlspec 0.27.0__py3-none-any.whl → 0.28.0__py3-none-any.whl

sqlspec/adapters/asyncpg/data_dictionary.py

@@ -117,7 +117,7 @@ class PostgresAsyncDataDictionary(AsyncDataDictionaryBase):
     async def get_columns(
         self, driver: AsyncDriverAdapterBase, table: str, schema: "str | None" = None
     ) -> "list[dict[str, Any]]":
-        """Get column information for a table using information_schema.
+        """Get column information for a table using pg_catalog.
 
         Args:
             driver: AsyncPG driver instance
@@ -130,25 +130,32 @@ class PostgresAsyncDataDictionary(AsyncDataDictionaryBase):
                 - data_type: PostgreSQL data type
                 - is_nullable: Whether column allows NULL (YES/NO)
                 - column_default: Default value if any
+
+        Notes:
+            Uses pg_catalog instead of information_schema to avoid potential
+            issues with PostgreSQL 'name' type in some drivers.
         """
         asyncpg_driver = cast("AsyncpgDriver", driver)
 
-        if schema:
-            sql = f"""
-                SELECT column_name, data_type, is_nullable, column_default
-                FROM information_schema.columns
-                WHERE table_name = '{table}' AND table_schema = '{schema}'
-                ORDER BY ordinal_position
-            """
-        else:
-            sql = f"""
-                SELECT column_name, data_type, is_nullable, column_default
-                FROM information_schema.columns
-                WHERE table_name = '{table}' AND table_schema = 'public'
-                ORDER BY ordinal_position
-            """
-
-        result = await asyncpg_driver.execute(sql)
+        schema_name = schema or "public"
+        sql = """
+            SELECT
+                a.attname::text AS column_name,
+                pg_catalog.format_type(a.atttypid, a.atttypmod) AS data_type,
+                CASE WHEN a.attnotnull THEN 'NO' ELSE 'YES' END AS is_nullable,
+                pg_catalog.pg_get_expr(d.adbin, d.adrelid)::text AS column_default
+            FROM pg_catalog.pg_attribute a
+            JOIN pg_catalog.pg_class c ON a.attrelid = c.oid
+            JOIN pg_catalog.pg_namespace n ON c.relnamespace = n.oid
+            LEFT JOIN pg_catalog.pg_attrdef d ON a.attrelid = d.adrelid AND a.attnum = d.adnum
+            WHERE c.relname = $1
+                AND n.nspname = $2
+                AND a.attnum > 0
+                AND NOT a.attisdropped
+            ORDER BY a.attnum
+        """
+
+        result = await asyncpg_driver.execute(sql, (table, schema_name))
         return result.data or []
 
     def list_available_features(self) -> "list[str]":

sqlspec/adapters/bigquery/adk/store.py

@@ -8,7 +8,7 @@ from google.cloud.bigquery import QueryJobConfig, ScalarQueryParameter
 from sqlspec.extensions.adk import BaseAsyncADKStore, EventRecord, SessionRecord
 from sqlspec.utils.logging import get_logger
 from sqlspec.utils.serializers import from_json, to_json
-from sqlspec.utils.sync_tools import async_
+from sqlspec.utils.sync_tools import async_, run_
 
 if TYPE_CHECKING:
     from sqlspec.adapters.bigquery.config import BigQueryConfig
@@ -102,7 +102,7 @@ class BigQueryADKStore(BaseAsyncADKStore["BigQueryConfig"]):
             return f"`{self._dataset_id}.{table_name}`"
         return f"`{table_name}`"
 
-    def _get_create_sessions_table_sql(self) -> str:
+    async def _get_create_sessions_table_sql(self) -> str:
         """Get BigQuery CREATE TABLE SQL for sessions.
 
         Returns:
@@ -136,7 +136,7 @@ class BigQueryADKStore(BaseAsyncADKStore["BigQueryConfig"]):
         CLUSTER BY app_name, user_id
         """
 
-    def _get_create_events_table_sql(self) -> str:
+    async def _get_create_events_table_sql(self) -> str:
         """Get BigQuery CREATE TABLE SQL for events.
 
         Returns:
@@ -193,9 +193,9 @@ class BigQueryADKStore(BaseAsyncADKStore["BigQueryConfig"]):
 
     def _create_tables(self) -> None:
         """Synchronous implementation of create_tables."""
-        with self._config.provide_connection() as conn:
-            conn.query(self._get_create_sessions_table_sql()).result()
-            conn.query(self._get_create_events_table_sql()).result()
+        with self._config.provide_session() as driver:
+            driver.execute_script(run_(self._get_create_sessions_table_sql)())
+            driver.execute_script(run_(self._get_create_events_table_sql)())
         logger.debug("Created BigQuery ADK tables: %s, %s", self._session_table, self._events_table)
 
     async def create_tables(self) -> None:
@@ -351,20 +351,29 @@ class BigQueryADKStore(BaseAsyncADKStore["BigQueryConfig"]):
         """
         await async_(self._update_session_state)(session_id, state)
 
-    def _list_sessions(self, app_name: str, user_id: str) -> "list[SessionRecord]":
+    def _list_sessions(self, app_name: str, user_id: "str | None") -> "list[SessionRecord]":
         """Synchronous implementation of list_sessions."""
         table_name = self._get_full_table_name(self._session_table)
-        sql = f"""
-        SELECT id, app_name, user_id, JSON_VALUE(state) as state, create_time, update_time
-        FROM {table_name}
-        WHERE app_name = @app_name AND user_id = @user_id
-        ORDER BY update_time DESC
-        """
 
-        params = [
-            ScalarQueryParameter("app_name", "STRING", app_name),
-            ScalarQueryParameter("user_id", "STRING", user_id),
-        ]
+        if user_id is None:
+            sql = f"""
+            SELECT id, app_name, user_id, JSON_VALUE(state) as state, create_time, update_time
+            FROM {table_name}
+            WHERE app_name = @app_name
+            ORDER BY update_time DESC
+            """
+            params = [ScalarQueryParameter("app_name", "STRING", app_name)]
+        else:
+            sql = f"""
+            SELECT id, app_name, user_id, JSON_VALUE(state) as state, create_time, update_time
+            FROM {table_name}
+            WHERE app_name = @app_name AND user_id = @user_id
+            ORDER BY update_time DESC
+            """
+            params = [
+                ScalarQueryParameter("app_name", "STRING", app_name),
+                ScalarQueryParameter("user_id", "STRING", user_id),
+            ]
 
         with self._config.provide_connection() as conn:
             job_config = QueryJobConfig(query_parameters=params)
@@ -383,18 +392,18 @@ class BigQueryADKStore(BaseAsyncADKStore["BigQueryConfig"]):
                 for row in results
             ]
 
-    async def list_sessions(self, app_name: str, user_id: str) -> "list[SessionRecord]":
-        """List all sessions for a user in an app.
+    async def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[SessionRecord]":
+        """List sessions for an app, optionally filtered by user.
 
         Args:
             app_name: Application name.
-            user_id: User identifier.
+            user_id: User identifier. If None, lists all sessions for the app.
 
         Returns:
             List of session records ordered by update_time DESC.
 
         Notes:
-            Uses clustering on (app_name, user_id) for efficiency.
+            Uses clustering on (app_name, user_id) when user_id is provided for efficiency.
         """
         return await async_(self._list_sessions)(app_name, user_id)
 

sqlspec/adapters/bigquery/config.py

@@ -67,6 +67,17 @@ class BigQueryDriverFeatures(TypedDict):
     """BigQuery driver-specific features configuration.
 
     Only non-standard BigQuery client parameters that are SQLSpec-specific extensions.
+
+    Attributes:
+        connection_instance: Pre-existing BigQuery connection instance to use.
+        on_job_start: Callback invoked when a query job starts.
+        on_job_complete: Callback invoked when a query job completes.
+        on_connection_create: Callback invoked when connection is created.
+        json_serializer: Custom JSON serializer for dict/list parameter conversion.
+            Defaults to sqlspec.utils.serializers.to_json if not provided.
+        enable_uuid_conversion: Enable automatic UUID string conversion.
+            When True (default), UUID strings are automatically converted to UUID objects.
+            When False, UUID strings are treated as regular strings.
     """
 
     connection_instance: NotRequired["BigQueryConnection"]

sqlspec/adapters/bigquery/driver.py

@@ -33,8 +33,11 @@ if TYPE_CHECKING:
     from collections.abc import Callable
     from contextlib import AbstractContextManager
 
-    from sqlspec.core import SQL, SQLResult
+    from sqlspec.builder import QueryBuilder
+    from sqlspec.core import SQL, SQLResult, Statement, StatementFilter
+    from sqlspec.core.result import ArrowResult
     from sqlspec.driver import SyncDataDictionaryBase
+    from sqlspec.typing import StatementParameters
 
 logger = logging.getLogger(__name__)
 
@@ -758,3 +761,137 @@ class BigQueryDriver(SyncDriverAdapterBase):
 
             self._data_dictionary = BigQuerySyncDataDictionary()
         return self._data_dictionary
+
+    def _storage_api_available(self) -> bool:
+        """Check if BigQuery Storage API is available.
+
+        Returns:
+            True if Storage API is available and working, False otherwise
+        """
+        try:
+            from google.cloud import bigquery_storage_v1  # type: ignore[attr-defined]
+
+            # Try to create client (will fail if API not enabled or credentials missing)
+            _ = bigquery_storage_v1.BigQueryReadClient()
+        except ImportError:
+            # Package not installed
+            return False
+        except Exception:
+            # API not enabled or permissions issue
+            return False
+        else:
+            return True
+
+    def select_to_arrow(
+        self,
+        statement: "Statement | QueryBuilder",
+        /,
+        *parameters: "StatementParameters | StatementFilter",
+        statement_config: "StatementConfig | None" = None,
+        return_format: str = "table",
+        native_only: bool = False,
+        batch_size: int | None = None,
+        arrow_schema: Any = None,
+        **kwargs: Any,
+    ) -> "ArrowResult":
+        """Execute query and return results as Apache Arrow (BigQuery native with Storage API).
+
+        BigQuery provides native Arrow via Storage API (query_job.to_arrow()).
+        Requires google-cloud-bigquery-storage package and API enabled.
+        Falls back to dict conversion if Storage API not available.
+
+        Args:
+            statement: SQL statement, string, or QueryBuilder
+            *parameters: Query parameters or filters
+            statement_config: Optional statement configuration override
+            return_format: "table" for pyarrow.Table (default), "batch" for RecordBatch
+            native_only: If True, raise error if Storage API unavailable (default: False)
+            batch_size: Batch size hint (for future streaming implementation)
+            arrow_schema: Optional pyarrow.Schema for type casting
+            **kwargs: Additional keyword arguments
+
+        Returns:
+            ArrowResult with native Arrow data (if Storage API available) or converted data
+
+        Raises:
+            MissingDependencyError: If pyarrow not installed, or if Storage API not available and native_only=True
+            SQLExecutionError: If query execution fails
+
+        Example:
+            >>> # Will use native Arrow if Storage API available, otherwise converts
+            >>> result = driver.select_to_arrow(
+            ...     "SELECT * FROM dataset.users WHERE age > @age",
+            ...     {"age": 18},
+            ... )
+            >>> df = result.to_pandas()
+
+            >>> # Force native Arrow (raises if Storage API unavailable)
+            >>> result = driver.select_to_arrow(
+            ...     "SELECT * FROM dataset.users", native_only=True
+            ... )
+        """
+        from sqlspec.utils.module_loader import ensure_pyarrow
+
+        ensure_pyarrow()
+
+        # Check Storage API availability
+        if not self._storage_api_available():
+            if native_only:
+                from sqlspec.exceptions import MissingDependencyError
+
+                msg = (
+                    "BigQuery native Arrow requires Storage API.\n"
+                    "1. Install: pip install google-cloud-bigquery-storage\n"
+                    "2. Enable API: https://console.cloud.google.com/apis/library/bigquerystorage.googleapis.com\n"
+                    "3. Grant permissions: roles/bigquery.dataViewer"
+                )
+                raise MissingDependencyError(
+                    package="google-cloud-bigquery-storage", install_package="google-cloud-bigquery-storage"
+                ) from RuntimeError(msg)
+
+            # Fallback to conversion path
+            result: ArrowResult = super().select_to_arrow(
+                statement,
+                *parameters,
+                statement_config=statement_config,
+                return_format=return_format,
+                native_only=native_only,
+                batch_size=batch_size,
+                arrow_schema=arrow_schema,
+                **kwargs,
+            )
+            return result
+
+        # Use native path with Storage API
+        import pyarrow as pa
+
+        from sqlspec.core.result import create_arrow_result
+
+        # Prepare statement
+        config = statement_config or self.statement_config
+        prepared_statement = self.prepare_statement(statement, parameters, statement_config=config, kwargs=kwargs)
+
+        # Get compiled SQL and parameters
+        sql, driver_params = self._get_compiled_sql(prepared_statement, config)
+
+        # Execute query using existing _run_query_job method
+        with self.handle_database_exceptions():
+            query_job = self._run_query_job(sql, driver_params)
+            query_job.result()  # Wait for completion
+
+            # Native Arrow via Storage API
+            arrow_table = query_job.to_arrow()
+
+            # Apply schema casting if requested
+            if arrow_schema is not None:
+                arrow_table = arrow_table.cast(arrow_schema)
+
+            # Convert to batch if requested
+            if return_format == "batch":
+                batches = arrow_table.to_batches()
+                arrow_data: Any = batches[0] if batches else pa.RecordBatch.from_pydict({})
+            else:
+                arrow_data = arrow_table
+
+        # Create ArrowResult
+        return create_arrow_result(statement=prepared_statement, data=arrow_data, rows_affected=arrow_data.num_rows)

sqlspec/adapters/duckdb/adk/store.py

@@ -315,29 +315,39 @@ class DuckdbADKStore(BaseSyncADKStore["DuckDBConfig"]):
             conn.execute(delete_session_sql, (session_id,))
             conn.commit()
 
-    def list_sessions(self, app_name: str, user_id: str) -> "list[SessionRecord]":
-        """List all sessions for a user in an app.
+    def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[SessionRecord]":
+        """List sessions for an app, optionally filtered by user.
 
         Args:
             app_name: Application name.
-            user_id: User identifier.
+            user_id: User identifier. If None, lists all sessions for the app.
 
         Returns:
             List of session records ordered by update_time DESC.
 
         Notes:
-            Uses composite index on (app_name, user_id).
-        """
-        sql = f"""
-        SELECT id, app_name, user_id, state, create_time, update_time
-        FROM {self._session_table}
-        WHERE app_name = ? AND user_id = ?
-        ORDER BY update_time DESC
+            Uses composite index on (app_name, user_id) when user_id is provided.
         """
+        if user_id is None:
+            sql = f"""
+            SELECT id, app_name, user_id, state, create_time, update_time
+            FROM {self._session_table}
+            WHERE app_name = ?
+            ORDER BY update_time DESC
+            """
+            params: tuple[str, ...] = (app_name,)
+        else:
+            sql = f"""
+            SELECT id, app_name, user_id, state, create_time, update_time
+            FROM {self._session_table}
+            WHERE app_name = ? AND user_id = ?
+            ORDER BY update_time DESC
+            """
+            params = (app_name, user_id)
 
         try:
             with self._config.provide_connection() as conn:
-                cursor = conn.execute(sql, (app_name, user_id))
+                cursor = conn.execute(sql, params)
                 rows = cursor.fetchall()
 
                 return [

sqlspec/adapters/duckdb/driver.py

@@ -15,6 +15,7 @@ from sqlspec.core.statement import SQL, StatementConfig
 from sqlspec.driver import SyncDriverAdapterBase
 from sqlspec.exceptions import (
     CheckViolationError,
+    DatabaseConnectionError,
     DataError,
     ForeignKeyViolationError,
     IntegrityError,
@@ -32,9 +33,12 @@ if TYPE_CHECKING:
     from contextlib import AbstractContextManager
 
     from sqlspec.adapters.duckdb._types import DuckDBConnection
-    from sqlspec.core.result import SQLResult
+    from sqlspec.builder import QueryBuilder
+    from sqlspec.core import Statement, StatementFilter
+    from sqlspec.core.result import ArrowResult, SQLResult
     from sqlspec.driver import ExecutionResult
     from sqlspec.driver._sync import SyncDataDictionaryBase
+    from sqlspec.typing import StatementParameters
 
 __all__ = ("DuckDBCursor", "DuckDBDriver", "DuckDBExceptionHandler", "duckdb_statement_config")
 
@@ -447,3 +451,85 @@ class DuckDBDriver(SyncDriverAdapterBase):
         if self._data_dictionary is None:
             self._data_dictionary = DuckDBSyncDataDictionary()
         return self._data_dictionary
+
+    def select_to_arrow(
+        self,
+        statement: "Statement | QueryBuilder",
+        /,
+        *parameters: "StatementParameters | StatementFilter",
+        statement_config: "StatementConfig | None" = None,
+        return_format: str = "table",
+        native_only: bool = False,
+        batch_size: int | None = None,
+        arrow_schema: Any = None,
+        **kwargs: Any,
+    ) -> "ArrowResult":
+        """Execute query and return results as Apache Arrow (DuckDB native path).
+
+        DuckDB provides native Arrow support via cursor.arrow().
+        This is the fastest path due to DuckDB's columnar architecture.
+
+        Args:
+            statement: SQL statement, string, or QueryBuilder
+            *parameters: Query parameters or filters
+            statement_config: Optional statement configuration override
+            return_format: "table" for pyarrow.Table (default), "batch" for RecordBatch
+            native_only: Ignored for DuckDB (always uses native path)
+            batch_size: Batch size hint (for future streaming implementation)
+            arrow_schema: Optional pyarrow.Schema for type casting
+            **kwargs: Additional keyword arguments
+
+        Returns:
+            ArrowResult with native Arrow data
+
+        Raises:
+            MissingDependencyError: If pyarrow not installed
+            SQLExecutionError: If query execution fails
+
+        Example:
+            >>> result = driver.select_to_arrow(
+            ...     "SELECT * FROM users WHERE age > ?", 18
+            ... )
+            >>> df = result.to_pandas()  # Fast zero-copy conversion
+        """
+        from sqlspec.utils.module_loader import ensure_pyarrow
+
+        ensure_pyarrow()
+
+        import pyarrow as pa
+
+        from sqlspec.core.result import create_arrow_result
+
+        # Prepare statement
+        config = statement_config or self.statement_config
+        prepared_statement = self.prepare_statement(statement, parameters, statement_config=config, kwargs=kwargs)
+
+        # Execute query and get native Arrow
+        with self.with_cursor(self.connection) as cursor, self.handle_database_exceptions():
+            if cursor is None:
+                msg = "Failed to create cursor"
+                raise DatabaseConnectionError(msg)
+
+            # Get compiled SQL and parameters
+            sql, driver_params = self._get_compiled_sql(prepared_statement, config)
+
+            # Execute query
+            cursor.execute(sql, driver_params or ())
+
+            # DuckDB native Arrow (zero-copy!)
+            arrow_reader = cursor.arrow()
+            arrow_table = arrow_reader.read_all()
+
+            # Apply schema casting if requested
+            if arrow_schema is not None:
+                arrow_table = arrow_table.cast(arrow_schema)
+
+            # Convert to batch if requested
+            if return_format == "batch":
+                batches = arrow_table.to_batches()
+                arrow_data: Any = batches[0] if batches else pa.RecordBatch.from_pydict({})
+            else:
+                arrow_data = arrow_table
+
+        # Create ArrowResult
+        return create_arrow_result(statement=prepared_statement, data=arrow_data, rows_affected=arrow_data.num_rows)