sqlspec 0.11.1__py3-none-any.whl → 0.12.1__py3-none-any.whl

sqlspec/adapters/asyncmy/driver.py

@@ -1,462 +1,244 @@
-# type: ignore
 import logging
-import re
-from collections.abc import AsyncGenerator
+from collections.abc import AsyncGenerator, Sequence
 from contextlib import asynccontextmanager
-from typing import TYPE_CHECKING, Any, Optional, Union, overload
+from typing import TYPE_CHECKING, Any, ClassVar, Optional, Union, cast
 
 from asyncmy import Connection
-
-from sqlspec.base import AsyncDriverAdapterProtocol
-from sqlspec.exceptions import ParameterStyleMismatchError
-from sqlspec.filters import StatementFilter
-from sqlspec.mixins import ResultConverter, SQLTranslatorMixin
-from sqlspec.statement import SQLStatement
-from sqlspec.typing import is_dict
+from typing_extensions import TypeAlias
+
+from sqlspec.driver import AsyncDriverAdapterProtocol
+from sqlspec.driver.mixins import (
+    AsyncPipelinedExecutionMixin,
+    AsyncStorageMixin,
+    SQLTranslatorMixin,
+    ToSchemaMixin,
+    TypeCoercionMixin,
+)
+from sqlspec.statement.parameters import ParameterStyle
+from sqlspec.statement.result import DMLResultDict, ScriptResultDict, SelectResultDict, SQLResult
+from sqlspec.statement.sql import SQL, SQLConfig
+from sqlspec.typing import DictRow, ModelDTOT, RowT
 
 if TYPE_CHECKING:
-    from collections.abc import Mapping, Sequence
-
-    from asyncmy.cursors import Cursor
-
-    from sqlspec.typing import ModelDTOT, StatementParameterType, T
+    from asyncmy.cursors import Cursor, DictCursor
+    from sqlglot.dialects.dialect import DialectType
 
-__all__ = ("AsyncmyDriver",)
-
-AsyncmyConnection = Connection
+__all__ = ("AsyncmyConnection", "AsyncmyDriver")
 
 logger = logging.getLogger("sqlspec")
 
-# Pattern to identify MySQL-style placeholders (%s) for proper conversion
-MYSQL_PLACEHOLDER_PATTERN = re.compile(r"(?<!%)%s")
+AsyncmyConnection: TypeAlias = Connection
 
 
 class AsyncmyDriver(
-    SQLTranslatorMixin["AsyncmyConnection"],
-    AsyncDriverAdapterProtocol["AsyncmyConnection"],
-    ResultConverter,
+    AsyncDriverAdapterProtocol[AsyncmyConnection, RowT],
+    SQLTranslatorMixin,
+    TypeCoercionMixin,
+    AsyncStorageMixin,
+    AsyncPipelinedExecutionMixin,
+    ToSchemaMixin,
 ):
-    """Asyncmy MySQL/MariaDB Driver Adapter."""
+    """Asyncmy MySQL/MariaDB Driver Adapter. Modern protocol implementation."""
 
-    connection: "AsyncmyConnection"
-    dialect: str = "mysql"
+    dialect: "DialectType" = "mysql"
+    supported_parameter_styles: "tuple[ParameterStyle, ...]" = (ParameterStyle.POSITIONAL_PYFORMAT,)
+    default_parameter_style: ParameterStyle = ParameterStyle.POSITIONAL_PYFORMAT
+    __supports_arrow__: ClassVar[bool] = True
+    __supports_parquet__: ClassVar[bool] = False
+    __slots__ = ()
 
-    def __init__(self, connection: "AsyncmyConnection") -> None:
-        self.connection = connection
+    def __init__(
+        self,
+        connection: AsyncmyConnection,
+        config: Optional[SQLConfig] = None,
+        default_row_type: type[DictRow] = DictRow,
+    ) -> None:
+        super().__init__(connection=connection, config=config, default_row_type=default_row_type)
 
-    @staticmethod
     @asynccontextmanager
-    async def _with_cursor(connection: "AsyncmyConnection") -> AsyncGenerator["Cursor", None]:
-        cursor = connection.cursor()
+    async def _get_cursor(
+        self, connection: "Optional[AsyncmyConnection]" = None
+    ) -> "AsyncGenerator[Union[Cursor, DictCursor], None]":
+        conn = self._connection(connection)
+        cursor = await conn.cursor()
         try:
             yield cursor
         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 using SQLStatement with dialect support.
-
-        Args:
-            sql: The SQL statement to process.
-            parameters: The parameters to bind to the statement. Can be data or a StatementFilter.
-            *filters: Statement filters to apply.
-            **kwargs: Additional keyword arguments.
-
-        Raises:
-            ParameterStyleMismatchError: If the parameter style is not supported.
-
-        Returns:
-            A tuple of (sql, parameters) ready for execution.
-        """
-        # Convert filters tuple to a list to allow modification
-        current_filters: list[StatementFilter] = list(filters)
-        actual_parameters: Optional[Union[Mapping[str, Any], Sequence[Any]]] = None
-
-        if parameters is not None:
-            if isinstance(parameters, StatementFilter):
-                current_filters.insert(0, parameters)
-                # actual_parameters remains None
-            else:
-                actual_parameters = parameters  # type: ignore[assignment]
-
-        # Handle MySQL-specific placeholders (%s) which SQLGlot doesn't parse well
-        # If %s placeholders are present, handle them directly
-        mysql_placeholders_count = len(MYSQL_PLACEHOLDER_PATTERN.findall(sql))
-
-        if mysql_placeholders_count > 0:
-            # For MySQL format placeholders, minimal processing is needed
-            if actual_parameters is None:
-                if mysql_placeholders_count > 0:
-                    msg = f"asyncmy: SQL statement contains {mysql_placeholders_count} format placeholders ('%s'), but no parameters were provided. SQL: {sql}"
-                    raise ParameterStyleMismatchError(msg)
-                return sql, None
-
-            # Convert dict to tuple if needed
-            if is_dict(actual_parameters):
-                # MySQL's %s placeholders require positional params
-                msg = "asyncmy: Dictionary parameters provided with '%s' placeholders. MySQL format placeholders require tuple/list parameters."
-                raise ParameterStyleMismatchError(msg)
-
-            # Convert to tuple (handles both scalar and sequence cases)
-            if not isinstance(actual_parameters, (list, tuple)):
-                # Scalar parameter case
-                return sql, (actual_parameters,)
-
-            # Sequence parameter case - ensure appropriate length
-            if len(actual_parameters) != mysql_placeholders_count:  # type: ignore[arg-type]
-                msg = f"asyncmy: Parameter count mismatch. SQL expects {mysql_placeholders_count} '%s' placeholders, but {len(actual_parameters)} parameters were provided. SQL: {sql}"  # type: ignore[arg-type]
-                raise ParameterStyleMismatchError(msg)
-
-            return sql, tuple(actual_parameters)  # type: ignore[arg-type]
-
-        # Create a SQLStatement with MySQL dialect
-        statement = SQLStatement(sql, actual_parameters, kwargs=kwargs, dialect=self.dialect)
-
-        # Apply any filters
-        for filter_obj in current_filters:  # Use the modified list of filters
-            statement = statement.apply_filter(filter_obj)
-
-        # Process the statement for execution
-        processed_sql, processed_params, _ = statement.process()
-
-        # Convert parameters to the format expected by MySQL
-        if processed_params is None:
-            return processed_sql, None
-
-        # For MySQL, ensure parameters are in the right format
-        if is_dict(processed_params):
-            # Dictionary parameters are not well supported by asyncmy
-            msg = "asyncmy: Dictionary parameters are not supported for MySQL placeholders. Use sequence parameters."
-            raise ParameterStyleMismatchError(msg)
-
-        # For sequence parameters, ensure they're a tuple
-        if isinstance(processed_params, (list, tuple)):
-            return processed_sql, tuple(processed_params)
-
-        # For scalar parameter, wrap in a tuple
-        return processed_sql, (processed_params,)
-
-    # --- Public API Methods --- #
-    @overload
-    async def select(
-        self,
-        sql: str,
-        parameters: "Optional[StatementParameterType]" = None,
-        *filters: "StatementFilter",
-        connection: "Optional[AsyncmyConnection]" = None,
-        schema_type: None = None,
-        **kwargs: Any,
-    ) -> "Sequence[dict[str, Any]]": ...
-    @overload
-    async def select(
-        self,
-        sql: str,
-        parameters: "Optional[StatementParameterType]" = None,
-        *filters: "StatementFilter",
-        connection: "Optional[AsyncmyConnection]" = None,
-        schema_type: "type[ModelDTOT]",
-        **kwargs: Any,
-    ) -> "Sequence[ModelDTOT]": ...
-    async def select(
-        self,
-        sql: str,
-        parameters: "Optional[StatementParameterType]" = None,
-        *filters: "StatementFilter",
-        connection: "Optional[AsyncmyConnection]" = None,
-        schema_type: "Optional[type[ModelDTOT]]" = None,
-        **kwargs: 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, *filters, **kwargs)
-        async with self._with_cursor(connection) as cursor:
-            await cursor.execute(sql, parameters)
-            results = await cursor.fetchall()
-            if not results:
-                return []
-            column_names = [c[0] for c in cursor.description or []]
-            return self.to_schema([dict(zip(column_names, row)) for row in results], schema_type=schema_type)
-
-    @overload
-    async def select_one(
-        self,
-        sql: str,
-        parameters: "Optional[StatementParameterType]" = None,
-        *filters: "StatementFilter",
-        connection: "Optional[AsyncmyConnection]" = None,
-        schema_type: None = None,
-        **kwargs: Any,
-    ) -> "dict[str, Any]": ...
-    @overload
-    async def select_one(
-        self,
-        sql: str,
-        parameters: "Optional[StatementParameterType]" = None,
-        *filters: "StatementFilter",
-        connection: "Optional[AsyncmyConnection]" = None,
-        schema_type: "type[ModelDTOT]",
-        **kwargs: Any,
-    ) -> "ModelDTOT": ...
-    async def select_one(
-        self,
-        sql: str,
-        parameters: "Optional[StatementParameterType]" = None,
-        *filters: "StatementFilter",
-        connection: "Optional[AsyncmyConnection]" = None,
-        schema_type: "Optional[type[ModelDTOT]]" = None,
-        **kwargs: 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, *filters, **kwargs)
-        async with self._with_cursor(connection) as cursor:
-            await cursor.execute(sql, parameters)
-            result = await cursor.fetchone()
-            result = self.check_not_found(result)
-            column_names = [c[0] for c in cursor.description or []]
-            return self.to_schema(dict(zip(column_names, result)), schema_type=schema_type)
-
-    @overload
-    async def select_one_or_none(
-        self,
-        sql: str,
-        parameters: "Optional[StatementParameterType]" = None,
-        *filters: "StatementFilter",
-        connection: "Optional[AsyncmyConnection]" = 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,
-        *filters: "StatementFilter",
-        connection: "Optional[AsyncmyConnection]" = None,
-        schema_type: "type[ModelDTOT]",
-        **kwargs: Any,
-    ) -> "Optional[ModelDTOT]": ...
-    async def select_one_or_none(
-        self,
-        sql: str,
-        parameters: "Optional[StatementParameterType]" = None,
-        *filters: "StatementFilter",
-        connection: "Optional[AsyncmyConnection]" = None,
-        schema_type: "Optional[type[ModelDTOT]]" = None,
-        **kwargs: 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, *filters, **kwargs)
-        async with self._with_cursor(connection) as cursor:
+    async def _execute_statement(
+        self, statement: SQL, connection: "Optional[AsyncmyConnection]" = None, **kwargs: Any
+    ) -> Union[SelectResultDict, DMLResultDict, ScriptResultDict]:
+        if statement.is_script:
+            sql, _ = statement.compile(placeholder_style=ParameterStyle.STATIC)
+            return await self._execute_script(sql, connection=connection, **kwargs)
+
+        # Let the SQL object handle parameter style conversion based on dialect support
+        sql, params = statement.compile(placeholder_style=self.default_parameter_style)
+
+        if statement.is_many:
+            # Process parameter list through type coercion
+            params = self._process_parameters(params)
+            return await self._execute_many(sql, params, connection=connection, **kwargs)
+
+        # Process parameters through type coercion
+        params = self._process_parameters(params)
+        return await self._execute(sql, params, statement, connection=connection, **kwargs)
+
+    async def _execute(
+        self, sql: str, parameters: Any, statement: SQL, connection: "Optional[AsyncmyConnection]" = None, **kwargs: Any
+    ) -> Union[SelectResultDict, DMLResultDict]:
+        conn = self._connection(connection)
+        # AsyncMy doesn't like empty lists/tuples, convert to None
+        if not parameters:
+            parameters = None
+        async with self._get_cursor(conn) as cursor:
+            # AsyncMy expects list/tuple parameters or dict for named params
             await cursor.execute(sql, parameters)
-            result = await cursor.fetchone()
-            if result is None:
-                return None
-            column_names = [c[0] for c in cursor.description or []]
-            return self.to_schema(dict(zip(column_names, result)), schema_type=schema_type)
-
-    @overload
-    async def select_value(
-        self,
-        sql: str,
-        parameters: "Optional[StatementParameterType]" = None,
-        *filters: "StatementFilter",
-        connection: "Optional[AsyncmyConnection]" = None,
-        schema_type: None = None,
-        **kwargs: Any,
-    ) -> "Any": ...
-    @overload
-    async def select_value(
-        self,
-        sql: str,
-        parameters: "Optional[StatementParameterType]" = None,
-        *filters: "StatementFilter",
-        connection: "Optional[AsyncmyConnection]" = None,
-        schema_type: "type[T]",
-        **kwargs: Any,
-    ) -> "T": ...
-    async def select_value(
-        self,
-        sql: str,
-        parameters: "Optional[StatementParameterType]" = None,
-        *filters: "StatementFilter",
-        connection: "Optional[AsyncmyConnection]" = None,
-        schema_type: "Optional[type[T]]" = None,
-        **kwargs: Any,
-    ) -> "Union[T, Any]":
-        """Fetch a single value from the database.
-
-        Returns:
-            The first value from the first row of results.
-        """
-        connection = self._connection(connection)
-        sql, parameters = self._process_sql_params(sql, parameters, *filters, **kwargs)
-        async with self._with_cursor(connection) as cursor:
-            await cursor.execute(sql, parameters)
-            result = await cursor.fetchone()
-            result = self.check_not_found(result)
-            value = result[0]
-            if schema_type is not None:
-                return schema_type(value)  # type: ignore[call-arg]
-            return value
-
-    @overload
-    async def select_value_or_none(
-        self,
-        sql: str,
-        parameters: "Optional[StatementParameterType]" = None,
-        *filters: "StatementFilter",
-        connection: "Optional[AsyncmyConnection]" = None,
-        schema_type: None = None,
-        **kwargs: Any,
-    ) -> "Optional[Any]": ...
-    @overload
-    async def select_value_or_none(
-        self,
-        sql: str,
-        parameters: "Optional[StatementParameterType]" = None,
-        *filters: "StatementFilter",
-        connection: "Optional[AsyncmyConnection]" = None,
-        schema_type: "type[T]",
-        **kwargs: Any,
-    ) -> "Optional[T]": ...
-    async def select_value_or_none(
-        self,
-        sql: str,
-        parameters: "Optional[StatementParameterType]" = None,
-        *filters: "StatementFilter",
-        connection: "Optional[AsyncmyConnection]" = 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, *filters, **kwargs)
-        async with self._with_cursor(connection) as cursor:
-            await cursor.execute(sql, parameters)
-            result = await cursor.fetchone()
-            if result is None:
-                return None
-            value = result[0]
-            if schema_type is not None:
-                return schema_type(value)  # type: ignore[call-arg]
-            return value
-
-    async def insert_update_delete(
-        self,
-        sql: str,
-        parameters: "Optional[StatementParameterType]" = None,
-        *filters: "StatementFilter",
-        connection: "Optional[AsyncmyConnection]" = None,
-        **kwargs: Any,
-    ) -> int:
-        """Insert, update, or delete data from the database.
-
-        Returns:
-            Row count affected by the operation.
-        """
-        connection = self._connection(connection)
-        sql, parameters = self._process_sql_params(sql, parameters, *filters, **kwargs)
-        async with self._with_cursor(connection) as cursor:
-            await cursor.execute(sql, parameters)
-            return cursor.rowcount
-
-    @overload
-    async def insert_update_delete_returning(
-        self,
-        sql: str,
-        parameters: "Optional[StatementParameterType]" = None,
-        *filters: "StatementFilter",
-        connection: "Optional[AsyncmyConnection]" = None,
-        schema_type: None = None,
-        **kwargs: Any,
-    ) -> "dict[str, Any]": ...
-    @overload
-    async def insert_update_delete_returning(
-        self,
-        sql: str,
-        parameters: "Optional[StatementParameterType]" = None,
-        *filters: "StatementFilter",
-        connection: "Optional[AsyncmyConnection]" = None,
-        schema_type: "type[ModelDTOT]",
-        **kwargs: Any,
-    ) -> "ModelDTOT": ...
-    async def insert_update_delete_returning(
-        self,
-        sql: str,
-        parameters: "Optional[StatementParameterType]" = None,
-        *filters: "StatementFilter",
-        connection: "Optional[AsyncmyConnection]" = None,
-        schema_type: "Optional[type[ModelDTOT]]" = None,
-        **kwargs: Any,
-    ) -> "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, *filters, **kwargs)
-
-        async with self._with_cursor(connection) as cursor:
-            await cursor.execute(sql, parameters)
-            result = await cursor.fetchone()
-            if result is None:
-                return None
-            column_names = [c[0] for c in cursor.description or []]
-
-            # Convert to dict and 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,
-        connection: "Optional[AsyncmyConnection]" = None,
-        **kwargs: Any,
-    ) -> str:
-        """Execute a script.
-
-        Returns:
-            Status message for the operation.
-        """
-        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)
-            return f"Script executed successfully. Rows affected: {cursor.rowcount}"
-
-    def _connection(self, connection: "Optional[AsyncmyConnection]" = None) -> "AsyncmyConnection":
-        """Get the connection to use for the operation.
-
-        Args:
-            connection: Optional connection to use.
 
-        Returns:
-            The connection to use.
-        """
+            if self.returns_rows(statement.expression):
+                # For SELECT queries, fetch data and return SelectResultDict
+                data = await cursor.fetchall()
+                column_names = [desc[0] for desc in cursor.description or []]
+                result: SelectResultDict = {"data": data, "column_names": column_names, "rows_affected": len(data)}
+                return result
+
+            # For DML/DDL queries, return DMLResultDict
+            dml_result: DMLResultDict = {
+                "rows_affected": cursor.rowcount if cursor.rowcount is not None else -1,
+                "status_message": "OK",
+            }
+            return dml_result
+
+    async def _execute_many(
+        self, sql: str, param_list: Any, connection: "Optional[AsyncmyConnection]" = None, **kwargs: Any
+    ) -> DMLResultDict:
+        conn = self._connection(connection)
+
+        # Convert parameter list to proper format for executemany
+        params_list: list[Union[list[Any], tuple[Any, ...]]] = []
+        if param_list and isinstance(param_list, Sequence):
+            for param_set in param_list:
+                if isinstance(param_set, (list, tuple)):
+                    params_list.append(param_set)
+                elif param_set is None:
+                    params_list.append([])
+                else:
+                    params_list.append([param_set])
+
+        async with self._get_cursor(conn) as cursor:
+            await cursor.executemany(sql, params_list)
+            result: DMLResultDict = {
+                "rows_affected": cursor.rowcount if cursor.rowcount != -1 else len(params_list),
+                "status_message": "OK",
+            }
+            return result
+
+    async def _execute_script(
+        self, script: str, connection: "Optional[AsyncmyConnection]" = None, **kwargs: Any
+    ) -> ScriptResultDict:
+        conn = self._connection(connection)
+        # AsyncMy may not support multi-statement scripts without CLIENT_MULTI_STATEMENTS flag
+        # Use the shared implementation to split and execute statements individually
+        statements = self._split_script_statements(script)
+        statements_executed = 0
+
+        async with self._get_cursor(conn) as cursor:
+            for statement_str in statements:
+                if statement_str:
+                    await cursor.execute(statement_str)
+                    statements_executed += 1
+
+        result: ScriptResultDict = {"statements_executed": statements_executed, "status_message": "SCRIPT EXECUTED"}
+        return result
+
+    async def _ingest_arrow_table(self, table: "Any", table_name: str, mode: str = "append", **options: Any) -> int:
+        self._ensure_pyarrow_installed()
+        conn = self._connection(None)
+
+        async with self._get_cursor(conn) as cursor:
+            if mode == "replace":
+                await cursor.execute(f"TRUNCATE TABLE {table_name}")
+            elif mode == "create":
+                msg = "'create' mode is not supported for asyncmy ingestion."
+                raise NotImplementedError(msg)
+
+            data_for_ingest = table.to_pylist()
+            if not data_for_ingest:
+                return 0
+
+            # Generate column placeholders: %s, %s, etc.
+            num_columns = len(data_for_ingest[0])
+            placeholders = ", ".join("%s" for _ in range(num_columns))
+            sql = f"INSERT INTO {table_name} VALUES ({placeholders})"
+            await cursor.executemany(sql, data_for_ingest)
+            return cursor.rowcount if cursor.rowcount is not None else -1
+
+    async def _wrap_select_result(
+        self, statement: SQL, result: SelectResultDict, schema_type: "Optional[type[ModelDTOT]]" = None, **kwargs: Any
+    ) -> "Union[SQLResult[ModelDTOT], SQLResult[RowT]]":
+        data = result["data"]
+        column_names = result["column_names"]
+        rows_affected = result["rows_affected"]
+
+        if not data:
+            return SQLResult[RowT](
+                statement=statement, data=[], column_names=column_names, rows_affected=0, operation_type="SELECT"
+            )
+
+        rows_as_dicts = [dict(zip(column_names, row)) for row in data]
+
+        if schema_type:
+            converted_data = self.to_schema(data=rows_as_dicts, schema_type=schema_type)
+            return SQLResult[ModelDTOT](
+                statement=statement,
+                data=list(converted_data),
+                column_names=column_names,
+                rows_affected=rows_affected,
+                operation_type="SELECT",
+            )
+
+        return SQLResult[RowT](
+            statement=statement,
+            data=rows_as_dicts,
+            column_names=column_names,
+            rows_affected=rows_affected,
+            operation_type="SELECT",
+        )
+
+    async def _wrap_execute_result(
+        self, statement: SQL, result: Union[DMLResultDict, ScriptResultDict], **kwargs: Any
+    ) -> SQLResult[RowT]:
+        operation_type = "UNKNOWN"
+        if statement.expression:
+            operation_type = str(statement.expression.key).upper()
+
+        # Handle script results
+        if "statements_executed" in result:
+            script_result = cast("ScriptResultDict", result)
+            return SQLResult[RowT](
+                statement=statement,
+                data=[],
+                rows_affected=0,
+                operation_type="SCRIPT",
+                metadata={
+                    "status_message": script_result.get("status_message", ""),
+                    "statements_executed": script_result.get("statements_executed", -1),
+                },
+            )
+
+        # Handle DML results
+        dml_result = cast("DMLResultDict", result)
+        rows_affected = dml_result.get("rows_affected", -1)
+        status_message = dml_result.get("status_message", "")
+        return SQLResult[RowT](
+            statement=statement,
+            data=[],
+            rows_affected=rows_affected,
+            operation_type=operation_type,
+            metadata={"status_message": status_message},
+        )
+
+    def _connection(self, connection: Optional["AsyncmyConnection"] = None) -> "AsyncmyConnection":
+        """Get the connection to use for the operation."""
         return connection or self.connection