sqlspec 0.8.0__py3-none-any.whl → 0.9.1__py3-none-any.whl

sqlspec/adapters/psycopg/driver.py

@@ -1,29 +1,142 @@
+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.rows import dict_row
 
-from sqlspec.base import PARAM_REGEX, AsyncDriverAdapterProtocol, SyncDriverAdapterProtocol, T
+from sqlspec.base import AsyncDriverAdapterProtocol, SyncDriverAdapterProtocol, T
+from sqlspec.exceptions import SQLParsingError
+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
 
+logger = logging.getLogger("sqlspec")
+
 __all__ = ("PsycopgAsyncDriver", "PsycopgSyncDriver")
 
 
-class PsycopgSyncDriver(SyncDriverAdapterProtocol["Connection"]):
+class PsycopgParameterParser:
+    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)
+
+            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(PsycopgParameterParser, SyncDriverAdapterProtocol["Connection"]):
     """Psycopg Sync Driver Adapter."""
 
     connection: "Connection"
-    param_style: str = "%s"
+    dialect: str = "postgres"
 
     def __init__(self, connection: "Connection") -> 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
     @contextmanager
     def _with_cursor(connection: "Connection") -> "Generator[Any, None, None]":
@@ -33,83 +146,46 @@ class PsycopgSyncDriver(SyncDriverAdapterProtocol["Connection"]):
         finally:
             cursor.close()
 
-    def _process_sql_params(
-        self, sql: str, parameters: "Optional[StatementParameterType]" = None
-    ) -> "tuple[str, Optional[Union[tuple[Any, ...], list[Any], dict[str, Any]]]]":
-        """Process SQL query and parameters for DB-API execution.
-
-        Converts named parameters (:name) to positional parameters (%s)
-        if the input parameters are a dictionary.
-
-        Args:
-            sql: The SQL query string.
-            parameters: The parameters for the query (dict, tuple, list, or None).
-
-        Returns:
-            A tuple containing the processed SQL string and the processed parameters
-            (always a tuple or None if the input was a dictionary, otherwise the original type).
-
-        Raises:
-            ValueError: If a named parameter in the SQL is not found in the dictionary
-                        or if a parameter in the dictionary is not used in the SQL.
-        """
-        if not isinstance(parameters, dict) or not parameters:
-            # If parameters are not a dict, or empty dict, assume positional/no params
-            # Let the underlying driver handle tuples/lists directly
-            return sql, parameters
-
-        processed_sql = ""
-        processed_params_list: list[Any] = []
-        last_end = 0
-        found_params: set[str] = set()
-
-        for match in PARAM_REGEX.finditer(sql):
-            if match.group("dquote") is not None or match.group("squote") is not None:
-                # Skip placeholders within quotes
-                continue
-
-            var_name = match.group("var_name")
-            if var_name is None:  # Should not happen with the regex, but safeguard
-                continue
-
-            if var_name not in parameters:
-                msg = f"Named parameter ':{var_name}' found in SQL but not provided in parameters dictionary."
-                raise ValueError(msg)
-
-            # Append segment before the placeholder + the driver's positional placeholder
-            processed_sql += sql[last_end : match.start("var_name") - 1] + "%s"
-            processed_params_list.append(parameters[var_name])
-            found_params.add(var_name)
-            last_end = match.end("var_name")
-
-        # Append the rest of the SQL string
-        processed_sql += sql[last_end:]
-
-        # Check if all provided parameters were used
-        unused_params = set(parameters.keys()) - found_params
-        if unused_params:
-            msg = f"Parameters provided but not found in SQL: {unused_params}"
-            # Depending on desired strictness, this could be a warning or an error
-            # For now, let's raise an error for clarity
-            raise ValueError(msg)
-
-        return processed_sql, tuple(processed_params_list)
-
+    # --- Public API Methods --- #
+    @overload
     def select(
         self,
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
+        *,
         connection: "Optional[Connection]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "Sequence[dict[str, Any]]": ...
+    @overload
+    def select(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[Connection]" = None,
+        schema_type: "type[ModelDTOT]",
+        **kwargs: Any,
+    ) -> "Sequence[ModelDTOT]": ...
+    def select(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
         schema_type: "Optional[type[ModelDTOT]]" = None,
-    ) -> "list[Union[ModelDTOT, dict[str, Any]]]":
+        connection: "Optional[Connection]" = None,
+        **kwargs: Any,
+    ) -> "Sequence[Union[ModelDTOT, dict[str, Any]]]":
         """Fetch data from the database.
 
         Returns:
             List of row data as either model instances or dictionaries.
         """
         connection = self._connection(connection)
-        sql, parameters = self._process_sql_params(sql, parameters)
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
         with self._with_cursor(connection) as cursor:
             cursor.execute(sql, parameters)
             results = cursor.fetchall()
@@ -120,13 +196,37 @@ 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,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "dict[str, Any]": ...
+    @overload
+    def select_one(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[Connection]" = None,
+        schema_type: "type[ModelDTOT]",
+        **kwargs: Any,
+    ) -> "ModelDTOT": ...
     def select_one(
         self,
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
+        *,
         connection: "Optional[Connection]" = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
+        **kwargs: Any,
     ) -> "Union[ModelDTOT, dict[str, Any]]":
         """Fetch one row from the database.
 
@@ -134,8 +234,7 @@ class PsycopgSyncDriver(SyncDriverAdapterProtocol["Connection"]):
             The first row of the query results.
         """
         connection = self._connection(connection)
-        sql, parameters = self._process_sql_params(sql, parameters)
-
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
         with self._with_cursor(connection) as cursor:
             cursor.execute(sql, parameters)
             row = cursor.fetchone()
@@ -144,13 +243,37 @@ 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[Connection]" = 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[Connection]" = 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,
         schema_type: "Optional[type[ModelDTOT]]" = None,
+        **kwargs: Any,
     ) -> "Optional[Union[ModelDTOT, dict[str, Any]]]":
         """Fetch one row from the database.
 
@@ -158,8 +281,7 @@ class PsycopgSyncDriver(SyncDriverAdapterProtocol["Connection"]):
             The first row of the query results.
         """
         connection = self._connection(connection)
-        sql, parameters = self._process_sql_params(sql, parameters)
-
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
         with self._with_cursor(connection) as cursor:
             cursor.execute(sql, parameters)
             row = cursor.fetchone()
@@ -169,13 +291,37 @@ 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,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "Any": ...
+    @overload
     def select_value(
         self,
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
+        *,
+        connection: "Optional[Connection]" = None,
+        schema_type: "type[T]",
+        **kwargs: Any,
+    ) -> "T": ...
+    def select_value(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
         connection: "Optional[Connection]" = None,
         schema_type: "Optional[type[T]]" = None,
+        **kwargs: Any,
     ) -> "Union[T, Any]":
         """Fetch a single value from the database.
 
@@ -183,24 +329,48 @@ class PsycopgSyncDriver(SyncDriverAdapterProtocol["Connection"]):
             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)
-
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
         with self._with_cursor(connection) as cursor:
             cursor.execute(sql, parameters)
             row = cursor.fetchone()
             row = self.check_not_found(row)
-            val = next(iter(row))
+            val = next(iter(row.values())) if row else None
+            val = self.check_not_found(val)
             if schema_type is not None:
                 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,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "Optional[Any]": ...
+    @overload
+    def select_value_or_none(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[Connection]" = None,
+        schema_type: "type[T]",
+        **kwargs: Any,
+    ) -> "Optional[T]": ...
+    def select_value_or_none(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
         connection: "Optional[Connection]" = None,
         schema_type: "Optional[type[T]]" = None,
+        **kwargs: Any,
     ) -> "Optional[Union[T, Any]]":
         """Fetch a single value from the database.
 
@@ -208,14 +378,15 @@ class PsycopgSyncDriver(SyncDriverAdapterProtocol["Connection"]):
             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)
-
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
         with self._with_cursor(connection) as cursor:
             cursor.execute(sql, parameters)
             row = cursor.fetchone()
             if row is None:
                 return None
-            val = next(iter(row))
+            val = next(iter(row.values())) if row else None
+            if val is None:
+                return None
             if schema_type is not None:
                 return schema_type(val)  # type: ignore[call-arg]
             return val
@@ -225,27 +396,52 @@ class PsycopgSyncDriver(SyncDriverAdapterProtocol["Connection"]):
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
+        *,
         connection: "Optional[Connection]" = None,
+        **kwargs: Any,
     ) -> int:
-        """Insert, update, or delete data from the database.
+        """Execute an INSERT, UPDATE, or DELETE query and return the number of affected rows.
 
         Returns:
-            Row count affected by the operation.
+            The number of rows affected by the operation.
         """
         connection = self._connection(connection)
-        sql, parameters = self._process_sql_params(sql, parameters)
-
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
         with self._with_cursor(connection) as cursor:
             cursor.execute(sql, parameters)
-            return cursor.rowcount if hasattr(cursor, "rowcount") else -1
+            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,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "dict[str, Any]": ...
+    @overload
+    def insert_update_delete_returning(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[Connection]" = None,
+        schema_type: "type[ModelDTOT]",
+        **kwargs: Any,
+    ) -> "ModelDTOT": ...
+    def insert_update_delete_returning(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
         connection: "Optional[Connection]" = 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 result.
 
@@ -253,8 +449,7 @@ class PsycopgSyncDriver(SyncDriverAdapterProtocol["Connection"]):
             The first row of results.
         """
         connection = self._connection(connection)
-        sql, parameters = self._process_sql_params(sql, parameters)
-
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
         with self._with_cursor(connection) as cursor:
             cursor.execute(sql, parameters)
             result = cursor.fetchone()
@@ -271,7 +466,9 @@ class PsycopgSyncDriver(SyncDriverAdapterProtocol["Connection"]):
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
+        *,
         connection: "Optional[Connection]" = None,
+        **kwargs: Any,
     ) -> str:
         """Execute a script.
 
@@ -279,45 +476,17 @@ class PsycopgSyncDriver(SyncDriverAdapterProtocol["Connection"]):
             Status message for the operation.
         """
         connection = self._connection(connection)
-        sql, parameters = self._process_sql_params(sql, parameters)
-
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
         with self._with_cursor(connection) as cursor:
             cursor.execute(sql, parameters)
-            return str(cursor.rowcount)
-
-    def execute_script_returning(
-        self,
-        sql: str,
-        parameters: "Optional[StatementParameterType]" = None,
-        /,
-        connection: "Optional[Connection]" = None,
-        schema_type: "Optional[type[ModelDTOT]]" = None,
-    ) -> "Optional[Union[dict[str, Any], ModelDTOT]]":
-        """Execute a script and return result.
-
-        Returns:
-            The first row of results.
-        """
-        connection = self._connection(connection)
-        sql, parameters = self._process_sql_params(sql, parameters)
-
-        with self._with_cursor(connection) as cursor:
-            cursor.execute(sql, parameters)
-            result = cursor.fetchone()
-
-            if result is None:
-                return None
-
-            if schema_type is not None:
-                return cast("ModelDTOT", schema_type(**result))  # pyright: ignore[reportUnknownArgumentType]
-            return cast("dict[str, Any]", result)  # pyright: ignore[reportUnknownArgumentType,reportUnknownVariableType]
+            return str(cursor.statusmessage) if cursor.statusmessage is not None else "DONE"
 
 
-class PsycopgAsyncDriver(AsyncDriverAdapterProtocol["AsyncConnection"]):
+class PsycopgAsyncDriver(PsycopgParameterParser, AsyncDriverAdapterProtocol["AsyncConnection"]):
     """Psycopg Async Driver Adapter."""
 
     connection: "AsyncConnection"
-    param_style: str = "%s"
+    dialect: str = "postgres"
 
     def __init__(self, connection: "AsyncConnection") -> None:
         self.connection = connection
@@ -331,83 +500,46 @@ class PsycopgAsyncDriver(AsyncDriverAdapterProtocol["AsyncConnection"]):
         finally:
             await cursor.close()
 
-    def _process_sql_params(
-        self, sql: str, parameters: "Optional[StatementParameterType]" = None
-    ) -> "tuple[str, Optional[Union[tuple[Any, ...], list[Any], dict[str, Any]]]]":
-        """Process SQL query and parameters for DB-API execution.
-
-        Converts named parameters (:name) to positional parameters (%s)
-        if the input parameters are a dictionary.
-
-        Args:
-            sql: The SQL query string.
-            parameters: The parameters for the query (dict, tuple, list, or None).
-
-        Returns:
-            A tuple containing the processed SQL string and the processed parameters
-            (always a tuple or None if the input was a dictionary, otherwise the original type).
-
-        Raises:
-            ValueError: If a named parameter in the SQL is not found in the dictionary
-                        or if a parameter in the dictionary is not used in the SQL.
-        """
-        if not isinstance(parameters, dict) or not parameters:
-            # If parameters are not a dict, or empty dict, assume positional/no params
-            # Let the underlying driver handle tuples/lists directly
-            return sql, parameters
-
-        processed_sql = ""
-        processed_params_list: list[Any] = []
-        last_end = 0
-        found_params: set[str] = set()
-
-        for match in PARAM_REGEX.finditer(sql):
-            if match.group("dquote") is not None or match.group("squote") is not None:
-                # Skip placeholders within quotes
-                continue
-
-            var_name = match.group("var_name")
-            if var_name is None:  # Should not happen with the regex, but safeguard
-                continue
-
-            if var_name not in parameters:
-                msg = f"Named parameter ':{var_name}' found in SQL but not provided in parameters dictionary."
-                raise ValueError(msg)
-
-            # Append segment before the placeholder + the driver's positional placeholder
-            processed_sql += sql[last_end : match.start("var_name") - 1] + "%s"
-            processed_params_list.append(parameters[var_name])
-            found_params.add(var_name)
-            last_end = match.end("var_name")
-
-        # Append the rest of the SQL string
-        processed_sql += sql[last_end:]
-
-        # Check if all provided parameters were used
-        unused_params = set(parameters.keys()) - found_params
-        if unused_params:
-            msg = f"Parameters provided but not found in SQL: {unused_params}"
-            # Depending on desired strictness, this could be a warning or an error
-            # For now, let's raise an error for clarity
-            raise ValueError(msg)
-
-        return processed_sql, tuple(processed_params_list)
-
+    # --- Public API Methods --- #
+    @overload
     async def select(
         self,
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
+        *,
+        connection: "Optional[AsyncConnection]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "Sequence[dict[str, Any]]": ...
+    @overload
+    async def select(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AsyncConnection]" = None,
+        schema_type: "type[ModelDTOT]",
+        **kwargs: Any,
+    ) -> "Sequence[ModelDTOT]": ...
+    async def select(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
         connection: "Optional[AsyncConnection]" = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
-    ) -> "list[Union[ModelDTOT, dict[str, Any]]]":
+        **kwargs: Any,
+    ) -> "Sequence[Union[ModelDTOT, dict[str, Any]]]":
         """Fetch data from the database.
 
         Returns:
             List of row data as either model instances or dictionaries.
         """
         connection = self._connection(connection)
-        sql, parameters = self._process_sql_params(sql, parameters)
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
         results: list[Union[ModelDTOT, dict[str, Any]]] = []
 
         async with self._with_cursor(connection) as cursor:
@@ -419,13 +551,37 @@ 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,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "dict[str, Any]": ...
+    @overload
+    async def select_one(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AsyncConnection]" = None,
+        schema_type: "type[ModelDTOT]",
+        **kwargs: Any,
+    ) -> "ModelDTOT": ...
+    async def select_one(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
         connection: "Optional[AsyncConnection]" = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
+        **kwargs: Any,
     ) -> "Union[ModelDTOT, dict[str, Any]]":
         """Fetch one row from the database.
 
@@ -433,7 +589,7 @@ class PsycopgAsyncDriver(AsyncDriverAdapterProtocol["AsyncConnection"]):
             The first row of the query results.
         """
         connection = self._connection(connection)
-        sql, parameters = self._process_sql_params(sql, parameters)
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
 
         async with self._with_cursor(connection) as cursor:
             await cursor.execute(sql, parameters)
@@ -443,13 +599,37 @@ 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[AsyncConnection]" = 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[AsyncConnection]" = None,
+        schema_type: "type[ModelDTOT]",
+        **kwargs: Any,
+    ) -> "Optional[ModelDTOT]": ...
+    async def select_one_or_none(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
         schema_type: "Optional[type[ModelDTOT]]" = None,
+        connection: "Optional[AsyncConnection]" = None,
+        **kwargs: Any,
     ) -> "Optional[Union[ModelDTOT, dict[str, Any]]]":
         """Fetch one row from the database.
 
@@ -457,7 +637,7 @@ class PsycopgAsyncDriver(AsyncDriverAdapterProtocol["AsyncConnection"]):
             The first row of the query results.
         """
         connection = self._connection(connection)
-        sql, parameters = self._process_sql_params(sql, parameters)
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
 
         async with self._with_cursor(connection) as cursor:
             await cursor.execute(sql, parameters)
@@ -468,27 +648,52 @@ 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[AsyncConnection]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "Any": ...
+    @overload
+    async def select_value(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AsyncConnection]" = None,
+        schema_type: "type[T]",
+        **kwargs: Any,
+    ) -> "T": ...
+    async def select_value(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
         connection: "Optional[AsyncConnection]" = None,
         schema_type: "Optional[type[T]]" = None,
-    ) -> "Optional[Union[T, Any]]":
+        **kwargs: Any,
+    ) -> "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)
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
 
         async with self._with_cursor(connection) as cursor:
             await cursor.execute(sql, parameters)
             row = await cursor.fetchone()
             row = self.check_not_found(row)
-            val = next(iter(row))
+            val = next(iter(row.values())) if row else None
+            val = self.check_not_found(val)
             if schema_type is not None:
                 return schema_type(val)  # type: ignore[call-arg]
             return val
@@ -498,8 +703,10 @@ class PsycopgAsyncDriver(AsyncDriverAdapterProtocol["AsyncConnection"]):
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
+        *,
         connection: "Optional[AsyncConnection]" = None,
         schema_type: "Optional[type[T]]" = None,
+        **kwargs: Any,
     ) -> "Optional[Union[T, Any]]":
         """Fetch a single value from the database.
 
@@ -507,14 +714,16 @@ class PsycopgAsyncDriver(AsyncDriverAdapterProtocol["AsyncConnection"]):
             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)
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
 
         async with self._with_cursor(connection) as cursor:
             await cursor.execute(sql, parameters)
             row = await cursor.fetchone()
             if row is None:
                 return None
-            val = next(iter(row))
+            val = next(iter(row.values())) if row else None
+            if val is None:
+                return None
             if schema_type is not None:
                 return schema_type(val)  # type: ignore[call-arg]
             return val
@@ -524,15 +733,17 @@ class PsycopgAsyncDriver(AsyncDriverAdapterProtocol["AsyncConnection"]):
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
+        *,
         connection: "Optional[AsyncConnection]" = None,
+        **kwargs: Any,
     ) -> int:
-        """Insert, update, or delete data from the database.
+        """Execute an INSERT, UPDATE, or DELETE query and return the number of affected rows.
 
         Returns:
-            Row count affected by the operation.
+            The number of rows affected by the operation.
         """
         connection = self._connection(connection)
-        sql, parameters = self._process_sql_params(sql, parameters)
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
 
         async with self._with_cursor(connection) as cursor:
             await cursor.execute(sql, parameters)
@@ -542,13 +753,37 @@ class PsycopgAsyncDriver(AsyncDriverAdapterProtocol["AsyncConnection"]):
                 rowcount = -1
             return rowcount
 
+    @overload
+    async def insert_update_delete_returning(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AsyncConnection]" = 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[AsyncConnection]" = 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,
         schema_type: "Optional[type[ModelDTOT]]" = None,
+        **kwargs: Any,
     ) -> "Optional[Union[dict[str, Any], ModelDTOT]]":
         """Insert, update, or delete data from the database and return result.
 
@@ -556,7 +791,7 @@ class PsycopgAsyncDriver(AsyncDriverAdapterProtocol["AsyncConnection"]):
             The first row of results.
         """
         connection = self._connection(connection)
-        sql, parameters = self._process_sql_params(sql, parameters)
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
 
         async with self._with_cursor(connection) as cursor:
             await cursor.execute(sql, parameters)
@@ -574,7 +809,9 @@ class PsycopgAsyncDriver(AsyncDriverAdapterProtocol["AsyncConnection"]):
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
+        *,
         connection: "Optional[AsyncConnection]" = None,
+        **kwargs: Any,
     ) -> str:
         """Execute a script.
 
@@ -582,35 +819,8 @@ class PsycopgAsyncDriver(AsyncDriverAdapterProtocol["AsyncConnection"]):
             Status message for the operation.
         """
         connection = self._connection(connection)
-        sql, parameters = self._process_sql_params(sql, parameters)
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
 
         async with self._with_cursor(connection) as cursor:
             await cursor.execute(sql, parameters)
-            return str(cursor.rowcount)
-
-    async def execute_script_returning(
-        self,
-        sql: str,
-        parameters: "Optional[StatementParameterType]" = None,
-        /,
-        connection: "Optional[AsyncConnection]" = None,
-        schema_type: "Optional[type[ModelDTOT]]" = None,
-    ) -> "Optional[Union[dict[str, Any], ModelDTOT]]":
-        """Execute a script and return result.
-
-        Returns:
-            The first row of results.
-        """
-        connection = self._connection(connection)
-        sql, parameters = self._process_sql_params(sql, parameters)
-
-        async with self._with_cursor(connection) as cursor:
-            await cursor.execute(sql, parameters)
-            result = await cursor.fetchone()
-
-            if result is None:
-                return None
-
-            if schema_type is not None:
-                return cast("ModelDTOT", schema_type(**result))  # pyright: ignore[reportUnknownArgumentType]
-            return cast("dict[str, Any]", result)  # pyright: ignore[reportUnknownArgumentType,reportUnknownVariableType]
+            return str(cursor.statusmessage) if cursor.statusmessage is not None else "DONE"