fix(litestar): signature namespace (#41)

Corrects a few issues with Litestar signature namespace support

Author: cofin

docs/examples/litestar_asyncpg.py

@@ -32,13 +32,8 @@
 @get("/")
 async def hello_world(db_session: AsyncpgDriver) -> dict[str, Any]:
     """Simple endpoint that returns a greeting from the database."""
-    # Execute a simple query that doesn't require real database tables
     result = await db_session.execute(SQL("SELECT 'Hello from AsyncPG!' as greeting"))
-
-    # Return the first row as a dictionary
-    if result.data:
-        return result.data[0]
-    return {"greeting": "No data returned"}
+    return result.get_first() or {"greeting": "No data returned"}
 
 
 @get("/version")

docs/examples/litestar_duckllm.py

@@ -1,4 +1,3 @@
-# type: ignore
 """Litestar DuckLLM
 
 This example demonstrates how to use the Litestar framework with the DuckLLM extension.
@@ -27,7 +26,8 @@ class ChatMessage(Struct):
 
 @post("/chat", sync_to_thread=True)
 def duckllm_chat(db_session: DuckDBDriver, data: ChatMessage) -> ChatMessage:
-    return db_session.execute("SELECT open_prompt(?)", data.message, schema_type=ChatMessage).get_first()
+    results = db_session.execute("SELECT open_prompt(?)", data.message, schema_type=ChatMessage).get_first()
+    return results or ChatMessage(message="No response from DuckLLM")
 
 
 sqlspec = SQLSpec(

docs/examples/litestar_multi_db.py

@@ -1,4 +1,3 @@
-# type: ignore
 """Litestar Multi DB
 
 This example demonstrates how to use multiple databases in a Litestar application.
@@ -19,22 +18,21 @@
 from sqlspec.adapters.aiosqlite import AiosqliteConfig, AiosqliteDriver
 from sqlspec.adapters.duckdb import DuckDBConfig, DuckDBDriver
 from sqlspec.extensions.litestar import DatabaseConfig, SQLSpec
+from sqlspec.statement.sql import SQL
 
 
 @get("/test", sync_to_thread=True)
 def simple_select(etl_session: DuckDBDriver) -> dict[str, str]:
-    from sqlspec.statement.sql import SQL
-
     result = etl_session.execute(SQL("SELECT 'Hello, ETL world!' AS greeting"))
     greeting = result.get_first()
     return {"greeting": greeting["greeting"] if greeting is not None else "hi"}
 
 
 @get("/")
 async def simple_sqlite(db_session: AiosqliteDriver) -> dict[str, str]:
-    from sqlspec.statement.sql import SQL
-
-    return await db_session.select_one(SQL("SELECT 'Hello, world!' AS greeting"))
+    result = await db_session.execute("SELECT 'Hello, world!' AS greeting")
+    greeting = result.get_first()
+    return {"greeting": greeting["greeting"] if greeting is not None else "hi"}
 
 
 sqlspec = SQLSpec(

docs/examples/unified_storage_demo.py

@@ -37,7 +37,7 @@ def demo_unified_storage_architecture() -> None:
                 range AS id,
                 'Product_' || (range % 10) AS product_name,
                 (random() * 1000)::int AS amount,
-                DATE '2024-01-01' + (range % 365) AS sale_date
+                DATE '2024-01-01' + (range % 365)::INTEGER AS sale_date
             FROM range(1000)
         """)
         )

sqlspec/adapters/aiosqlite/config.py

@@ -186,3 +186,16 @@ async def provide_session(self, *args: Any, **kwargs: Any) -> AsyncGenerator[Aio
     async def provide_pool(self, *args: Any, **kwargs: Any) -> None:
         """Aiosqlite doesn't support pooling."""
         return
+
+    def get_signature_namespace(self) -> "dict[str, type[Any]]":
+        """Get the signature namespace for Aiosqlite types.
+
+        This provides all Aiosqlite-specific types that Litestar needs to recognize
+        to avoid serialization attempts.
+
+        Returns:
+            Dictionary mapping type names to types.
+        """
+        namespace = super().get_signature_namespace()
+        namespace.update({"AiosqliteConnection": AiosqliteConnection})
+        return namespace

sqlspec/adapters/asyncmy/config.py

@@ -6,6 +6,7 @@
 from typing import TYPE_CHECKING, Any, ClassVar, Optional, Union
 
 import asyncmy
+from asyncmy.pool import Pool as AsyncmyPool
 
 from sqlspec.adapters.asyncmy.driver import AsyncmyConnection, AsyncmyDriver
 from sqlspec.config import AsyncDatabaseConfig
@@ -283,3 +284,16 @@ async def provide_pool(self, *args: Any, **kwargs: Any) -> "Pool":  # pyright: i
         if not self.pool_instance:
             self.pool_instance = await self.create_pool()
         return self.pool_instance
+
+    def get_signature_namespace(self) -> "dict[str, type[Any]]":
+        """Get the signature namespace for Asyncmy types.
+
+        This provides all Asyncmy-specific types that Litestar needs to recognize
+        to avoid serialization attempts.
+
+        Returns:
+            Dictionary mapping type names to types.
+        """
+        namespace = super().get_signature_namespace()
+        namespace.update({"AsyncmyConnection": AsyncmyConnection, "AsyncmyPool": AsyncmyPool})
+        return namespace

sqlspec/adapters/asyncpg/config.py

@@ -5,8 +5,10 @@
 from contextlib import asynccontextmanager
 from typing import TYPE_CHECKING, Any, ClassVar, TypedDict
 
-from asyncpg import Record
+from asyncpg import Connection, Record
 from asyncpg import create_pool as asyncpg_create_pool
+from asyncpg.connection import ConnectionMeta
+from asyncpg.pool import Pool, PoolConnectionProxy, PoolConnectionProxyMeta
 from typing_extensions import NotRequired, Unpack
 
 from sqlspec.adapters.asyncpg.driver import AsyncpgConnection, AsyncpgDriver
@@ -18,7 +20,6 @@
 if TYPE_CHECKING:
     from asyncio.events import AbstractEventLoop
 
-    from asyncpg.pool import Pool
     from sqlglot.dialects.dialect import DialectType
 
 
@@ -347,24 +348,15 @@ def get_signature_namespace(self) -> "dict[str, type[Any]]":
             Dictionary mapping type names to types.
         """
         namespace = super().get_signature_namespace()
-
-        try:
-            from asyncpg import Connection, Record
-            from asyncpg.connection import ConnectionMeta
-            from asyncpg.pool import Pool, PoolConnectionProxy, PoolConnectionProxyMeta
-
-            namespace.update(
-                {
-                    "Connection": Connection,
-                    "Pool": Pool,
-                    "PoolConnectionProxy": PoolConnectionProxy,
-                    "PoolConnectionProxyMeta": PoolConnectionProxyMeta,
-                    "ConnectionMeta": ConnectionMeta,
-                    "Record": Record,
-                    "AsyncpgConnection": type(AsyncpgConnection),  # The Union type alias
-                }
-            )
-        except ImportError:
-            logger.warning("Failed to import AsyncPG types for signature namespace")
-
+        namespace.update(
+            {
+                "Connection": Connection,
+                "Pool": Pool,
+                "PoolConnectionProxy": PoolConnectionProxy,
+                "PoolConnectionProxyMeta": PoolConnectionProxyMeta,
+                "ConnectionMeta": ConnectionMeta,
+                "Record": Record,
+                "AsyncpgConnection": type(AsyncpgConnection),
+            }
+        )
         return namespace

sqlspec/adapters/oracledb/config.py

@@ -315,6 +315,19 @@ def provide_pool(self, *args: Any, **kwargs: Any) -> "ConnectionPool":
             self.pool_instance = self.create_pool()
         return self.pool_instance
 
+    def get_signature_namespace(self) -> "dict[str, type[Any]]":
+        """Get the signature namespace for OracleDB types.
+
+        This provides all OracleDB-specific types that Litestar needs to recognize
+        to avoid serialization attempts.
+
+        Returns:
+            Dictionary mapping type names to types.
+        """
+        namespace = super().get_signature_namespace()
+        namespace.update({"OracleSyncConnection": OracleSyncConnection, "OracleAsyncConnection": OracleAsyncConnection})
+        return namespace
+
     @property
     def connection_config_dict(self) -> dict[str, Any]:
         """Return the connection configuration as a dict for Oracle operations.
@@ -624,3 +637,16 @@ async def provide_pool(self, *args: Any, **kwargs: Any) -> "AsyncConnectionPool"
         if not self.pool_instance:
             self.pool_instance = await self.create_pool()
         return self.pool_instance
+
+    def get_signature_namespace(self) -> "dict[str, type[Any]]":
+        """Get the signature namespace for OracleDB async types.
+
+        This provides all OracleDB async-specific types that Litestar needs to recognize
+        to avoid serialization attempts.
+
+        Returns:
+            Dictionary mapping type names to types.
+        """
+        namespace = super().get_signature_namespace()
+        namespace.update({"OracleSyncConnection": OracleSyncConnection, "OracleAsyncConnection": OracleAsyncConnection})
+        return namespace

sqlspec/adapters/psqlpy/config.py

@@ -324,7 +324,7 @@ def pool_config_dict(self) -> dict[str, Any]:
 
         return config
 
-    async def _create_pool(self) -> ConnectionPool:
+    async def _create_pool(self) -> "ConnectionPool":
         """Create the actual async connection pool."""
         logger.info("Creating psqlpy connection pool", extra={"adapter": "psqlpy"})
 
@@ -351,7 +351,7 @@ async def _close_pool(self) -> None:
             logger.exception("Failed to close psqlpy connection pool", extra={"adapter": "psqlpy", "error": str(e)})
             raise
 
-    async def create_connection(self) -> PsqlpyConnection:
+    async def create_connection(self) -> "PsqlpyConnection":
         """Create a single async connection (not from pool).
 
         Returns:
@@ -413,3 +413,16 @@ async def provide_pool(self, *args: Any, **kwargs: Any) -> ConnectionPool:
         if not self.pool_instance:
             self.pool_instance = await self.create_pool()
         return self.pool_instance
+
+    def get_signature_namespace(self) -> "dict[str, type[Any]]":
+        """Get the signature namespace for Psqlpy types.
+
+        This provides all Psqlpy-specific types that Litestar needs to recognize
+        to avoid serialization attempts.
+
+        Returns:
+            Dictionary mapping type names to types.
+        """
+        namespace = super().get_signature_namespace()
+        namespace.update({"PsqlpyConnection": PsqlpyConnection})
+        return namespace

sqlspec/adapters/psycopg/config.py

@@ -400,6 +400,19 @@ def provide_pool(self, *args: Any, **kwargs: Any) -> "ConnectionPool":
             self.pool_instance = self.create_pool()
         return self.pool_instance
 
+    def get_signature_namespace(self) -> "dict[str, type[Any]]":
+        """Get the signature namespace for Psycopg types.
+
+        This provides all Psycopg-specific types that Litestar needs to recognize
+        to avoid serialization attempts.
+
+        Returns:
+            Dictionary mapping type names to types.
+        """
+        namespace = super().get_signature_namespace()
+        namespace.update({"PsycopgSyncConnection": PsycopgSyncConnection})
+        return namespace
+
 
 class PsycopgAsyncConfig(AsyncDatabaseConfig[PsycopgAsyncConnection, AsyncConnectionPool, PsycopgAsyncDriver]):
     """Configuration for Psycopg asynchronous database connections with direct field-based configuration."""
@@ -727,3 +740,16 @@ async def provide_pool(self, *args: Any, **kwargs: Any) -> "AsyncConnectionPool"
         if not self.pool_instance:
             self.pool_instance = await self.create_pool()
         return self.pool_instance
+
+    def get_signature_namespace(self) -> "dict[str, type[Any]]":
+        """Get the signature namespace for Psycopg async types.
+
+        This provides all Psycopg async-specific types that Litestar needs to recognize
+        to avoid serialization attempts.
+
+        Returns:
+            Dictionary mapping type names to types.
+        """
+        namespace = super().get_signature_namespace()
+        namespace.update({"PsycopgAsyncConnection": PsycopgAsyncConnection})
+        return namespace