sqlspec 0.9.1__py3-none-any.whl → 0.10.1__py3-none-any.whl

sqlspec/adapters/aiosqlite/driver.py

@@ -1,33 +1,38 @@
 from contextlib import asynccontextmanager
 from typing import TYPE_CHECKING, Any, Optional, Union, cast, overload
 
-from sqlspec.base import AsyncDriverAdapterProtocol, T
+import aiosqlite
+
+from sqlspec.base import AsyncDriverAdapterProtocol
+from sqlspec.mixins import SQLTranslatorMixin
 
 if TYPE_CHECKING:
     from collections.abc import AsyncGenerator, Sequence
 
-    from aiosqlite import Connection, Cursor
-
-    from sqlspec.typing import ModelDTOT, StatementParameterType
+    from sqlspec.typing import ModelDTOT, StatementParameterType, T
 
-__all__ = ("AiosqliteDriver",)
+__all__ = ("AiosqliteConnection", "AiosqliteDriver")
+AiosqliteConnection = aiosqlite.Connection
 
 
-class AiosqliteDriver(AsyncDriverAdapterProtocol["Connection"]):
+class AiosqliteDriver(
+    SQLTranslatorMixin["AiosqliteConnection"],
+    AsyncDriverAdapterProtocol["AiosqliteConnection"],
+):
     """SQLite Async Driver Adapter."""
 
-    connection: "Connection"
+    connection: "AiosqliteConnection"
     dialect: str = "sqlite"
 
-    def __init__(self, connection: "Connection") -> None:
+    def __init__(self, connection: "AiosqliteConnection") -> None:
         self.connection = connection
 
     @staticmethod
-    async def _cursor(connection: "Connection", *args: Any, **kwargs: Any) -> "Cursor":
+    async def _cursor(connection: "AiosqliteConnection", *args: Any, **kwargs: Any) -> "aiosqlite.Cursor":
         return await connection.cursor(*args, **kwargs)
 
     @asynccontextmanager
-    async def _with_cursor(self, connection: "Connection") -> "AsyncGenerator[Cursor, None]":
+    async def _with_cursor(self, connection: "AiosqliteConnection") -> "AsyncGenerator[aiosqlite.Cursor, None]":
         cursor = await self._cursor(connection)
         try:
             yield cursor
@@ -42,7 +47,7 @@ class AiosqliteDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[AiosqliteConnection]" = None,
         schema_type: None = None,
         **kwargs: Any,
     ) -> "Sequence[dict[str, Any]]": ...
@@ -53,7 +58,7 @@ class AiosqliteDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[AiosqliteConnection]" = None,
         schema_type: "type[ModelDTOT]",
         **kwargs: Any,
     ) -> "Sequence[ModelDTOT]": ...
@@ -63,7 +68,7 @@ class AiosqliteDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: Optional["StatementParameterType"] = None,
         /,
         *,
-        connection: Optional["Connection"] = None,
+        connection: Optional["AiosqliteConnection"] = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
         **kwargs: Any,
     ) -> "Sequence[Union[ModelDTOT, dict[str, Any]]]":
@@ -91,7 +96,7 @@ class AiosqliteDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[AiosqliteConnection]" = None,
         schema_type: None = None,
         **kwargs: Any,
     ) -> "dict[str, Any]": ...
@@ -102,7 +107,7 @@ class AiosqliteDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[AiosqliteConnection]" = None,
         schema_type: "type[ModelDTOT]",
         **kwargs: Any,
     ) -> "ModelDTOT": ...
@@ -112,7 +117,7 @@ class AiosqliteDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: Optional["StatementParameterType"] = None,
         /,
         *,
-        connection: Optional["Connection"] = None,
+        connection: Optional["AiosqliteConnection"] = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
         **kwargs: Any,
     ) -> "Union[ModelDTOT, dict[str, Any]]":
@@ -139,7 +144,7 @@ class AiosqliteDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[AiosqliteConnection]" = None,
         schema_type: None = None,
         **kwargs: Any,
     ) -> "Optional[dict[str, Any]]": ...
@@ -150,7 +155,7 @@ class AiosqliteDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[AiosqliteConnection]" = None,
         schema_type: "type[ModelDTOT]",
         **kwargs: Any,
     ) -> "Optional[ModelDTOT]": ...
@@ -160,7 +165,7 @@ class AiosqliteDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: Optional["StatementParameterType"] = None,
         /,
         *,
-        connection: Optional["Connection"] = None,
+        connection: Optional["AiosqliteConnection"] = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
         **kwargs: Any,
     ) -> "Optional[Union[ModelDTOT, dict[str, Any]]]":
@@ -188,7 +193,7 @@ class AiosqliteDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[AiosqliteConnection]" = None,
         schema_type: None = None,
         **kwargs: Any,
     ) -> "Any": ...
@@ -199,7 +204,7 @@ class AiosqliteDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[AiosqliteConnection]" = None,
         schema_type: "type[T]",
         **kwargs: Any,
     ) -> "T": ...
@@ -209,7 +214,7 @@ class AiosqliteDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[AiosqliteConnection]" = None,
         schema_type: "Optional[type[T]]" = None,
         **kwargs: Any,
     ) -> "Union[T, Any]":
@@ -235,7 +240,7 @@ class AiosqliteDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[AiosqliteConnection]" = None,
         schema_type: None = None,
         **kwargs: Any,
     ) -> "Optional[Any]": ...
@@ -246,7 +251,7 @@ class AiosqliteDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[AiosqliteConnection]" = None,
         schema_type: "type[T]",
         **kwargs: Any,
     ) -> "Optional[T]": ...
@@ -256,7 +261,7 @@ class AiosqliteDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[AiosqliteConnection]" = None,
         schema_type: "Optional[type[T]]" = None,
         **kwargs: Any,
     ) -> "Optional[Union[T, Any]]":
@@ -282,7 +287,7 @@ class AiosqliteDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: Optional["StatementParameterType"] = None,
         /,
         *,
-        connection: Optional["Connection"] = None,
+        connection: Optional["AiosqliteConnection"] = None,
         **kwargs: Any,
     ) -> int:
         """Insert, update, or delete data from the database.
@@ -304,7 +309,7 @@ class AiosqliteDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[AiosqliteConnection]" = None,
         schema_type: None = None,
         **kwargs: Any,
     ) -> "dict[str, Any]": ...
@@ -315,7 +320,7 @@ class AiosqliteDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[AiosqliteConnection]" = None,
         schema_type: "type[ModelDTOT]",
         **kwargs: Any,
     ) -> "ModelDTOT": ...
@@ -325,7 +330,7 @@ class AiosqliteDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: Optional["StatementParameterType"] = None,
         /,
         *,
-        connection: Optional["Connection"] = None,
+        connection: Optional["AiosqliteConnection"] = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
         **kwargs: Any,
     ) -> "Optional[Union[dict[str, Any], ModelDTOT]]":
@@ -353,7 +358,7 @@ class AiosqliteDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: Optional["StatementParameterType"] = None,
         /,
         *,
-        connection: Optional["Connection"] = None,
+        connection: Optional["AiosqliteConnection"] = None,
         **kwargs: Any,
     ) -> str:
         """Execute a script.
@@ -374,7 +379,7 @@ class AiosqliteDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: Optional["StatementParameterType"] = None,
         /,
         *,
-        connection: Optional["Connection"] = None,
+        connection: Optional["AiosqliteConnection"] = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
         **kwargs: Any,
     ) -> "Optional[Union[dict[str, Any], ModelDTOT]]":

sqlspec/adapters/asyncmy/__init__.py

@@ -1,8 +1,9 @@
 from sqlspec.adapters.asyncmy.config import AsyncmyConfig, AsyncmyPoolConfig
-from sqlspec.adapters.asyncmy.driver import AsyncmyDriver  # type: ignore[attr-defined]
+from sqlspec.adapters.asyncmy.driver import AsyncmyConnection, AsyncmyDriver  # type: ignore[attr-defined]
 
 __all__ = (
     "AsyncmyConfig",
+    "AsyncmyConnection",
     "AsyncmyDriver",
     "AsyncmyPoolConfig",
 )

sqlspec/adapters/asyncmy/driver.py

@@ -3,33 +3,36 @@ from collections.abc import AsyncGenerator, Sequence
 from contextlib import asynccontextmanager
 from typing import TYPE_CHECKING, Any, Optional, Union, cast, overload
 
-from sqlspec.base import AsyncDriverAdapterProtocol, T
+from asyncmy import Connection
+
+from sqlspec.base import AsyncDriverAdapterProtocol
+from sqlspec.mixins import SQLTranslatorMixin
 
 if TYPE_CHECKING:
-    from asyncmy import Connection
     from asyncmy.cursors import Cursor
 
-    from sqlspec.typing import ModelDTOT, StatementParameterType
+    from sqlspec.typing import ModelDTOT, StatementParameterType, T
 
 __all__ = ("AsyncmyDriver",)
 
+AsyncmyConnection = Connection
+
 
-class AsyncmyDriver(AsyncDriverAdapterProtocol["Connection"]):
+class AsyncmyDriver(
+    SQLTranslatorMixin["AsyncmyConnection"],
+    AsyncDriverAdapterProtocol["AsyncmyConnection"],
+):
     """Asyncmy MySQL/MariaDB Driver Adapter."""
 
-    connection: "Connection"
+    connection: "AsyncmyConnection"
     dialect: str = "mysql"
 
-    def __init__(self, connection: "Connection") -> None:
+    def __init__(self, connection: "AsyncmyConnection") -> None:
         self.connection = connection
 
-    @staticmethod
-    async def _cursor(connection: "Connection") -> "Cursor":
-        return await connection.cursor()
-
     @staticmethod
     @asynccontextmanager
-    async def _with_cursor(connection: "Connection") -> AsyncGenerator["Cursor", None]:
+    async def _with_cursor(connection: "AsyncmyConnection") -> AsyncGenerator["Cursor", None]:
         cursor = connection.cursor()
         try:
             yield cursor
@@ -44,7 +47,7 @@ class AsyncmyDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[AsyncmyConnection]" = None,
         schema_type: None = None,
         **kwargs: Any,
     ) -> "Sequence[dict[str, Any]]": ...
@@ -55,7 +58,7 @@ class AsyncmyDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[AsyncmyConnection]" = None,
         schema_type: "type[ModelDTOT]",
         **kwargs: Any,
     ) -> "Sequence[ModelDTOT]": ...
@@ -65,7 +68,7 @@ class AsyncmyDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: Optional["StatementParameterType"] = None,
         /,
         *,
-        connection: Optional["Connection"] = None,
+        connection: Optional["AsyncmyConnection"] = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
         **kwargs: Any,
     ) -> "Sequence[Union[ModelDTOT, dict[str, Any]]]":
@@ -93,7 +96,7 @@ class AsyncmyDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[AsyncmyConnection]" = None,
         schema_type: None = None,
         **kwargs: Any,
     ) -> "dict[str, Any]": ...
@@ -104,7 +107,7 @@ class AsyncmyDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[AsyncmyConnection]" = None,
         schema_type: "type[ModelDTOT]",
         **kwargs: Any,
     ) -> "ModelDTOT": ...
@@ -114,7 +117,7 @@ class AsyncmyDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: Optional["StatementParameterType"] = None,
         /,
         *,
-        connection: Optional["Connection"] = None,
+        connection: Optional["AsyncmyConnection"] = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
         **kwargs: Any,
     ) -> "Union[ModelDTOT, dict[str, Any]]":
@@ -141,7 +144,7 @@ class AsyncmyDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[AsyncmyConnection]" = None,
         schema_type: None = None,
         **kwargs: Any,
     ) -> "Optional[dict[str, Any]]": ...
@@ -152,7 +155,7 @@ class AsyncmyDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[AsyncmyConnection]" = None,
         schema_type: "type[ModelDTOT]",
         **kwargs: Any,
     ) -> "Optional[ModelDTOT]": ...
@@ -162,7 +165,7 @@ class AsyncmyDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: Optional["StatementParameterType"] = None,
         /,
         *,
-        connection: Optional["Connection"] = None,
+        connection: Optional["AsyncmyConnection"] = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
         **kwargs: Any,
     ) -> "Optional[Union[ModelDTOT, dict[str, Any]]]":
@@ -190,7 +193,7 @@ class AsyncmyDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[AsyncmyConnection]" = None,
         schema_type: None = None,
         **kwargs: Any,
     ) -> "Any": ...
@@ -201,7 +204,7 @@ class AsyncmyDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[AsyncmyConnection]" = None,
         schema_type: "type[T]",
         **kwargs: Any,
     ) -> "T": ...
@@ -211,7 +214,7 @@ class AsyncmyDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[AsyncmyConnection]" = None,
         schema_type: "Optional[type[T]]" = None,
         **kwargs: Any,
     ) -> "Union[T, Any]":
@@ -240,7 +243,7 @@ class AsyncmyDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[AsyncmyConnection]" = None,
         schema_type: None = None,
         **kwargs: Any,
     ) -> "Optional[Any]": ...
@@ -251,7 +254,7 @@ class AsyncmyDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[AsyncmyConnection]" = None,
         schema_type: "type[T]",
         **kwargs: Any,
     ) -> "Optional[T]": ...
@@ -261,7 +264,7 @@ class AsyncmyDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[AsyncmyConnection]" = None,
         schema_type: "Optional[type[T]]" = None,
         **kwargs: Any,
     ) -> "Optional[Union[T, Any]]":
@@ -291,7 +294,7 @@ class AsyncmyDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: Optional["StatementParameterType"] = None,
         /,
         *,
-        connection: Optional["Connection"] = None,
+        connection: Optional["AsyncmyConnection"] = None,
         **kwargs: Any,
     ) -> int:
         """Insert, update, or delete data from the database.
@@ -313,7 +316,7 @@ class AsyncmyDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[AsyncmyConnection]" = None,
         schema_type: None = None,
         **kwargs: Any,
     ) -> "dict[str, Any]": ...
@@ -324,7 +327,7 @@ class AsyncmyDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[AsyncmyConnection]" = None,
         schema_type: "type[ModelDTOT]",
         **kwargs: Any,
     ) -> "ModelDTOT": ...
@@ -334,7 +337,7 @@ class AsyncmyDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: Optional["StatementParameterType"] = None,
         /,
         *,
-        connection: Optional["Connection"] = None,
+        connection: Optional["AsyncmyConnection"] = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
         **kwargs: Any,
     ) -> "Optional[Union[dict[str, Any], ModelDTOT]]":
@@ -363,7 +366,7 @@ class AsyncmyDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: Optional["StatementParameterType"] = None,
         /,
         *,
-        connection: Optional["Connection"] = None,
+        connection: Optional["AsyncmyConnection"] = None,
         **kwargs: Any,
     ) -> str:
         """Execute a script.

sqlspec/adapters/asyncpg/config.py

@@ -190,9 +190,7 @@ class AsyncpgConfig(AsyncDatabaseConfig["AsyncpgConnection", "Pool", "AsyncpgDri
             raise ImproperConfigurationError(msg) from e
 
     @asynccontextmanager
-    async def provide_connection(
-        self, *args: "Any", **kwargs: "Any"
-    ) -> "AsyncGenerator[PoolConnectionProxy[Any], None]":  # pyright: ignore[reportMissingTypeArgument,reportUnknownParameterType]
+    async def provide_connection(self, *args: "Any", **kwargs: "Any") -> "AsyncGenerator[AsyncpgConnection, None]":  # pyright: ignore[reportMissingTypeArgument,reportUnknownParameterType]
         """Create a connection instance.
 
         Yields:

sqlspec/adapters/asyncpg/driver.py

@@ -5,8 +5,9 @@ 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:
@@ -15,7 +16,7 @@ if TYPE_CHECKING:
     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")
 
@@ -35,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"

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)