feat: additional progress on new adapters

Author: cofin

.pre-commit-config.yaml

@@ -17,7 +17,7 @@ repos:
       - id: mixed-line-ending
       - id: trailing-whitespace
   - repo: https://github.com/charliermarsh/ruff-pre-commit
-    rev: "v0.7.2"
+    rev: "v0.7.3"
     hooks:
       - id: ruff
         args: ["--fix"]

pyproject.toml

@@ -9,6 +9,7 @@ requires-python = ">=3.9, <4.0"
 version = "0.2.0"
 
 [project.optional-dependencies]
+adbc = ["adbc-driver-manager", "pyarrow"]
 aioodbc = ["aioodbc"]
 aiosqlite = ["aiosqlite"]
 asyncmy = ["asyncmy"]
@@ -28,7 +29,7 @@ fastapi = ["fastapi"]
 flask = ["flask"]
 
 [dependency-groups]
-dev = [{ include-group = "lint" }, { include-group = "doc" }, { include-group = "test" }]
+dev = ["adbc-driver-sqlite", "adbc-driver-postgresql", "adbc-driver-flightsql", { include-group = "lint" }, { include-group = "doc" }, { include-group = "test" }]
 doc = [
   "auto-pytabs[sphinx]>=0.5.0",
   "git-cliff>=2.6.1",

sqlspec/adapters/adbc/__init__.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+from contextlib import contextmanager
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, TypeVar
+
+from sqlspec.types.configs import GenericDatabaseConfig
+from sqlspec.types.empty import Empty
+
+if TYPE_CHECKING:
+    from collections.abc import Generator
+    from typing import Any
+
+    from adbc_driver_manager.dbapi import Connection, Cursor
+
+    from sqlspec.types.empty import EmptyType
+
+__all__ = ("AdbcDatabaseConfig",)
+
+ConnectionT = TypeVar("ConnectionT", bound="Connection")
+CursorT = TypeVar("CursorT", bound="Cursor")
+
+
+@dataclass
+class AdbcDatabaseConfig(GenericDatabaseConfig):
+    """Configuration for ADBC connections.
+
+    This class provides configuration options for ADBC database connections using the
+    ADBC Driver Manager.([1](https://arrow.apache.org/adbc/current/python/api/adbc_driver_manager.html))
+    """
+
+    uri: str | EmptyType = Empty
+    """Database URI"""
+    driver_name: str | EmptyType = Empty
+    """Name of the ADBC driver to use"""
+    db_kwargs: dict[str, Any] | None = None
+    """Additional database-specific connection parameters"""
+
+    @property
+    def connection_params(self) -> dict[str, Any]:
+        """Return the connection parameters as a dict."""
+        return {
+            k: v
+            for k, v in {"uri": self.uri, "driver": self.driver_name, **(self.db_kwargs or {})}.items()
+            if v is not Empty
+        }
+
+    @contextmanager
+    def provide_connection(self, *args: Any, **kwargs: Any) -> Generator[Connection, None, None]:
+        """Create and provide a database connection."""
+        from adbc_driver_manager.dbapi import connect
+
+        with connect(**self.connection_params) as connection:
+            yield connection

sqlspec/adapters/adbc/config.py

@@ -0,0 +1,142 @@
+from __future__ import annotations
+
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Any
+
+from sqlspec.types.protocols import StatementType, SyncDriverAdapterProtocol
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Generator, Iterable
+
+    from adbc_driver_manager.dbapi import Connection, Cursor
+
+__all__ = ("AdbcAdapter",)
+
+
+class AdbcAdapter(SyncDriverAdapterProtocol):
+    """A synchronous ADBC SQLSpec Adapter."""
+
+    is_async: bool = False
+
+    def process_sql(self, op_type: StatementType, sql: str) -> str:
+        """Process SQL query."""
+        return sql
+
+    def _cursor(self, connection: Connection) -> Cursor:
+        """Get a cursor from a connection."""
+        return connection.cursor()
+
+    def _process_row(self, row: Any, column_names: list[str], record_class: Callable | None = None) -> Any:
+        """Process a row into the desired format."""
+        if record_class is None:
+            return row
+        return record_class(**{str(k): v for k, v in zip(column_names, row, strict=False)})
+
+    def select(
+        self,
+        connection: Connection,
+        sql: str,
+        parameters: list | dict,
+        record_class: Callable | None,
+    ) -> Iterable[Any]:
+        """Handle a relation-returning SELECT."""
+        cur = self._cursor(connection)
+        try:
+            cur.execute(sql, parameters)
+            if record_class is None:
+                yield from cur
+            else:
+                column_names = [desc[0] for desc in cur.description]
+                for row in cur:
+                    yield self._process_row(row, column_names, record_class)
+        finally:
+            cur.close()
+
+    def select_one(
+        self,
+        connection: Connection,
+        sql: str,
+        parameters: list | dict,
+        record_class: Callable | None,
+    ) -> Any | None:
+        """Handle a single-row-returning SELECT."""
+        cur = self._cursor(connection)
+        try:
+            cur.execute(sql, parameters)
+            row = cur.fetchone()
+            if row is None:
+                return None
+            if record_class is None:
+                return row
+            column_names = [desc[0] for desc in cur.description or []]
+            return self._process_row(row, column_names, record_class)
+        finally:
+            cur.close()
+
+    def select_scalar(
+        self,
+        connection: Connection,
+        sql: str,
+        parameters: list | dict,
+    ) -> Any | None:
+        """Handle a scalar-returning SELECT."""
+        cur = self._cursor(connection)
+        try:
+            cur.execute(sql, parameters)
+            row = cur.fetchone()
+            return row[0] if row else None
+        finally:
+            cur.close()
+
+    @contextmanager
+    def with_cursor(
+        self,
+        connection: Connection,
+        sql: str,
+        parameters: list | dict,
+    ) -> Generator[Cursor, None, None]:
+        """Execute a query and yield the cursor."""
+        cur = self._cursor(connection)
+        try:
+            cur.execute(sql, parameters)
+            yield cur
+        finally:
+            cur.close()
+
+    def insert_update_delete(
+        self,
+        connection: Connection,
+        sql: str,
+        parameters: list | dict,
+    ) -> int:
+        """Handle an INSERT, UPDATE, or DELETE."""
+        cur = self._cursor(connection)
+        try:
+            cur.execute(sql, parameters)
+            return cur.rowcount
+        finally:
+            cur.close()
+
+    def insert_update_delete_returning(
+        self,
+        connection: Connection,
+        sql: str,
+        parameters: list | dict,
+        record_class: Callable | None = None,
+    ) -> Any:
+        """Handle an INSERT, UPDATE, or DELETE with RETURNING clause."""
+        return self.select_one(connection, sql, parameters, record_class)
+
+    def insert_update_delete_many(
+        self,
+        connection: Connection,
+        sql: str,
+        parameters: list | dict,
+    ) -> int:
+        """Handle multiple INSERT, UPDATE, or DELETE operations."""
+        cur = self._cursor(connection)
+        try:
+            cur.executemany(sql, parameters)
+            return cur.rowcount
+        finally:
+            cur.close()

sqlspec/adapters/adbc/driver.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from collections import defaultdict
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, cast
 
 from sqlspec.sql.patterns import VAR_REF
 from sqlspec.types.protocols import AsyncDriverAdapterProtocol, StatementType
@@ -15,22 +15,23 @@
 __all__ = ("AsyncpgAdapter",)
 
 
-class MaybeAcquire:
+class ManagedConnection:
     """Context manager for handling connection acquisition from pools or direct connections."""
 
     def __init__(self, client: Connection | Pool) -> None:
         self.client = client
         self._managed_conn: Connection | None = None
 
     async def __aenter__(self) -> Connection:
-        if hasattr(self.client, "acquire"):
-            self._managed_conn = await self.client.acquire()
-            return self._managed_conn
-        return self.client
+        if "acquire" in dir(self.client):
+            self._managed_conn = await self.client.acquire()  # pyright: ignore[reportAttributeAccessIssue]
+            return cast("Connection", self._managed_conn)
+        self._managed_conn = None
+        return cast("Connection", self.client)
 
     async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> None:
         if self._managed_conn is not None:
-            await self.client.release(self._managed_conn)
+            await self.client.release(self._managed_conn)  # pyright: ignore[reportAttributeAccessIssue]
 
 
 class AsyncpgAdapter(AsyncDriverAdapterProtocol):
@@ -99,7 +100,7 @@ async def select(
         record_class: Callable | None,
     ) -> Iterable[Any]:
         """Handle a relation-returning SELECT."""
-        async with MaybeAcquire(connection) as conn:
+        async with ManagedConnection(connection) as conn:
             results = await conn.fetch(sql, *parameters)
             if record_class is not None:
                 return [record_class(**dict(rec)) for rec in results]
@@ -113,7 +114,7 @@ async def select_one(
         record_class: Callable | None,
     ) -> Any | None:
         """Handle a single-row-returning SELECT."""
-        async with MaybeAcquire(connection) as conn:
+        async with ManagedConnection(connection) as conn:
             result = await conn.fetchrow(sql, *parameters)
             if result is not None and record_class is not None:
                 return record_class(**dict(result))
@@ -126,7 +127,7 @@ async def select_scalar(
         parameters: list | dict,
     ) -> Any | None:
         """Handle a scalar-returning SELECT."""
-        async with MaybeAcquire(connection) as conn:
+        async with ManagedConnection(connection) as conn:
             return await conn.fetchval(sql, *parameters)
 
     async def with_cursor(
@@ -136,7 +137,7 @@ async def with_cursor(
         parameters: list | dict,
     ) -> AsyncGenerator[Cursor | CursorFactory, None]:
         """Execute a query and yield the cursor."""
-        async with MaybeAcquire(connection) as conn:
+        async with ManagedConnection(connection) as conn:
             stmt = await conn.prepare(sql)
             async with conn.transaction():
                 yield stmt.cursor(*parameters)
@@ -148,7 +149,7 @@ async def insert_update_delete(
         parameters: list | dict,
     ) -> int:
         """Handle an INSERT, UPDATE, or DELETE."""
-        async with MaybeAcquire(connection) as conn:
+        async with ManagedConnection(connection) as conn:
             result = await conn.execute(sql, *parameters)
             if isinstance(result, str):
                 return int(result.split()[-1])
@@ -172,7 +173,7 @@ async def execute_script(
         record_class: Callable | None = None,
     ) -> Any:
         """Execute a SQL script."""
-        async with MaybeAcquire(connection) as conn:
+        async with ManagedConnection(connection) as conn:
             if parameters:
                 result = await conn.fetch(sql, *parameters)
             else:

sqlspec/adapters/asyncmy/__init__.py

@@ -36,18 +36,22 @@ async def select(
         record_class: Callable | None,
     ) -> Iterable[Any]:
         """Handle a relation-returning SELECT."""
-        cur = await self._cursor(connection)
+        cur = self._cursor(connection)
         try:
             await cur.execute(sql, parameters)
+            column_names = [desc[0] for desc in cur.description]
+            cur.rowfactory = lambda *args: dict(zip(column_names, args))
+            data = await cur.fetchall()
             if record_class is None:
                 async for row in cur:
                     yield row
             else:
                 column_names = [desc[0] for desc in cur.description]
+                cur.rowfactory = lambda *args: dict(zip(column_names, args))
                 async for row in cur:
                     yield record_class(**{str(k): v for k, v in zip(column_names, row, strict=False)})
         finally:
-            await cur.close()
+            cur.close()
 
     async def select_one(
         self,
@@ -90,7 +94,7 @@ async def with_cursor(
         parameters: list | dict,
     ) -> AsyncGenerator[AsyncCursor, None]:
         """Execute a query and yield the cursor."""
-        cur = await self._cursor(connection)
+        cur = self._cursor(connection)
         try:
             await cur.execute(sql, parameters)
             yield cur
@@ -104,7 +108,7 @@ async def insert_update_delete(
         parameters: list | dict,
     ) -> int:
         """Handle an INSERT, UPDATE, or DELETE."""
-        cur = await self._cursor(connection)
+        cur = self._cursor(connection)
         try:
             await cur.execute(sql, parameters)
             return cur.rowcount