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

sqlspec/adapters/psycopg/driver.py

@@ -1,32 +1,95 @@
 import logging
 from contextlib import asynccontextmanager, contextmanager
-from typing import TYPE_CHECKING, Any, Optional, Union, cast
+from typing import TYPE_CHECKING, Any, Optional, Union, cast, overload
 
+from psycopg import AsyncConnection, Connection
 from psycopg.rows import dict_row
 
-from sqlspec.base import AsyncDriverAdapterProtocol, SyncDriverAdapterProtocol, T
+from sqlspec.base import AsyncDriverAdapterProtocol, SyncDriverAdapterProtocol
 from sqlspec.exceptions import SQLParsingError
+from sqlspec.mixins import SQLTranslatorMixin
 from sqlspec.statement import PARAM_REGEX, SQLStatement
 
 if TYPE_CHECKING:
-    from collections.abc import AsyncGenerator, Generator
+    from collections.abc import AsyncGenerator, Generator, Sequence
 
-    from psycopg import AsyncConnection, Connection
-
-    from sqlspec.typing import ModelDTOT, StatementParameterType
+    from sqlspec.typing import ModelDTOT, StatementParameterType, T
 
 logger = logging.getLogger("sqlspec")
 
-__all__ = ("PsycopgAsyncDriver", "PsycopgSyncDriver")
+__all__ = ("PsycopgAsyncConnection", "PsycopgAsyncDriver", "PsycopgSyncConnection", "PsycopgSyncDriver")
+
+PsycopgSyncConnection = Connection
+PsycopgAsyncConnection = AsyncConnection
+
+
+class PsycopgDriverBase:
+    dialect: str
+
+    def _process_sql_params(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        **kwargs: Any,
+    ) -> "tuple[str, Optional[Union[tuple[Any, ...], list[Any], dict[str, Any]]]]":
+        """Process SQL and parameters, converting :name -> %(name)s if needed."""
+        stmt = SQLStatement(sql=sql, parameters=parameters, dialect=self.dialect, kwargs=kwargs or None)
+        processed_sql, processed_params = stmt.process()
+
+        if isinstance(processed_params, dict):
+            parameter_dict = processed_params
+            processed_sql_parts: list[str] = []
+            last_end = 0
+            found_params_regex: list[str] = []
+
+            for match in PARAM_REGEX.finditer(processed_sql):
+                if match.group("dquote") or match.group("squote") or match.group("comment"):
+                    continue
+
+                if match.group("var_name"):
+                    var_name = match.group("var_name")
+                    found_params_regex.append(var_name)
+                    start = match.start("var_name") - 1
+                    end = match.end("var_name")
+
+                    if var_name not in parameter_dict:
+                        msg = (
+                            f"Named parameter ':{var_name}' found in SQL but missing from processed parameters. "
+                            f"Processed SQL: {processed_sql}"
+                        )
+                        raise SQLParsingError(msg)
+
+                    processed_sql_parts.extend((processed_sql[last_end:start], f"%({var_name})s"))
+                    last_end = end
 
+            processed_sql_parts.append(processed_sql[last_end:])
+            final_sql = "".join(processed_sql_parts)
 
-class PsycopgSyncDriver(SyncDriverAdapterProtocol["Connection"]):
+            if not found_params_regex and parameter_dict:
+                logger.warning(
+                    "Dict params provided (%s), but no :name placeholders found. SQL: %s",
+                    list(parameter_dict.keys()),
+                    processed_sql,
+                )
+                return processed_sql, parameter_dict
+
+            return final_sql, parameter_dict
+
+        return processed_sql, processed_params
+
+
+class PsycopgSyncDriver(
+    PsycopgDriverBase,
+    SQLTranslatorMixin["PsycopgSyncConnection"],
+    SyncDriverAdapterProtocol["PsycopgSyncConnection"],
+):
     """Psycopg Sync Driver Adapter."""
 
-    connection: "Connection"
+    connection: "PsycopgSyncConnection"
     dialect: str = "postgres"
 
-    def __init__(self, connection: "Connection") -> None:
+    def __init__(self, connection: "PsycopgSyncConnection") -> None:
         self.connection = connection
 
     def _process_sql_params(
@@ -36,7 +99,6 @@ class PsycopgSyncDriver(SyncDriverAdapterProtocol["Connection"]):
         /,
         **kwargs: Any,
     ) -> "tuple[str, Optional[Union[tuple[Any, ...], list[Any], dict[str, Any]]]]":
-        """Process SQL and parameters, converting :name -> %(name)s if needed."""
         stmt = SQLStatement(sql=sql, parameters=parameters, dialect=self.dialect, kwargs=kwargs or None)
         processed_sql, processed_params = stmt.process()
 
@@ -83,13 +145,36 @@ class PsycopgSyncDriver(SyncDriverAdapterProtocol["Connection"]):
 
     @staticmethod
     @contextmanager
-    def _with_cursor(connection: "Connection") -> "Generator[Any, None, None]":
+    def _with_cursor(connection: "PsycopgSyncConnection") -> "Generator[Any, None, None]":
         cursor = connection.cursor(row_factory=dict_row)
         try:
             yield cursor
         finally:
             cursor.close()
 
+    # --- Public API Methods --- #
+    @overload
+    def select(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[PsycopgSyncConnection]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "Sequence[dict[str, Any]]": ...
+    @overload
+    def select(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[PsycopgSyncConnection]" = None,
+        schema_type: "type[ModelDTOT]",
+        **kwargs: Any,
+    ) -> "Sequence[ModelDTOT]": ...
     def select(
         self,
         sql: str,
@@ -97,9 +182,9 @@ class PsycopgSyncDriver(SyncDriverAdapterProtocol["Connection"]):
         /,
         *,
         schema_type: "Optional[type[ModelDTOT]]" = None,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[PsycopgSyncConnection]" = None,
         **kwargs: Any,
-    ) -> "list[Union[ModelDTOT, dict[str, Any]]]":
+    ) -> "Sequence[Union[ModelDTOT, dict[str, Any]]]":
         """Fetch data from the database.
 
         Returns:
@@ -117,13 +202,35 @@ class PsycopgSyncDriver(SyncDriverAdapterProtocol["Connection"]):
                 return [cast("ModelDTOT", schema_type(**row)) for row in results]  # pyright: ignore[reportUnknownArgumentType]
             return [cast("dict[str,Any]", row) for row in results]  # pyright: ignore[reportUnknownArgumentType]
 
+    @overload
     def select_one(
         self,
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[PsycopgSyncConnection]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "dict[str, Any]": ...
+    @overload
+    def select_one(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[PsycopgSyncConnection]" = None,
+        schema_type: "type[ModelDTOT]",
+        **kwargs: Any,
+    ) -> "ModelDTOT": ...
+    def select_one(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[PsycopgSyncConnection]" = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
         **kwargs: Any,
     ) -> "Union[ModelDTOT, dict[str, Any]]":
@@ -142,13 +249,35 @@ class PsycopgSyncDriver(SyncDriverAdapterProtocol["Connection"]):
                 return cast("ModelDTOT", schema_type(**cast("dict[str,Any]", row)))
             return cast("dict[str,Any]", row)
 
+    @overload
+    def select_one_or_none(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[PsycopgSyncConnection]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "Optional[dict[str, Any]]": ...
+    @overload
+    def select_one_or_none(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[PsycopgSyncConnection]" = None,
+        schema_type: "type[ModelDTOT]",
+        **kwargs: Any,
+    ) -> "Optional[ModelDTOT]": ...
     def select_one_or_none(
         self,
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[PsycopgSyncConnection]" = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
         **kwargs: Any,
     ) -> "Optional[Union[ModelDTOT, dict[str, Any]]]":
@@ -168,13 +297,35 @@ class PsycopgSyncDriver(SyncDriverAdapterProtocol["Connection"]):
                 return cast("ModelDTOT", schema_type(**cast("dict[str,Any]", row)))
             return cast("dict[str,Any]", row)
 
+    @overload
     def select_value(
         self,
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[PsycopgSyncConnection]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "Any": ...
+    @overload
+    def select_value(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[PsycopgSyncConnection]" = None,
+        schema_type: "type[T]",
+        **kwargs: Any,
+    ) -> "T": ...
+    def select_value(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[PsycopgSyncConnection]" = None,
         schema_type: "Optional[type[T]]" = None,
         **kwargs: Any,
     ) -> "Union[T, Any]":
@@ -195,13 +346,35 @@ class PsycopgSyncDriver(SyncDriverAdapterProtocol["Connection"]):
                 return schema_type(val)  # type: ignore[call-arg]
             return val
 
+    @overload
     def select_value_or_none(
         self,
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[PsycopgSyncConnection]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "Optional[Any]": ...
+    @overload
+    def select_value_or_none(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[PsycopgSyncConnection]" = None,
+        schema_type: "type[T]",
+        **kwargs: Any,
+    ) -> "Optional[T]": ...
+    def select_value_or_none(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[PsycopgSyncConnection]" = None,
         schema_type: "Optional[type[T]]" = None,
         **kwargs: Any,
     ) -> "Optional[Union[T, Any]]":
@@ -230,7 +403,7 @@ class PsycopgSyncDriver(SyncDriverAdapterProtocol["Connection"]):
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[PsycopgSyncConnection]" = None,
         **kwargs: Any,
     ) -> int:
         """Execute an INSERT, UPDATE, or DELETE query and return the number of affected rows.
@@ -244,13 +417,35 @@ class PsycopgSyncDriver(SyncDriverAdapterProtocol["Connection"]):
             cursor.execute(sql, parameters)
             return getattr(cursor, "rowcount", -1)  # pyright: ignore[reportUnknownMemberType]
 
+    @overload
     def insert_update_delete_returning(
         self,
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[PsycopgSyncConnection]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "dict[str, Any]": ...
+    @overload
+    def insert_update_delete_returning(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[PsycopgSyncConnection]" = None,
+        schema_type: "type[ModelDTOT]",
+        **kwargs: Any,
+    ) -> "ModelDTOT": ...
+    def insert_update_delete_returning(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[PsycopgSyncConnection]" = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
         **kwargs: Any,
     ) -> "Optional[Union[dict[str, Any], ModelDTOT]]":
@@ -278,7 +473,7 @@ class PsycopgSyncDriver(SyncDriverAdapterProtocol["Connection"]):
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[PsycopgSyncConnection]" = None,
         **kwargs: Any,
     ) -> str:
         """Execute a script.
@@ -293,86 +488,61 @@ class PsycopgSyncDriver(SyncDriverAdapterProtocol["Connection"]):
             return str(cursor.statusmessage) if cursor.statusmessage is not None else "DONE"
 
 
-class PsycopgAsyncDriver(AsyncDriverAdapterProtocol["AsyncConnection"]):
+class PsycopgAsyncDriver(
+    PsycopgDriverBase,
+    SQLTranslatorMixin["PsycopgAsyncConnection"],
+    AsyncDriverAdapterProtocol["PsycopgAsyncConnection"],
+):
     """Psycopg Async Driver Adapter."""
 
-    connection: "AsyncConnection"
+    connection: "PsycopgAsyncConnection"
     dialect: str = "postgres"
 
-    def __init__(self, connection: "AsyncConnection") -> None:
+    def __init__(self, connection: "PsycopgAsyncConnection") -> None:
         self.connection = connection
 
-    def _process_sql_params(
-        self,
-        sql: str,
-        parameters: "Optional[StatementParameterType]" = None,
-        /,
-        **kwargs: Any,
-    ) -> "tuple[str, Optional[Union[tuple[Any, ...], list[Any], dict[str, Any]]]]":
-        """Process SQL and parameters, converting :name -> %(name)s if needed."""
-        stmt = SQLStatement(sql=sql, parameters=parameters, dialect=self.dialect, kwargs=kwargs or None)
-        processed_sql, processed_params = stmt.process()
-
-        if isinstance(processed_params, dict):
-            parameter_dict = processed_params
-            processed_sql_parts: list[str] = []
-            last_end = 0
-            found_params_regex: list[str] = []
-
-            for match in PARAM_REGEX.finditer(processed_sql):
-                if match.group("dquote") or match.group("squote") or match.group("comment"):
-                    continue
-
-                if match.group("var_name"):
-                    var_name = match.group("var_name")
-                    found_params_regex.append(var_name)
-                    start = match.start("var_name") - 1
-                    end = match.end("var_name")
-
-                    if var_name not in parameter_dict:
-                        msg = (
-                            f"Named parameter ':{var_name}' found in SQL but missing from processed parameters. "
-                            f"Processed SQL: {processed_sql}"
-                        )
-                        raise SQLParsingError(msg)
-
-                    processed_sql_parts.extend((processed_sql[last_end:start], f"%({var_name})s"))
-                    last_end = end
-
-            processed_sql_parts.append(processed_sql[last_end:])
-            final_sql = "".join(processed_sql_parts)
-
-            if not found_params_regex and parameter_dict:
-                logger.warning(
-                    "Dict params provided (%s), but no :name placeholders found. SQL: %s",
-                    list(parameter_dict.keys()),
-                    processed_sql,
-                )
-                return processed_sql, parameter_dict
-
-            return final_sql, parameter_dict
-
-        return processed_sql, processed_params
-
     @staticmethod
     @asynccontextmanager
-    async def _with_cursor(connection: "AsyncConnection") -> "AsyncGenerator[Any, None]":
+    async def _with_cursor(connection: "PsycopgAsyncConnection") -> "AsyncGenerator[Any, None]":
         cursor = connection.cursor(row_factory=dict_row)
         try:
             yield cursor
         finally:
             await cursor.close()
 
+    # --- Public API Methods --- #
+    @overload
     async def select(
         self,
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[AsyncConnection]" = None,
+        connection: "Optional[PsycopgAsyncConnection]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "Sequence[dict[str, Any]]": ...
+    @overload
+    async def select(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[PsycopgAsyncConnection]" = None,
+        schema_type: "type[ModelDTOT]",
+        **kwargs: Any,
+    ) -> "Sequence[ModelDTOT]": ...
+    async def select(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[PsycopgAsyncConnection]" = 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.
 
         Returns:
@@ -391,13 +561,35 @@ class PsycopgAsyncDriver(AsyncDriverAdapterProtocol["AsyncConnection"]):
                 return [cast("ModelDTOT", schema_type(**cast("dict[str,Any]", row))) for row in results]  # pyright: ignore[reportUnknownArgumentType]
             return [cast("dict[str,Any]", row) for row in results]  # pyright: ignore[reportUnknownArgumentType]
 
+    @overload
     async def select_one(
         self,
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[AsyncConnection]" = None,
+        connection: "Optional[PsycopgAsyncConnection]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "dict[str, Any]": ...
+    @overload
+    async def select_one(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[PsycopgAsyncConnection]" = None,
+        schema_type: "type[ModelDTOT]",
+        **kwargs: Any,
+    ) -> "ModelDTOT": ...
+    async def select_one(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[PsycopgAsyncConnection]" = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
         **kwargs: Any,
     ) -> "Union[ModelDTOT, dict[str, Any]]":
@@ -417,6 +609,28 @@ class PsycopgAsyncDriver(AsyncDriverAdapterProtocol["AsyncConnection"]):
                 return cast("ModelDTOT", schema_type(**cast("dict[str,Any]", row)))
             return cast("dict[str,Any]", row)
 
+    @overload
+    async def select_one_or_none(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[PsycopgAsyncConnection]" = 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[PsycopgAsyncConnection]" = None,
+        schema_type: "type[ModelDTOT]",
+        **kwargs: Any,
+    ) -> "Optional[ModelDTOT]": ...
     async def select_one_or_none(
         self,
         sql: str,
@@ -424,7 +638,7 @@ class PsycopgAsyncDriver(AsyncDriverAdapterProtocol["AsyncConnection"]):
         /,
         *,
         schema_type: "Optional[type[ModelDTOT]]" = None,
-        connection: "Optional[AsyncConnection]" = None,
+        connection: "Optional[PsycopgAsyncConnection]" = None,
         **kwargs: Any,
     ) -> "Optional[Union[ModelDTOT, dict[str, Any]]]":
         """Fetch one row from the database.
@@ -444,13 +658,35 @@ class PsycopgAsyncDriver(AsyncDriverAdapterProtocol["AsyncConnection"]):
                 return cast("ModelDTOT", schema_type(**cast("dict[str,Any]", row)))
             return cast("dict[str,Any]", row)
 
+    @overload
+    async def select_value(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[PsycopgAsyncConnection]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "Any": ...
+    @overload
+    async def select_value(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[PsycopgAsyncConnection]" = None,
+        schema_type: "type[T]",
+        **kwargs: Any,
+    ) -> "T": ...
     async def select_value(
         self,
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[AsyncConnection]" = None,
+        connection: "Optional[PsycopgAsyncConnection]" = None,
         schema_type: "Optional[type[T]]" = None,
         **kwargs: Any,
     ) -> "Union[T, Any]":
@@ -478,7 +714,7 @@ class PsycopgAsyncDriver(AsyncDriverAdapterProtocol["AsyncConnection"]):
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[AsyncConnection]" = None,
+        connection: "Optional[PsycopgAsyncConnection]" = None,
         schema_type: "Optional[type[T]]" = None,
         **kwargs: Any,
     ) -> "Optional[Union[T, Any]]":
@@ -508,7 +744,7 @@ class PsycopgAsyncDriver(AsyncDriverAdapterProtocol["AsyncConnection"]):
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[AsyncConnection]" = None,
+        connection: "Optional[PsycopgAsyncConnection]" = None,
         **kwargs: Any,
     ) -> int:
         """Execute an INSERT, UPDATE, or DELETE query and return the number of affected rows.
@@ -527,13 +763,35 @@ class PsycopgAsyncDriver(AsyncDriverAdapterProtocol["AsyncConnection"]):
                 rowcount = -1
             return rowcount
 
+    @overload
+    async def insert_update_delete_returning(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[PsycopgAsyncConnection]" = 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[PsycopgAsyncConnection]" = None,
+        schema_type: "type[ModelDTOT]",
+        **kwargs: Any,
+    ) -> "ModelDTOT": ...
     async def insert_update_delete_returning(
         self,
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[AsyncConnection]" = None,
+        connection: "Optional[PsycopgAsyncConnection]" = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
         **kwargs: Any,
     ) -> "Optional[Union[dict[str, Any], ModelDTOT]]":
@@ -562,7 +820,7 @@ class PsycopgAsyncDriver(AsyncDriverAdapterProtocol["AsyncConnection"]):
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[AsyncConnection]" = None,
+        connection: "Optional[PsycopgAsyncConnection]" = None,
         **kwargs: Any,
     ) -> str:
         """Execute a script.

sqlspec/adapters/sqlite/__init__.py

@@ -1,7 +1,8 @@
 from sqlspec.adapters.sqlite.config import SqliteConfig
-from sqlspec.adapters.sqlite.driver import SqliteDriver
+from sqlspec.adapters.sqlite.driver import SqliteConnection, SqliteDriver
 
 __all__ = (
     "SqliteConfig",
+    "SqliteConnection",
     "SqliteDriver",
 )