sqlspec 0.10.1__py3-none-any.whl → 0.11.0__py3-none-any.whl

sqlspec/adapters/aiosqlite/driver.py

@@ -1,23 +1,31 @@
+import logging
 from contextlib import asynccontextmanager
-from typing import TYPE_CHECKING, Any, Optional, Union, cast, overload
+from typing import TYPE_CHECKING, Any, Optional, Union, overload
 
 import aiosqlite
+from sqlglot import exp
 
 from sqlspec.base import AsyncDriverAdapterProtocol
-from sqlspec.mixins import SQLTranslatorMixin
+from sqlspec.mixins import ResultConverter, SQLTranslatorMixin
+from sqlspec.statement import SQLStatement
+from sqlspec.typing import is_dict
 
 if TYPE_CHECKING:
     from collections.abc import AsyncGenerator, Sequence
 
+    from sqlspec.filters import StatementFilter
     from sqlspec.typing import ModelDTOT, StatementParameterType, T
 
 __all__ = ("AiosqliteConnection", "AiosqliteDriver")
 AiosqliteConnection = aiosqlite.Connection
 
+logger = logging.getLogger("sqlspec")
+
 
 class AiosqliteDriver(
     SQLTranslatorMixin["AiosqliteConnection"],
     AsyncDriverAdapterProtocol["AiosqliteConnection"],
+    ResultConverter,
 ):
     """SQLite Async Driver Adapter."""
 
@@ -39,6 +47,74 @@ class AiosqliteDriver(
         finally:
             await cursor.close()
 
+    def _process_sql_params(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *filters: "StatementFilter",
+        **kwargs: Any,
+    ) -> "tuple[str, Optional[Union[tuple[Any, ...], list[Any], dict[str, Any]]]]":
+        """Process SQL and parameters for aiosqlite using SQLStatement.
+
+        aiosqlite supports both named (:name) and positional (?) parameters.
+        This method processes the SQL with dialect-aware parsing and handles
+        parameters appropriately for aiosqlite.
+
+        Args:
+            sql: SQL statement.
+            parameters: Query parameters.
+            *filters: Statement filters to apply.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            Tuple of processed SQL and parameters.
+        """
+        statement = SQLStatement(sql, parameters, kwargs=kwargs, dialect=self.dialect)
+
+        # Apply any filters
+        for filter_obj in filters:
+            statement = statement.apply_filter(filter_obj)
+
+        processed_sql, processed_params, parsed_expr = statement.process()
+        if processed_params is None:
+            return processed_sql, None
+
+        if is_dict(processed_params):
+            # For dict parameters, we need to use ordered ? placeholders
+            # but only if we have a parsed expression to work with
+            if parsed_expr:
+                # Collect named parameters in the order they appear in the SQL
+                named_params = []
+                for node in parsed_expr.find_all(exp.Parameter, exp.Placeholder):
+                    if isinstance(node, exp.Parameter) and node.name and node.name in processed_params:
+                        named_params.append(node.name)
+                    elif (
+                        isinstance(node, exp.Placeholder)
+                        and isinstance(node.this, str)
+                        and node.this in processed_params
+                    ):
+                        named_params.append(node.this)
+
+                if named_params:
+                    # Transform SQL to use ? placeholders
+                    def _convert_to_qmark(node: exp.Expression) -> exp.Expression:
+                        if (isinstance(node, exp.Parameter) and node.name and node.name in processed_params) or (
+                            isinstance(node, exp.Placeholder)
+                            and isinstance(node.this, str)
+                            and node.this in processed_params
+                        ):
+                            return exp.Placeholder()  # ? placeholder
+                        return node
+
+                    return parsed_expr.transform(_convert_to_qmark, copy=True).sql(dialect=self.dialect), tuple(
+                        processed_params[name] for name in named_params
+                    )
+            return processed_sql, processed_params
+        if isinstance(processed_params, (list, tuple)):
+            return processed_sql, tuple(processed_params)
+        return processed_sql, (processed_params,)
+
     # --- Public API Methods --- #
     @overload
     async def select(
@@ -46,7 +122,7 @@ class AiosqliteDriver(
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
-        *,
+        *filters: "StatementFilter",
         connection: "Optional[AiosqliteConnection]" = None,
         schema_type: None = None,
         **kwargs: Any,
@@ -57,7 +133,7 @@ class AiosqliteDriver(
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
-        *,
+        *filters: "StatementFilter",
         connection: "Optional[AiosqliteConnection]" = None,
         schema_type: "type[ModelDTOT]",
         **kwargs: Any,
@@ -65,29 +141,33 @@ class AiosqliteDriver(
     async def select(
         self,
         sql: str,
-        parameters: Optional["StatementParameterType"] = None,
+        parameters: "Optional[StatementParameterType]" = None,
         /,
-        *,
-        connection: Optional["AiosqliteConnection"] = None,
+        *filters: "StatementFilter",
+        connection: "Optional[AiosqliteConnection]" = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
         **kwargs: Any,
-    ) -> "Sequence[Union[ModelDTOT, dict[str, Any]]]":
+    ) -> "Sequence[Union[dict[str, Any], ModelDTOT]]":
         """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, **kwargs)
-        async with self._with_cursor(connection) as cursor:
-            await cursor.execute(sql, parameters)  # pyright: ignore[reportUnknownMemberType]
-            results = await cursor.fetchall()  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
-            if not results:
-                return []
-            column_names = [c[0] for c in cursor.description or []]  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
-            if schema_type is None:
-                return [dict(zip(column_names, row)) for row in results]  # pyright: ignore[reportUnknownArgumentType]
-            return [cast("ModelDTOT", schema_type(**dict(zip(column_names, row)))) for row in results]  # pyright: ignore[reportUnknownArgumentType]
+        sql, parameters = self._process_sql_params(sql, parameters, *filters, **kwargs)
+
+        # Execute the query
+        cursor = await connection.execute(sql, parameters or ())
+        results = await cursor.fetchall()
+        if not results:
+            return []
+
+        # Get column names
+        column_names = [column[0] for column in cursor.description]
+
+        # Convert to dicts first
+        dict_results = [dict(zip(column_names, row)) for row in results]
+        return self.to_schema(dict_results, schema_type=schema_type)
 
     @overload
     async def select_one(
@@ -95,7 +175,7 @@ class AiosqliteDriver(
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
-        *,
+        *filters: "StatementFilter",
         connection: "Optional[AiosqliteConnection]" = None,
         schema_type: None = None,
         **kwargs: Any,
@@ -106,7 +186,7 @@ class AiosqliteDriver(
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
-        *,
+        *filters: "StatementFilter",
         connection: "Optional[AiosqliteConnection]" = None,
         schema_type: "type[ModelDTOT]",
         **kwargs: Any,
@@ -114,28 +194,32 @@ class AiosqliteDriver(
     async def select_one(
         self,
         sql: str,
-        parameters: Optional["StatementParameterType"] = None,
+        parameters: "Optional[StatementParameterType]" = None,
         /,
-        *,
-        connection: Optional["AiosqliteConnection"] = None,
+        *filters: "StatementFilter",
+        connection: "Optional[AiosqliteConnection]" = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
         **kwargs: Any,
-    ) -> "Union[ModelDTOT, dict[str, Any]]":
+    ) -> "Union[dict[str, Any], ModelDTOT]":
         """Fetch one row from the database.
 
         Returns:
             The first row of the query results.
         """
         connection = self._connection(connection)
-        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
-        async with self._with_cursor(connection) as cursor:
-            await cursor.execute(sql, parameters)  # pyright: ignore[reportUnknownMemberType]
-            result = await cursor.fetchone()  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
-            result = self.check_not_found(result)
-            column_names = [c[0] for c in cursor.description or []]  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
-            if schema_type is None:
-                return dict(zip(column_names, result))  # pyright: ignore[reportUnknownArgumentType, reportUnknownVariableType]
-            return cast("ModelDTOT", schema_type(**dict(zip(column_names, result))))  # pyright: ignore[reportUnknownArgumentType]
+        sql, parameters = self._process_sql_params(sql, parameters, *filters, **kwargs)
+
+        # Execute the query
+        cursor = await connection.execute(sql, parameters or ())
+        result = await cursor.fetchone()
+        result = self.check_not_found(result)
+
+        # Get column names
+        column_names = [column[0] for column in cursor.description]
+
+        # Convert to dict and then use ResultConverter
+        dict_result = dict(zip(column_names, result))
+        return self.to_schema(dict_result, schema_type=schema_type)
 
     @overload
     async def select_one_or_none(
@@ -143,7 +227,7 @@ class AiosqliteDriver(
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
-        *,
+        *filters: "StatementFilter",
         connection: "Optional[AiosqliteConnection]" = None,
         schema_type: None = None,
         **kwargs: Any,
@@ -154,7 +238,7 @@ class AiosqliteDriver(
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
-        *,
+        *filters: "StatementFilter",
         connection: "Optional[AiosqliteConnection]" = None,
         schema_type: "type[ModelDTOT]",
         **kwargs: Any,
@@ -162,29 +246,33 @@ class AiosqliteDriver(
     async def select_one_or_none(
         self,
         sql: str,
-        parameters: Optional["StatementParameterType"] = None,
+        parameters: "Optional[StatementParameterType]" = None,
         /,
-        *,
-        connection: Optional["AiosqliteConnection"] = None,
+        *filters: "StatementFilter",
+        connection: "Optional[AiosqliteConnection]" = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
         **kwargs: Any,
-    ) -> "Optional[Union[ModelDTOT, dict[str, Any]]]":
+    ) -> "Optional[Union[dict[str, Any], ModelDTOT]]":
         """Fetch one row from the database.
 
         Returns:
             The first row of the query results.
         """
         connection = self._connection(connection)
-        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
-        async with self._with_cursor(connection) as cursor:
-            await cursor.execute(sql, parameters)  # pyright: ignore[reportUnknownMemberType]
-            result = await cursor.fetchone()  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
-            if result is None:
-                return None
-            column_names = [c[0] for c in cursor.description or []]  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
-            if schema_type is None:
-                return dict(zip(column_names, result))  # pyright: ignore[reportUnknownArgumentType, reportUnknownVariableType]
-            return cast("ModelDTOT", schema_type(**dict(zip(column_names, result))))  # pyright: ignore[reportUnknownArgumentType]
+        sql, parameters = self._process_sql_params(sql, parameters, *filters, **kwargs)
+
+        # Execute the query
+        cursor = await connection.execute(sql, parameters or ())
+        result = await cursor.fetchone()
+        if result is None:
+            return None
+
+        # Get column names
+        column_names = [column[0] for column in cursor.description]
+
+        # Convert to dict and then use ResultConverter
+        dict_result = dict(zip(column_names, result))
+        return self.to_schema(dict_result, schema_type=schema_type)
 
     @overload
     async def select_value(
@@ -192,7 +280,7 @@ class AiosqliteDriver(
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
-        *,
+        *filters: "StatementFilter",
         connection: "Optional[AiosqliteConnection]" = None,
         schema_type: None = None,
         **kwargs: Any,
@@ -203,7 +291,7 @@ class AiosqliteDriver(
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
-        *,
+        *filters: "StatementFilter",
         connection: "Optional[AiosqliteConnection]" = None,
         schema_type: "type[T]",
         **kwargs: Any,
@@ -213,7 +301,7 @@ class AiosqliteDriver(
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
-        *,
+        *filters: "StatementFilter",
         connection: "Optional[AiosqliteConnection]" = None,
         schema_type: "Optional[type[T]]" = None,
         **kwargs: Any,
@@ -221,17 +309,21 @@ class AiosqliteDriver(
         """Fetch a single value from the database.
 
         Returns:
-            The first value from the first row of results, or None if no results.
+            The first value from the first row of results.
         """
         connection = self._connection(connection)
-        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
-        async with self._with_cursor(connection) as cursor:
-            await cursor.execute(sql, parameters)  # pyright: ignore[reportUnknownMemberType]
-            result = await cursor.fetchone()  # pyright: ignore[reportUnknownMemberType]
-            result = self.check_not_found(result)
-            if schema_type is None:
-                return result[0]
-            return schema_type(result[0])  # type: ignore[call-arg]
+        sql, parameters = self._process_sql_params(sql, parameters, *filters, **kwargs)
+
+        # Execute the query
+        cursor = await connection.execute(sql, parameters or ())
+        result = await cursor.fetchone()
+        result = self.check_not_found(result)
+
+        # Return first value from the row
+        result_value = result[0]
+        if schema_type is None:
+            return result_value
+        return schema_type(result_value)  # type: ignore[call-arg]
 
     @overload
     async def select_value_or_none(
@@ -239,7 +331,7 @@ class AiosqliteDriver(
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
-        *,
+        *filters: "StatementFilter",
         connection: "Optional[AiosqliteConnection]" = None,
         schema_type: None = None,
         **kwargs: Any,
@@ -250,7 +342,7 @@ class AiosqliteDriver(
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
-        *,
+        *filters: "StatementFilter",
         connection: "Optional[AiosqliteConnection]" = None,
         schema_type: "type[T]",
         **kwargs: Any,
@@ -260,7 +352,7 @@ class AiosqliteDriver(
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
-        *,
+        *filters: "StatementFilter",
         connection: "Optional[AiosqliteConnection]" = None,
         schema_type: "Optional[type[T]]" = None,
         **kwargs: Any,
@@ -271,23 +363,27 @@ class AiosqliteDriver(
             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)
-        async with self._with_cursor(connection) as cursor:
-            await cursor.execute(sql, parameters)  # pyright: ignore[reportUnknownMemberType]
-            result = await cursor.fetchone()  # pyright: ignore[reportUnknownMemberType]
-            if result is None:
-                return None
-            if schema_type is None:
-                return result[0]
-            return schema_type(result[0])  # type: ignore[call-arg]
+        sql, parameters = self._process_sql_params(sql, parameters, *filters, **kwargs)
+
+        # Execute the query
+        cursor = await connection.execute(sql, parameters or ())
+        result = await cursor.fetchone()
+        if result is None:
+            return None
+
+        # Return first value from the row
+        result_value = result[0]
+        if schema_type is None:
+            return result_value
+        return schema_type(result_value)  # type: ignore[call-arg]
 
     async def insert_update_delete(
         self,
         sql: str,
-        parameters: Optional["StatementParameterType"] = None,
+        parameters: "Optional[StatementParameterType]" = None,
         /,
-        *,
-        connection: Optional["AiosqliteConnection"] = None,
+        *filters: "StatementFilter",
+        connection: "Optional[AiosqliteConnection]" = None,
         **kwargs: Any,
     ) -> int:
         """Insert, update, or delete data from the database.
@@ -296,11 +392,12 @@ class AiosqliteDriver(
             Row count affected by the operation.
         """
         connection = self._connection(connection)
-        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
+        sql, parameters = self._process_sql_params(sql, parameters, *filters, **kwargs)
 
-        async with self._with_cursor(connection) as cursor:
-            await cursor.execute(sql, parameters)  # pyright: ignore[reportUnknownMemberType]
-            return cursor.rowcount if hasattr(cursor, "rowcount") else -1  # pyright: ignore[reportUnknownVariableType, reportGeneralTypeIssues]
+        # Execute the query
+        cursor = await connection.execute(sql, parameters or ())
+        await connection.commit()
+        return cursor.rowcount
 
     @overload
     async def insert_update_delete_returning(
@@ -308,7 +405,7 @@ class AiosqliteDriver(
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
-        *,
+        *filters: "StatementFilter",
         connection: "Optional[AiosqliteConnection]" = None,
         schema_type: None = None,
         **kwargs: Any,
@@ -319,7 +416,7 @@ class AiosqliteDriver(
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
-        *,
+        *filters: "StatementFilter",
         connection: "Optional[AiosqliteConnection]" = None,
         schema_type: "type[ModelDTOT]",
         **kwargs: Any,
@@ -327,38 +424,42 @@ class AiosqliteDriver(
     async def insert_update_delete_returning(
         self,
         sql: str,
-        parameters: Optional["StatementParameterType"] = None,
+        parameters: "Optional[StatementParameterType]" = None,
         /,
-        *,
-        connection: Optional["AiosqliteConnection"] = None,
+        *filters: "StatementFilter",
+        connection: "Optional[AiosqliteConnection]" = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
         **kwargs: Any,
-    ) -> "Optional[Union[dict[str, Any], ModelDTOT]]":
+    ) -> "Union[dict[str, Any], ModelDTOT]":
         """Insert, update, or delete data from the database and return result.
 
         Returns:
             The first row of results.
         """
         connection = self._connection(connection)
-        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
+        sql, parameters = self._process_sql_params(sql, parameters, *filters, **kwargs)
+
+        # Execute the query
+        cursor = await connection.execute(sql, parameters or ())
+        result = await cursor.fetchone()
+        await connection.commit()
+        await cursor.close()
 
-        async with self._with_cursor(connection) as cursor:
-            await cursor.execute(sql, parameters)  # pyright: ignore[reportUnknownMemberType]
-            results = list(await cursor.fetchall())  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
-            if not results:  # Check if empty
-                return None
-            column_names = [c[0] for c in cursor.description or []]  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
-            if schema_type is not None:
-                return cast("ModelDTOT", schema_type(**dict(zip(column_names, results[0]))))  # pyright: ignore[reportUnknownArgumentType, reportUnknownVariableType]
-            return dict(zip(column_names, results[0]))  # pyright: ignore[reportUnknownArgumentType, reportUnknownVariableType]
+        result = self.check_not_found(result)
+
+        # Get column names
+        column_names = [column[0] for column in cursor.description]
+
+        # Convert to dict and then use ResultConverter
+        dict_result = dict(zip(column_names, result))
+        return self.to_schema(dict_result, schema_type=schema_type)
 
     async def execute_script(
         self,
         sql: str,
-        parameters: Optional["StatementParameterType"] = None,
+        parameters: "Optional[StatementParameterType]" = None,
         /,
-        *,
-        connection: Optional["AiosqliteConnection"] = None,
+        connection: "Optional[AiosqliteConnection]" = None,
         **kwargs: Any,
     ) -> str:
         """Execute a script.
@@ -369,34 +470,18 @@ class AiosqliteDriver(
         connection = self._connection(connection)
         sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
 
-        async with self._with_cursor(connection) as cursor:
-            await cursor.execute(sql, parameters)  # pyright: ignore[reportUnknownMemberType]
-            return "DONE"
+        # Execute the script
+        await connection.executescript(sql)
+        await connection.commit()
+        return "Script executed successfully."
 
-    async def execute_script_returning(
-        self,
-        sql: str,
-        parameters: Optional["StatementParameterType"] = None,
-        /,
-        *,
-        connection: Optional["AiosqliteConnection"] = None,
-        schema_type: "Optional[type[ModelDTOT]]" = None,
-        **kwargs: Any,
-    ) -> "Optional[Union[dict[str, Any], ModelDTOT]]":
-        """Execute a script and return result.
+    def _connection(self, connection: "Optional[AiosqliteConnection]" = None) -> "AiosqliteConnection":
+        """Get the connection to use for the operation.
+
+        Args:
+            connection: Optional connection to use.
 
         Returns:
-            The first row of results.
+            The connection to use.
         """
-        connection = self._connection(connection)
-        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
-
-        async with self._with_cursor(connection) as cursor:
-            await cursor.execute(sql, parameters)  # pyright: ignore[reportUnknownMemberType]
-            results = list(await cursor.fetchall())  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
-            if not results:  # Check if empty
-                return None
-            column_names = [c[0] for c in cursor.description or []]  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
-            if schema_type is not None:
-                return cast("ModelDTOT", schema_type(**dict(zip(column_names, results[0]))))  # pyright: ignore[reportUnknownArgumentType, reportUnknownVariableType]
-            return dict(zip(column_names, results[0]))  # pyright: ignore[reportUnknownArgumentType, reportUnknownVariableType]
+        return connection or self.connection