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

sqlspec/adapters/asyncpg/config.py

@@ -1,13 +1,13 @@
 from contextlib import asynccontextmanager
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, Optional, TypeVar, Union
 
 from asyncpg import Record
 from asyncpg import create_pool as asyncpg_create_pool
-from asyncpg.pool import Pool, PoolConnectionProxy
-from typing_extensions import TypeAlias
+from asyncpg.pool import PoolConnectionProxy
 
 from sqlspec._serialization import decode_json, encode_json
+from sqlspec.adapters.asyncpg.driver import AsyncpgConnection, AsyncpgDriver
 from sqlspec.base import AsyncDatabaseConfig, GenericPoolConfig
 from sqlspec.exceptions import ImproperConfigurationError
 from sqlspec.typing import Empty, EmptyType, dataclass_to_dict
@@ -17,21 +17,20 @@ if TYPE_CHECKING:
     from collections.abc import AsyncGenerator, Awaitable, Callable, Coroutine
 
     from asyncpg.connection import Connection
+    from asyncpg.pool import Pool
 
 
 __all__ = (
-    "AsyncPgConfig",
-    "AsyncPgPoolConfig",
+    "AsyncpgConfig",
+    "AsyncpgPoolConfig",
 )
 
 
 T = TypeVar("T")
 
-PgConnection: TypeAlias = "Union[Connection, PoolConnectionProxy]"  # pyright: ignore[reportMissingTypeArgument]
-
 
 @dataclass
-class AsyncPgPoolConfig(GenericPoolConfig):
+class AsyncpgPoolConfig(GenericPoolConfig):
     """Configuration for Asyncpg's :class:`Pool <asyncpg.pool.Pool>`.
 
     For details see: https://magicstack.github.io/asyncpg/current/api/index.html#connection-pools
@@ -52,7 +51,7 @@ class AsyncPgPoolConfig(GenericPoolConfig):
     min_size: "Union[int, EmptyType]" = Empty
     """The number of connections to keep open inside the connection pool."""
     max_size: "Union[int, EmptyType]" = Empty
-    """The number of connections to allow in connection pool “overflow”, that is connections that can be opened above
+    """The number of connections to allow in connection pool "overflow", that is connections that can be opened above
     and beyond the pool_size setting, which defaults to 10."""
 
     max_queries: "Union[int, EmptyType]" = Empty
@@ -71,23 +70,55 @@ class AsyncPgPoolConfig(GenericPoolConfig):
 
 
 @dataclass
-class AsyncPgConfig(AsyncDatabaseConfig[PgConnection, Pool]):  # pyright: ignore[reportMissingTypeArgument]
+class AsyncpgConfig(AsyncDatabaseConfig["AsyncpgConnection", "Pool", "AsyncpgDriver"]):  # pyright: ignore[reportMissingTypeArgument]
     """Asyncpg Configuration."""
 
-    pool_config: "Optional[AsyncPgPoolConfig]" = None
+    pool_config: "Optional[AsyncpgPoolConfig]" = field(default=None)
     """Asyncpg Pool configuration"""
-    json_deserializer: "Callable[[str], Any]" = decode_json
+    json_deserializer: "Callable[[str], Any]" = field(hash=False, default=decode_json)
     """For dialects that support the :class:`JSON <sqlalchemy.types.JSON>` datatype, this is a Python callable that will
     convert a JSON string to a Python object. By default, this is set to SQLSpec's
     :attr:`decode_json() <sqlspec._serialization.decode_json>` function."""
-    json_serializer: "Callable[[Any], str]" = encode_json
+    json_serializer: "Callable[[Any], str]" = field(hash=False, default=encode_json)
     """For dialects that support the JSON datatype, this is a Python callable that will render a given object as JSON.
     By default, SQLSpec's :attr:`encode_json() <sqlspec._serialization.encode_json>` is used."""
-    pool_instance: "Optional[Pool[Any]]" = None
-    """Optional pool to use.
+    connection_type: "type[AsyncpgConnection]" = field(
+        hash=False, init=False, default_factory=lambda: PoolConnectionProxy
+    )
+    """Type of the connection object"""
+    driver_type: "type[AsyncpgDriver]" = field(hash=False, init=False, default_factory=lambda: AsyncpgDriver)  # type: ignore[type-abstract,unused-ignore]
+    """Type of the driver object"""
+    pool_instance: "Optional[Pool[Any]]" = field(hash=False, default=None)
+    """The connection pool instance. If set, this will be used instead of creating a new pool."""
 
-    If set, the plugin will use the provided pool rather than instantiate one.
-    """
+    @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 asyncpg.connect function.
+
+        Raises:
+            ImproperConfigurationError: If the connection configuration is not provided.
+        """
+        if self.pool_config:
+            connect_dict: dict[str, Any] = {}
+
+            # Add dsn if available
+            if hasattr(self.pool_config, "dsn"):
+                connect_dict["dsn"] = self.pool_config.dsn
+
+            # Add any connect_kwargs if available
+            if (
+                hasattr(self.pool_config, "connect_kwargs")
+                and self.pool_config.connect_kwargs is not Empty
+                and isinstance(self.pool_config.connect_kwargs, dict)
+            ):
+                connect_dict.update(dict(self.pool_config.connect_kwargs.items()))
+
+            return connect_dict
+        msg = "You must provide a 'pool_config' for this adapter."
+        raise ImproperConfigurationError(msg)
 
     @property
     def pool_config_dict(self) -> "dict[str, Any]":
@@ -96,9 +127,17 @@ class AsyncPgConfig(AsyncDatabaseConfig[PgConnection, Pool]):  # pyright: ignore
         Returns:
             A string keyed dict of config kwargs for the Asyncpg :func:`create_pool <asyncpg.pool.create_pool>`
             function.
+
+        Raises:
+            ImproperConfigurationError: If no pool_config is provided but a pool_instance is set.
         """
         if self.pool_config:
-            return dataclass_to_dict(self.pool_config, exclude_empty=True, convert_nested=False)
+            return dataclass_to_dict(
+                self.pool_config,
+                exclude_empty=True,
+                exclude={"pool_instance", "driver_type", "connection_type"},
+                convert_nested=False,
+            )
         msg = "'pool_config' methods can not be used when a 'pool_instance' is provided."
         raise ImproperConfigurationError(msg)
 
@@ -107,6 +146,10 @@ class AsyncPgConfig(AsyncDatabaseConfig[PgConnection, Pool]):  # pyright: ignore
 
         Returns:
             Getter that returns the pool instance used by the plugin.
+
+        Raises:
+            ImproperConfigurationError: If neither pool_config nor pool_instance are provided,
+                or if the pool could not be configured.
         """
         if self.pool_instance is not None:
             return self.pool_instance
@@ -117,11 +160,9 @@ class AsyncPgConfig(AsyncDatabaseConfig[PgConnection, Pool]):  # pyright: ignore
 
         pool_config = self.pool_config_dict
         self.pool_instance = await asyncpg_create_pool(**pool_config)
-        if self.pool_instance is None:
-            msg = "Could not configure the 'pool_instance'. Please check your configuration."
-            raise ImproperConfigurationError(
-                msg,
-            )
+        if self.pool_instance is None:  # pyright: ignore[reportUnnecessaryComparison]
+            msg = "Could not configure the 'pool_instance'. Please check your configuration."  # type: ignore[unreachable]
+            raise ImproperConfigurationError(msg)
         return self.pool_instance
 
     def provide_pool(self, *args: "Any", **kwargs: "Any") -> "Awaitable[Pool]":  # pyright: ignore[reportMissingTypeArgument,reportUnknownParameterType]
@@ -132,13 +173,49 @@ class AsyncPgConfig(AsyncDatabaseConfig[PgConnection, Pool]):  # pyright: ignore
         """
         return self.create_pool()  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType]
 
+    async def create_connection(self) -> "AsyncpgConnection":
+        """Create and return a new asyncpg connection from the pool.
+
+        Returns:
+            A Connection instance.
+
+        Raises:
+            ImproperConfigurationError: If the connection could not be created.
+        """
+        try:
+            pool = await self.provide_pool()
+            return await pool.acquire()
+        except Exception as e:
+            msg = f"Could not configure the asyncpg connection. Error: {e!s}"
+            raise ImproperConfigurationError(msg) from e
+
     @asynccontextmanager
-    async def provide_connection(self, *args: "Any", **kwargs: "Any") -> "AsyncGenerator[PoolConnectionProxy, None]":  # pyright: ignore[reportMissingTypeArgument,reportUnknownParameterType]
+    async def provide_connection(
+        self, *args: "Any", **kwargs: "Any"
+    ) -> "AsyncGenerator[PoolConnectionProxy[Any], None]":  # pyright: ignore[reportMissingTypeArgument,reportUnknownParameterType]
         """Create a connection instance.
 
-        Returns:
+        Yields:
             A connection instance.
         """
         db_pool = await self.provide_pool(*args, **kwargs)  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType]
         async with db_pool.acquire() as connection:  # pyright: ignore[reportUnknownVariableType]
             yield connection
+
+    async def close_pool(self) -> None:
+        """Close the pool."""
+        if self.pool_instance is not None:
+            await self.pool_instance.close()
+            self.pool_instance = None
+
+    @asynccontextmanager
+    async def provide_session(self, *args: Any, **kwargs: Any) -> "AsyncGenerator[AsyncpgDriver, None]":
+        """Create and provide a database session.
+
+        Yields:
+            A Aiosqlite driver instance.
+
+
+        """
+        async with self.provide_connection(*args, **kwargs) as connection:
+            yield self.driver_type(connection)

sqlspec/adapters/asyncpg/driver.py

@@ -0,0 +1,444 @@
+import logging
+import re
+from typing import TYPE_CHECKING, Any, Optional, Union, cast
+
+from asyncpg import Connection
+from typing_extensions import TypeAlias
+
+from sqlspec.base import AsyncDriverAdapterProtocol, T
+from sqlspec.exceptions import SQLParsingError
+from sqlspec.statement import PARAM_REGEX, SQLStatement
+
+if TYPE_CHECKING:
+    from asyncpg.connection import Connection
+    from asyncpg.pool import PoolConnectionProxy
+
+    from sqlspec.typing import ModelDTOT, StatementParameterType
+
+__all__ = ("AsyncpgConnection", "AsyncpgDriver")
+
+logger = logging.getLogger("sqlspec")
+
+# Regex to find '?' placeholders, skipping those inside quotes or SQL comments
+# Simplified version, assumes standard SQL quoting/comments
+QMARK_REGEX = re.compile(
+    r"""(?P<dquote>"[^"]*") | # Double-quoted strings
+         (?P<squote>\'[^\']*\') | # Single-quoted strings
+         (?P<comment>--[^\n]*|/\*.*?\*/) | # SQL comments (single/multi-line)
+         (?P<qmark>\?) # The question mark placeholder
+      """,
+    re.VERBOSE | re.DOTALL,
+)
+
+AsyncpgConnection: TypeAlias = "Union[Connection[Any], PoolConnectionProxy[Any]]"  # pyright: ignore[reportMissingTypeArgument]
+
+
+class AsyncpgDriver(AsyncDriverAdapterProtocol["AsyncpgConnection"]):
+    """AsyncPG Postgres Driver Adapter."""
+
+    connection: "AsyncpgConnection"
+    dialect: str = "postgres"
+
+    def __init__(self, connection: "AsyncpgConnection") -> None:
+        self.connection = connection
+
+    def _process_sql_params(  # noqa: C901, PLR0912, PLR0915
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        **kwargs: Any,
+    ) -> "tuple[str, Optional[Union[tuple[Any, ...], list[Any], dict[str, Any]]]]":
+        # Use SQLStatement for parameter validation and merging first
+        # It also handles potential dialect-specific logic if implemented there.
+        stmt = SQLStatement(sql=sql, parameters=parameters, dialect=self.dialect, kwargs=kwargs or None)
+        sql, parameters = stmt.process()
+
+        # Case 1: Parameters are effectively a dictionary (either passed as dict or via kwargs merged by SQLStatement)
+        if isinstance(parameters, dict):
+            processed_sql_parts: list[str] = []
+            ordered_params = []
+            last_end = 0
+            param_index = 1
+            found_params_regex: list[str] = []
+
+            # Manually parse the PROCESSED SQL for :name -> $n conversion
+            for match in PARAM_REGEX.finditer(sql):
+                # Skip matches inside quotes or comments
+                if match.group("dquote") or match.group("squote") or match.group("comment"):
+                    continue
+
+                if match.group("var_name"):  # Finds :var_name
+                    var_name = match.group("var_name")
+                    found_params_regex.append(var_name)
+                    start = match.start("var_name") - 1  # Include the ':'
+                    end = match.end("var_name")
+
+                    # SQLStatement should have already validated parameter existence,
+                    # but we double-check here during ordering.
+                    if var_name not in parameters:
+                        # This should ideally not happen if SQLStatement validation is robust.
+                        msg = (
+                            f"Named parameter ':{var_name}' found in SQL but missing from processed parameters. "
+                            f"Processed SQL: {sql}"
+                        )
+                        raise SQLParsingError(msg)
+
+                    processed_sql_parts.extend((sql[last_end:start], f"${param_index}"))
+                    ordered_params.append(parameters[var_name])
+                    last_end = end
+                    param_index += 1
+
+            processed_sql_parts.append(sql[last_end:])
+            final_sql = "".join(processed_sql_parts)
+
+            # --- Validation ---
+            # Check if named placeholders were found if dict params were provided
+            # SQLStatement might handle this validation, but a warning here can be useful.
+            if not found_params_regex and parameters:
+                logger.warning(
+                    "Dict params provided (%s), but no :name placeholders found. SQL: %s",
+                    list(parameters.keys()),
+                    sql,
+                )
+                # If no placeholders, return original SQL from SQLStatement and empty tuple for asyncpg
+                return sql, ()
+
+            # Additional checks (potentially redundant if SQLStatement covers them):
+            # 1. Ensure all found placeholders have corresponding params (covered by check inside loop)
+            # 2. Ensure all provided params correspond to a placeholder
+            provided_keys = set(parameters.keys())
+            found_keys = set(found_params_regex)
+            unused_keys = provided_keys - found_keys
+            if unused_keys:
+                # SQLStatement might handle this, but log a warning just in case.
+                logger.warning(
+                    "Parameters provided but not used in SQL: %s. SQL: %s",
+                    unused_keys,
+                    sql,
+                )
+
+            return final_sql, tuple(ordered_params)  # asyncpg expects a sequence
+
+        # Case 2: Parameters are effectively a sequence/scalar (merged by SQLStatement)
+        if isinstance(parameters, (list, tuple)):
+            # Parameters are a sequence, need to convert ? -> $n
+            sequence_processed_parts: list[str] = []
+            param_index = 1
+            last_end = 0
+            qmark_found = False
+
+            # Manually parse the PROCESSED SQL to find '?' outside comments/quotes and convert to $n
+            for match in QMARK_REGEX.finditer(sql):
+                if match.group("dquote") or match.group("squote") or match.group("comment"):
+                    continue  # Skip quotes and comments
+
+                if match.group("qmark"):
+                    qmark_found = True
+                    start = match.start("qmark")
+                    end = match.end("qmark")
+                    sequence_processed_parts.extend((sql[last_end:start], f"${param_index}"))
+                    last_end = end
+                    param_index += 1
+
+            sequence_processed_parts.append(sql[last_end:])
+            final_sql = "".join(sequence_processed_parts)
+
+            # --- Validation ---
+            # Check if '?' was found if parameters were provided
+            if parameters and not qmark_found:
+                # SQLStatement might allow this, log a warning.
+                logger.warning(
+                    "Sequence/scalar parameters provided, but no '?' placeholders found. SQL: %s",
+                    sql,
+                )
+                # Return PROCESSED SQL from SQLStatement as no conversion happened here
+                return sql, parameters
+
+            # Check parameter count match (using count from manual parsing vs count from stmt)
+            expected_params = param_index - 1
+            actual_params = len(parameters)
+            if expected_params != actual_params:
+                msg = (
+                    f"Parameter count mismatch: Processed SQL expected {expected_params} parameters ('$n'), "
+                    f"but {actual_params} were provided by SQLStatement. "
+                    f"Final Processed SQL: {final_sql}"
+                )
+                raise SQLParsingError(msg)
+
+            return final_sql, parameters
+
+        # Case 3: Parameters are None (as determined by SQLStatement)
+        # processed_params is None
+        # Check if the SQL contains any placeholders unexpectedly
+        # Check for :name style
+        named_placeholders_found = False
+        for match in PARAM_REGEX.finditer(sql):
+            if not (match.group("dquote") or match.group("squote") or match.group("comment")) and match.group(
+                "var_name"
+            ):
+                named_placeholders_found = True
+                break
+        if named_placeholders_found:
+            msg = f"Processed SQL contains named parameters (:name) but no parameters were provided. SQL: {sql}"
+            raise SQLParsingError(msg)
+
+        # Check for ? style
+        qmark_placeholders_found = False
+        for match in QMARK_REGEX.finditer(sql):
+            if not (match.group("dquote") or match.group("squote") or match.group("comment")) and match.group("qmark"):
+                qmark_placeholders_found = True
+                break
+        if qmark_placeholders_found:
+            msg = f"Processed SQL contains positional parameters (?) but no parameters were provided. SQL: {sql}"
+            raise SQLParsingError(msg)
+
+        # No parameters provided and none found in SQL, return original SQL from SQLStatement and empty tuple
+        return sql, ()  # asyncpg expects a sequence, even if empty
+
+    async def select(
+        self,
+        sql: str,
+        parameters: Optional["StatementParameterType"] = None,
+        /,
+        *,
+        connection: Optional["AsyncpgConnection"] = None,
+        schema_type: "Optional[type[ModelDTOT]]" = None,
+        **kwargs: Any,
+    ) -> "list[Union[ModelDTOT, dict[str, Any]]]":
+        """Fetch data from the database.
+
+        Args:
+            sql: SQL statement.
+            parameters: Query parameters.
+            connection: Optional connection to use.
+            schema_type: Optional schema class for the result.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            List of row data as either model instances or dictionaries.
+        """
+        connection = self._connection(connection)
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
+        parameters = parameters if parameters is not None else {}
+
+        results = await connection.fetch(sql, *parameters)  # pyright: ignore
+        if not results:
+            return []
+        if schema_type is None:
+            return [dict(row.items()) for row in results]  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
+        return [cast("ModelDTOT", schema_type(**dict(row.items()))) for row in results]  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
+
+    async def select_one(
+        self,
+        sql: str,
+        parameters: Optional["StatementParameterType"] = None,
+        /,
+        *,
+        connection: Optional["AsyncpgConnection"] = None,
+        schema_type: "Optional[type[ModelDTOT]]" = None,
+        **kwargs: Any,
+    ) -> "Union[ModelDTOT, dict[str, Any]]":
+        """Fetch one row from the database.
+
+        Args:
+            sql: SQL statement.
+            parameters: Query parameters.
+            connection: Optional connection to use.
+            schema_type: Optional schema class for the result.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            The first row of the query results.
+        """
+        connection = self._connection(connection)
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
+        parameters = parameters if parameters is not None else {}
+        result = await connection.fetchrow(sql, *parameters)  # pyright: ignore
+        result = self.check_not_found(result)
+
+        if schema_type is None:
+            # Always return as dictionary
+            return dict(result.items())  # type: ignore[attr-defined]
+        return cast("ModelDTOT", schema_type(**dict(result.items())))  # type: ignore[attr-defined]
+
+    async def select_one_or_none(
+        self,
+        sql: str,
+        parameters: Optional["StatementParameterType"] = None,
+        /,
+        *,
+        connection: Optional["AsyncpgConnection"] = None,
+        schema_type: "Optional[type[ModelDTOT]]" = None,
+        **kwargs: Any,
+    ) -> "Optional[Union[ModelDTOT, dict[str, Any]]]":
+        """Fetch one row from the database.
+
+        Args:
+            sql: SQL statement.
+            parameters: Query parameters.
+            connection: Optional connection to use.
+            schema_type: Optional schema class for the result.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            The first row of the query results.
+        """
+        connection = self._connection(connection)
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
+        parameters = parameters if parameters is not None else {}
+        result = await connection.fetchrow(sql, *parameters)  # pyright: ignore
+        if result is None:
+            return None
+        if schema_type is None:
+            # Always return as dictionary
+            return dict(result.items())
+        return cast("ModelDTOT", schema_type(**dict(result.items())))
+
+    async def select_value(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AsyncpgConnection]" = None,
+        schema_type: "Optional[type[T]]" = None,
+        **kwargs: Any,
+    ) -> "Union[T, Any]":
+        """Fetch a single value from the database.
+
+        Args:
+            sql: SQL statement.
+            parameters: Query parameters.
+            connection: Optional connection to use.
+            schema_type: Optional schema class for the result.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            The first value from the first row of results, or None if no results.
+        """
+        connection = self._connection(connection)
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
+        parameters = parameters if parameters is not None else {}
+        result = await connection.fetchval(sql, *parameters)  # pyright: ignore
+        result = self.check_not_found(result)
+        if schema_type is None:
+            return result
+        return schema_type(result)  # type: ignore[call-arg]
+
+    async def select_value_or_none(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AsyncpgConnection]" = None,
+        schema_type: "Optional[type[T]]" = None,
+        **kwargs: Any,
+    ) -> "Optional[Union[T, Any]]":
+        """Fetch a single value from the database.
+
+        Returns:
+            The first value from the first row of results, or None if no results.
+        """
+        connection = self._connection(connection)
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
+        parameters = parameters if parameters is not None else {}
+        result = await connection.fetchval(sql, *parameters)  # pyright: ignore
+        if result is None:
+            return None
+        if schema_type is None:
+            return result
+        return schema_type(result)  # type: ignore[call-arg]
+
+    async def insert_update_delete(
+        self,
+        sql: str,
+        parameters: Optional["StatementParameterType"] = None,
+        /,
+        *,
+        connection: Optional["AsyncpgConnection"] = None,
+        **kwargs: Any,
+    ) -> int:
+        """Insert, update, or delete data from the database.
+
+        Args:
+            sql: SQL statement.
+            parameters: Query parameters.
+            connection: Optional connection to use.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            Row count affected by the operation.
+        """
+        connection = self._connection(connection)
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
+        parameters = parameters if parameters is not None else {}
+        status = await connection.execute(sql, *parameters)  # pyright: ignore
+        # AsyncPG returns a string like "INSERT 0 1" where the last number is the affected rows
+        try:
+            return int(status.split()[-1])  # pyright: ignore[reportUnknownMemberType]
+        except (ValueError, IndexError, AttributeError):
+            return -1  # Fallback if we can't parse the status
+
+    async def insert_update_delete_returning(
+        self,
+        sql: str,
+        parameters: Optional["StatementParameterType"] = None,
+        /,
+        *,
+        connection: Optional["AsyncpgConnection"] = None,
+        schema_type: "Optional[type[ModelDTOT]]" = None,
+        **kwargs: Any,
+    ) -> "Optional[Union[dict[str, Any], ModelDTOT]]":
+        """Insert, update, or delete data from the database and return the affected row.
+
+        Args:
+            sql: SQL statement.
+            parameters: Query parameters.
+            connection: Optional connection to use.
+            schema_type: Optional schema class for the result.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            The affected row data as either a model instance or dictionary.
+        """
+        connection = self._connection(connection)
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
+        parameters = parameters if parameters is not None else {}
+        result = await connection.fetchrow(sql, *parameters)  # pyright: ignore
+        if result is None:
+            return None
+        if schema_type is None:
+            # Always return as dictionary
+            return dict(result.items())  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
+        return cast("ModelDTOT", schema_type(**dict(result.items())))  # pyright: ignore[reportUnknownArgumentType, reportUnknownMemberType, reportUnknownVariableType]
+
+    async def execute_script(
+        self,
+        sql: str,
+        parameters: Optional["StatementParameterType"] = None,
+        /,
+        *,
+        connection: Optional["AsyncpgConnection"] = None,
+        **kwargs: Any,
+    ) -> str:
+        """Execute a script.
+
+        Args:
+            sql: SQL statement.
+            parameters: Query parameters.
+            connection: Optional connection to use.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            Status message for the operation.
+        """
+        connection = self._connection(connection)
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
+        parameters = parameters if parameters is not None else {}
+        return await connection.execute(sql, *parameters)  # pyright: ignore
+
+    def _connection(self, connection: Optional["AsyncpgConnection"] = None) -> "AsyncpgConnection":
+        """Return the connection to use. If None, use the default connection."""
+        return connection if connection is not None else self.connection

sqlspec/adapters/duckdb/__init__.py

@@ -1,3 +1,7 @@
 from sqlspec.adapters.duckdb.config import DuckDBConfig
+from sqlspec.adapters.duckdb.driver import DuckDBDriver
 
-__all__ = ("DuckDBConfig",)
+__all__ = (
+    "DuckDBConfig",
+    "DuckDBDriver",
+)