sqlspec 0.9.0__py3-none-any.whl → 0.10.0__py3-none-any.whl

sqlspec/adapters/asyncpg/driver.py

@@ -1,19 +1,22 @@
 import logging
 import re
-from typing import TYPE_CHECKING, Any, Optional, Union, cast
+from typing import TYPE_CHECKING, Any, Optional, Union, cast, overload
 
 from asyncpg import Connection
 from typing_extensions import TypeAlias
 
-from sqlspec.base import AsyncDriverAdapterProtocol, T
+from sqlspec.base import AsyncDriverAdapterProtocol
 from sqlspec.exceptions import SQLParsingError
+from sqlspec.mixins import SQLTranslatorMixin
 from sqlspec.statement import PARAM_REGEX, SQLStatement
 
 if TYPE_CHECKING:
+    from collections.abc import Sequence
+
     from asyncpg.connection import Connection
     from asyncpg.pool import PoolConnectionProxy
 
-    from sqlspec.typing import ModelDTOT, StatementParameterType
+    from sqlspec.typing import ModelDTOT, StatementParameterType, T
 
 __all__ = ("AsyncpgConnection", "AsyncpgDriver")
 
@@ -33,7 +36,10 @@ QMARK_REGEX = re.compile(
 AsyncpgConnection: TypeAlias = "Union[Connection[Any], PoolConnectionProxy[Any]]"  # pyright: ignore[reportMissingTypeArgument]
 
 
-class AsyncpgDriver(AsyncDriverAdapterProtocol["AsyncpgConnection"]):
+class AsyncpgDriver(
+    SQLTranslatorMixin["AsyncpgConnection"],
+    AsyncDriverAdapterProtocol["AsyncpgConnection"],
+):
     """AsyncPG Postgres Driver Adapter."""
 
     connection: "AsyncpgConnection"
@@ -196,6 +202,28 @@ class AsyncpgDriver(AsyncDriverAdapterProtocol["AsyncpgConnection"]):
         # No parameters provided and none found in SQL, return original SQL from SQLStatement and empty tuple
         return sql, ()  # asyncpg expects a sequence, even if empty
 
+    @overload
+    async def select(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AsyncpgConnection]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "Sequence[dict[str, Any]]": ...
+    @overload
+    async def select(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AsyncpgConnection]" = None,
+        schema_type: "type[ModelDTOT]",
+        **kwargs: Any,
+    ) -> "Sequence[ModelDTOT]": ...
     async def select(
         self,
         sql: str,
@@ -205,7 +233,7 @@ class AsyncpgDriver(AsyncDriverAdapterProtocol["AsyncpgConnection"]):
         connection: Optional["AsyncpgConnection"] = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
         **kwargs: Any,
-    ) -> "list[Union[ModelDTOT, dict[str, Any]]]":
+    ) -> "Sequence[Union[ModelDTOT, dict[str, Any]]]":
         """Fetch data from the database.
 
         Args:
@@ -229,6 +257,28 @@ class AsyncpgDriver(AsyncDriverAdapterProtocol["AsyncpgConnection"]):
             return [dict(row.items()) for row in results]  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
         return [cast("ModelDTOT", schema_type(**dict(row.items()))) for row in results]  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
 
+    @overload
+    async def select_one(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AsyncpgConnection]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "dict[str, Any]": ...
+    @overload
+    async def select_one(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AsyncpgConnection]" = None,
+        schema_type: "type[ModelDTOT]",
+        **kwargs: Any,
+    ) -> "ModelDTOT": ...
     async def select_one(
         self,
         sql: str,
@@ -262,6 +312,28 @@ class AsyncpgDriver(AsyncDriverAdapterProtocol["AsyncpgConnection"]):
             return dict(result.items())  # type: ignore[attr-defined]
         return cast("ModelDTOT", schema_type(**dict(result.items())))  # type: ignore[attr-defined]
 
+    @overload
+    async def select_one_or_none(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AsyncpgConnection]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "Optional[dict[str, Any]]": ...
+    @overload
+    async def select_one_or_none(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AsyncpgConnection]" = None,
+        schema_type: "type[ModelDTOT]",
+        **kwargs: Any,
+    ) -> "Optional[ModelDTOT]": ...
     async def select_one_or_none(
         self,
         sql: str,
@@ -295,6 +367,28 @@ class AsyncpgDriver(AsyncDriverAdapterProtocol["AsyncpgConnection"]):
             return dict(result.items())
         return cast("ModelDTOT", schema_type(**dict(result.items())))
 
+    @overload
+    async def select_value(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AsyncpgConnection]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "Any": ...
+    @overload
+    async def select_value(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AsyncpgConnection]" = None,
+        schema_type: "type[T]",
+        **kwargs: Any,
+    ) -> "T": ...
     async def select_value(
         self,
         sql: str,
@@ -326,6 +420,28 @@ class AsyncpgDriver(AsyncDriverAdapterProtocol["AsyncpgConnection"]):
             return result
         return schema_type(result)  # type: ignore[call-arg]
 
+    @overload
+    async def select_value_or_none(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AsyncpgConnection]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "Optional[Any]": ...
+    @overload
+    async def select_value_or_none(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AsyncpgConnection]" = None,
+        schema_type: "type[T]",
+        **kwargs: Any,
+    ) -> "Optional[T]": ...
     async def select_value_or_none(
         self,
         sql: str,
@@ -381,6 +497,28 @@ class AsyncpgDriver(AsyncDriverAdapterProtocol["AsyncpgConnection"]):
         except (ValueError, IndexError, AttributeError):
             return -1  # Fallback if we can't parse the status
 
+    @overload
+    async def insert_update_delete_returning(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AsyncpgConnection]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "dict[str, Any]": ...
+    @overload
+    async def insert_update_delete_returning(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AsyncpgConnection]" = None,
+        schema_type: "type[ModelDTOT]",
+        **kwargs: Any,
+    ) -> "ModelDTOT": ...
     async def insert_update_delete_returning(
         self,
         sql: str,

sqlspec/adapters/bigquery/__init__.py

@@ -0,0 +1,4 @@
+from sqlspec.adapters.bigquery.config import BigQueryConfig, BigQueryConnectionConfig
+from sqlspec.adapters.bigquery.driver import BigQueryConnection, BigQueryDriver
+
+__all__ = ("BigQueryConfig", "BigQueryConnection", "BigQueryConnectionConfig", "BigQueryDriver")

sqlspec/adapters/bigquery/config/__init__.py

@@ -0,0 +1,3 @@
+from sqlspec.adapters.bigquery.config._sync import BigQueryConfig, BigQueryConnectionConfig
+
+__all__ = ("BigQueryConfig", "BigQueryConnectionConfig")

sqlspec/adapters/bigquery/config/_common.py

@@ -0,0 +1,40 @@
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Optional
+
+from google.cloud.bigquery import LoadJobConfig, QueryJobConfig
+
+if TYPE_CHECKING:
+    from google.api_core.client_info import ClientInfo
+    from google.api_core.client_options import ClientOptions
+    from google.auth.credentials import Credentials
+
+__all__ = ("BigQueryConnectionConfigCommon",)
+
+
+@dataclass
+class BigQueryConnectionConfigCommon:
+    """Common configuration options for BigQuery."""
+
+    project: "Optional[str]" = field(default=None)
+    """Google Cloud project ID."""
+    location: "Optional[str]" = field(default=None)
+    """Default geographic location for jobs and datasets."""
+    credentials: "Optional[Credentials]" = field(default=None, hash=False)
+    """Credentials to use for authentication."""
+    dataset_id: "Optional[str]" = field(default=None)
+    """Default dataset ID to use if not specified in queries."""
+    credentials_path: "Optional[str]" = field(default=None)
+    """Path to Google Cloud service account key file (JSON). If None, attempts default authentication."""
+    client_options: "Optional[ClientOptions]" = field(default=None, hash=False)
+    """Client options used to set user options on the client (e.g., api_endpoint)."""
+    default_query_job_config: "Optional[QueryJobConfig]" = field(default=None, hash=False)
+    """Default QueryJobConfig settings."""
+    default_load_job_config: "Optional[LoadJobConfig]" = field(default=None, hash=False)
+    """Default LoadJobConfig settings."""
+    client_info: "Optional[ClientInfo]" = field(default=None, hash=False)
+    """Client info used to send a user-agent string along with API requests."""
+
+    def __post_init__(self) -> None:
+        """Post-initialization hook."""
+        if self.default_query_job_config is None:
+            self.default_query_job_config = QueryJobConfig(default_dataset=self.dataset_id)

sqlspec/adapters/bigquery/config/_sync.py

@@ -0,0 +1,87 @@
+import contextlib
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Optional
+
+from sqlspec.adapters.bigquery.config._common import BigQueryConnectionConfigCommon
+from sqlspec.adapters.bigquery.driver import BigQueryConnection, BigQueryDriver
+from sqlspec.base import NoPoolSyncConfig
+from sqlspec.typing import dataclass_to_dict
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
+__all__ = ("BigQueryConfig", "BigQueryConnectionConfig")
+
+
+class BigQueryConnectionConfig(BigQueryConnectionConfigCommon):
+    """BigQuery Connection Configuration."""
+
+
+@dataclass
+class BigQueryConfig(NoPoolSyncConfig["BigQueryConnection", "BigQueryDriver"]):
+    """BigQuery Synchronous Driver Configuration."""
+
+    connection_config: "BigQueryConnectionConfig" = field(default_factory=BigQueryConnectionConfig)
+    """BigQuery Connection Configuration."""
+    driver_type: "type[BigQueryDriver]" = field(init=False, repr=False, default=BigQueryDriver)
+    """BigQuery Driver Type."""
+    connection_type: "type[BigQueryConnection]" = field(init=False, repr=False, default=BigQueryConnection)
+    """BigQuery Connection Type."""
+    pool_instance: "None" = field(init=False, repr=False, default=None, hash=False)
+    """This is set to have a init=False since BigQuery does not support pooling."""
+    connection_instance: "Optional[BigQueryConnection]" = field(init=False, repr=False, default=None, hash=False)
+    """BigQuery Connection Instance."""
+
+    @property
+    def connection_config_dict(self) -> "dict[str, Any]":
+        """Return the connection configuration as a dict.
+
+        Returns:
+            A string keyed dict of config kwargs for the BigQueryConnection constructor.
+        """
+        return dataclass_to_dict(
+            self.connection_config,
+            exclude_empty=True,
+            exclude_none=True,
+            exclude={"dataset_id", "credentials_path"},
+        )
+
+    def create_connection(self) -> "BigQueryConnection":
+        """Create a BigQuery Client instance.
+
+        Returns:
+            A BigQuery Client instance.
+        """
+        if self.connection_instance is not None:
+            return self.connection_instance
+
+        self.connection_instance = self.connection_type(**self.connection_config_dict)
+        return self.connection_instance
+
+    @contextlib.contextmanager
+    def provide_connection(self, *args: Any, **kwargs: Any) -> "Iterator[BigQueryConnection]":
+        """Provide a BigQuery client within a context manager.
+
+        Args:
+            *args: Additional arguments to pass to the connection.
+            **kwargs: Additional keyword arguments to pass to the connection.
+
+        Yields:
+            An iterator of BigQuery Client instances.
+        """
+        conn = self.create_connection()
+        yield conn
+
+    @contextlib.contextmanager
+    def provide_session(self, *args: Any, **kwargs: Any) -> "Iterator[BigQueryDriver]":
+        """Provide a BigQuery driver session within a context manager.
+
+        Args:
+            *args: Additional arguments to pass to the driver.
+            **kwargs: Additional keyword arguments to pass to the driver.
+
+        Yields:
+            An iterator of BigQueryDriver instances.
+        """
+        conn = self.create_connection()
+        yield self.driver_type(connection=conn)