feat: implements a driver query protocol (#21)

This change implements the first iteration of a "Driver".   This is a class designed to make interacting with database connections more consistent across the various adapters.

Author: cofin

.github/workflows/ci.yml

@@ -108,51 +108,51 @@ jobs:
       - name: Test
         run: uv run pytest -m ""
 
-  test-windows:
-    runs-on: windows-latest
-    strategy:
-      fail-fast: true
-      matrix:
-        python-version: ["3.12", "3.13"]
-    timeout-minutes: 30
-    steps:
-      - name: Check out repository
-        uses: actions/checkout@v4
-
-      - name: Install uv
-        uses: astral-sh/setup-uv@v3
-
-      - name: Set up Python
-        run: uv python install ${{ matrix.python-version }}
-
-      - name: Install dependencies
-        run: uv sync --all-extras --dev
-
-      - name: Test
-        run: uv run pytest -m ""
-
-  test-osx:
-    runs-on: macos-latest
-    strategy:
-      fail-fast: true
-      matrix:
-        python-version: ["3.11", "3.12", "3.13"]
-    timeout-minutes: 30
-    steps:
-      - name: Check out repository
-        uses: actions/checkout@v4
-
-      - name: Install uv
-        uses: astral-sh/setup-uv@v3
-
-      - name: Set up Python
-        run: uv python install ${{ matrix.python-version }}
-
-      - name: Install dependencies
-        run: uv sync --all-extras --dev
-
-      - name: Test
-        run: uv run pytest -m ""
+  # test-windows:
+  #   runs-on: windows-latest
+  #   strategy:
+  #     fail-fast: true
+  #     matrix:
+  #       python-version: ["3.12", "3.13"]
+  #   timeout-minutes: 30
+  #   steps:
+  #     - name: Check out repository
+  #       uses: actions/checkout@v4
+
+  #     - name: Install uv
+  #       uses: astral-sh/setup-uv@v3
+
+  #     - name: Set up Python
+  #       run: uv python install ${{ matrix.python-version }}
+
+  #     - name: Install dependencies
+  #       run: uv sync --all-extras --dev
+
+  #     - name: Test
+  #       run: uv run pytest -m ""
+
+  # test-osx:
+  #   runs-on: macos-latest
+  #   strategy:
+  #     fail-fast: true
+  #     matrix:
+  #       python-version: ["3.11", "3.12", "3.13"]
+  #   timeout-minutes: 30
+  #   steps:
+  #     - name: Check out repository
+  #       uses: actions/checkout@v4
+
+  #     - name: Install uv
+  #       uses: astral-sh/setup-uv@v3
+
+  #     - name: Set up Python
+  #       run: uv python install ${{ matrix.python-version }}
+
+  #     - name: Install dependencies
+  #       run: uv sync --all-extras --dev
+
+  #     - name: Test
+  #       run: uv run pytest -m ""
 
   build-docs:
     needs:

.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.10.0"
+    rev: "v0.11.6"
     hooks:
       - id: ruff
         args: ["--fix"]

docs/examples/litestar_duckllm.py

@@ -0,0 +1,55 @@
+"""Litestar DuckLLM
+
+This example demonstrates how to use the Litestar framework with the DuckLLM extension.
+
+The example uses the `SQLSpec` extension to create a connection to the DuckDB database.
+The `DuckDB` adapter is used to create a connection to the database.
+"""
+
+# /// script
+# dependencies = [
+#   "sqlspec[duckdb,performance] @ git+https://github.com/litestar-org/sqlspec.git@query-service",
+#   "litestar[standard]",
+# ]
+# ///
+
+from duckdb import DuckDBPyConnection
+from litestar import Litestar, post
+from msgspec import Struct
+
+from sqlspec.adapters.duckdb import DuckDB
+from sqlspec.extensions.litestar import SQLSpec
+
+
+class ChatMessage(Struct):
+    message: str
+
+
+@post("/chat", sync_to_thread=True)
+def duckllm_chat(db_connection: DuckDBPyConnection, data: ChatMessage) -> ChatMessage:
+    result = db_connection.execute("SELECT open_prompt(?)", (data.message,)).fetchall()
+    return ChatMessage(message=result[0][0])
+
+
+sqlspec = SQLSpec(
+    config=DuckDB(
+        extensions=[{"name": "open_prompt"}],
+        secrets=[
+            {
+                "secret_type": "open_prompt",
+                "name": "open_prompt",
+                "value": {
+                    "api_url": "http://127.0.0.1:11434/v1/chat/completions",
+                    "model_name": "gemma3:1b",
+                    "api_timeout": "120",
+                },
+            }
+        ],
+    ),
+)
+app = Litestar(route_handlers=[duckllm_chat], plugins=[sqlspec], debug=True)
+
+if __name__ == "__main__":
+    import uvicorn
+
+    uvicorn.run(app, host="0.0.0.0", port=8000)

docs/examples/litestar_gemini.py

@@ -0,0 +1,57 @@
+"""Litestar DuckLLM
+
+This example demonstrates how to use the Litestar framework with the DuckLLM extension.
+
+The example uses the `SQLSpec` extension to create a connection to the DuckDB database.
+The `DuckDB` adapter is used to create a connection to the database.
+"""
+
+# /// script
+# dependencies = [
+#   "sqlspec[duckdb,performance] @ git+https://github.com/litestar-org/sqlspec.git@query-service",
+#   "litestar[standard]",
+# ]
+# ///
+
+import os
+
+from sqlspec import SQLSpec
+from sqlspec.adapters.duckdb import DuckDB
+
+EMBEDDING_MODEL = "gemini-embedding-exp-03-07"
+GOOGLE_API_KEY = os.environ.get("GOOGLE_API_KEY")
+API_URL = (
+    f"https://generativelanguage.googleapis.com/v1beta/models/{EMBEDDING_MODEL}:embedContent?key=${GOOGLE_API_KEY}"
+)
+
+sql = SQLSpec()
+etl_config = sql.add_config(
+    DuckDB(
+        extensions=[{"name": "vss"}, {"name": "http_client"}],
+        on_connection_create=lambda connection: connection.execute(f"""
+            CREATE IF NOT EXISTS MACRO generate_embedding(q) AS (
+                WITH  __request AS (
+                    SELECT http_post(
+                        '{API_URL}',
+                        headers => MAP {{
+                            'accept': 'application/json',
+                        }},
+                        params => MAP {{
+                            'model': 'models/{EMBEDDING_MODEL}',
+                            'parts': [{{ 'text': q }}],
+                            'taskType': 'SEMANTIC_SIMILARITY'
+                        }}
+                    ) AS response
+                )
+                SELECT *
+                FROM __request,
+            );
+        """),
+    )
+)
+
+
+if __name__ == "__main__":
+    with sql.get_connection(etl_config) as connection:
+        result = connection.execute("SELECT generate_embedding('example text')")
+        print(result.fetchall())

docs/examples/litestar_multi_db.py

@@ -2,27 +2,33 @@
 from duckdb import DuckDBPyConnection
 from litestar import Litestar, get
 
-from sqlspec.adapters.aiosqlite import AiosqliteConfig
-from sqlspec.adapters.duckdb import DuckDBConfig
+from sqlspec.adapters.aiosqlite import Aiosqlite
+from sqlspec.adapters.duckdb import DuckDB
 from sqlspec.extensions.litestar import DatabaseConfig, SQLSpec
 
 
 @get("/test", sync_to_thread=True)
-def simple_select(etl_session: DuckDBPyConnection) -> dict[str, str]:
-    result = etl_session.execute("SELECT 'Hello, world!' AS greeting").fetchall()
+def simple_select(etl_connection: DuckDBPyConnection) -> dict[str, str]:
+    result = etl_connection.execute("SELECT 'Hello, world!' AS greeting").fetchall()
     return {"greeting": result[0][0]}
 
 
 @get("/")
 async def simple_sqlite(db_connection: Connection) -> dict[str, str]:
     result = await db_connection.execute_fetchall("SELECT 'Hello, world!' AS greeting")
-    return {"greeting": result[0][0]}  # type: ignore  # noqa: PGH003
+    return {"greeting": result[0][0]}  # type: ignore
 
 
 sqlspec = SQLSpec(
     config=[
-        DatabaseConfig(config=AiosqliteConfig(), commit_mode="autocommit"),
-        DatabaseConfig(config=DuckDBConfig(), connection_key="etl_session"),
+        DatabaseConfig(config=Aiosqlite(), commit_mode="autocommit"),
+        DatabaseConfig(
+            config=DuckDB(
+                extensions=[{"name": "vss", "force_install": True}],
+                secrets=[{"secret_type": "s3", "name": "s3_secret", "value": {"key_id": "abcd"}}],
+            ),
+            connection_key="etl_connection",
+        ),
     ],
 )
 app = Litestar(route_handlers=[simple_sqlite, simple_select], plugins=[sqlspec])

docs/examples/litestar_single_db.py

@@ -1,20 +1,20 @@
 from aiosqlite import Connection
 from litestar import Litestar, get
 
-from sqlspec.adapters.aiosqlite import AiosqliteConfig
+from sqlspec.adapters.aiosqlite import Aiosqlite
 from sqlspec.extensions.litestar import SQLSpec
 
 
 @get("/")
-async def simple_sqlite(db_session: Connection) -> dict[str, str]:
+async def simple_sqlite(db_connection: Connection) -> dict[str, str]:
     """Simple select statement.
 
     Returns:
         dict[str, str]: The greeting.
     """
-    result = await db_session.execute_fetchall("SELECT 'Hello, world!' AS greeting")
-    return {"greeting": result[0][0]}  # type: ignore  # noqa: PGH003
+    result = await db_connection.execute_fetchall("SELECT 'Hello, world!' AS greeting")
+    return {"greeting": result[0][0]}  # type: ignore
 
 
-sqlspec = SQLSpec(config=AiosqliteConfig())
+sqlspec = SQLSpec(config=Aiosqlite())
 app = Litestar(route_handlers=[simple_sqlite], plugins=[sqlspec])

pyproject.toml

@@ -10,7 +10,7 @@ requires-python = ">=3.9, <4.0"
 version = "0.7.1"
 
 [project.optional-dependencies]
-adbc = ["adbc-driver-manager", "pyarrow"]
+adbc = ["adbc_driver_manager", "pyarrow"]
 aioodbc = ["aioodbc"]
 aiosqlite = ["aiosqlite"]
 asyncmy = ["asyncmy"]
@@ -35,9 +35,7 @@ uuid = ["uuid-utils>=0.6.1"]
 [dependency-groups]
 build = ["bump-my-version"]
 dev = [
-  "adbc-driver-sqlite",
-  "adbc-driver-postgresql",
-  "adbc-driver-flightsql",
+  { include-group = "extras" },
   { include-group = "lint" },
   { include-group = "doc" },
   { include-group = "test" },
@@ -59,6 +57,15 @@ doc = [
   "myst-parser",
   "sphinx-autodoc-typehints",
 ]
+extras = [
+  "adbc_driver_manager",
+  "pyarrow",
+  "polars",
+  "adbc_driver_sqlite",
+  "adbc_driver_postgresql",
+  "adbc_driver_flightsql",
+  "adbc_driver_bigquery",
+]
 lint = [
   "mypy>=1.13.0",
   "pre-commit>=3.5.0",
@@ -75,7 +82,7 @@ test = [
   "pytest>=8.0.0",
   "pytest-asyncio>=0.23.8",
   "pytest-cov>=5.0.0",
-  "pytest-databases>=0.10.0",
+  "pytest-databases[postgres,oracle,mysql,bigquery,spanner]>=0.12.2",
   "pytest-mock>=3.14.0",
   "pytest-sugar>=1.0.0",
   "pytest-xdist>=3.6.1",
@@ -217,10 +224,17 @@ module = [
 
 [tool.pyright]
 disableBytesTypePromotions = true
-exclude = ["tools", "docs"]
+exclude = ["**/node_modules", "**/__pycache__", ".venv", "tools", "docs"]
 include = ["sqlspec", "tests"]
 pythonVersion = "3.9"
+reportMissingTypeStubs = false
+reportPrivateImportUsage = false
+reportPrivateUsage = false
+reportUnknownArgumentType = false
+reportUnknownMemberType = false
+reportUnknownVariableType = false
 reportUnnecessaryTypeIgnoreComments = true
+root = "."
 
 
 [tool.slotscheck]
@@ -269,6 +283,7 @@ ignore = [
   "CPY001",  # pycodestyle - Missing Copywrite notice at the top of the file
   "RUF029",  # Ruff - function is declared as async but has no awaitable calls
   "COM812",  # flake8-comma - Missing trailing comma
+  "PGH003",  # Use Specific ignore for pyright
 ]
 select = ["ALL"]
 

sqlspec/__init__.py

@@ -0,0 +1,15 @@
+from sqlspec import adapters, base, exceptions, extensions, filters, typing, utils
+from sqlspec.__metadata__ import __version__
+from sqlspec.base import SQLSpec
+
+__all__ = (
+    "SQLSpec",
+    "__version__",
+    "adapters",
+    "base",
+    "exceptions",
+    "extensions",
+    "filters",
+    "typing",
+    "utils",
+)