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

sqlspec/adapters/duckdb/driver.py

@@ -2,12 +2,14 @@ import contextlib
 import uuid
 from collections.abc import Generator
 from contextlib import contextmanager
-from typing import TYPE_CHECKING, Any, ClassVar, Optional, Union, cast
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, ClassVar, Optional, Union
 
 from duckdb import DuckDBPyConnection
 from sqlglot import exp
 
 from sqlspec.driver import SyncDriverAdapterProtocol
+from sqlspec.driver.connection import managed_transaction_sync
 from sqlspec.driver.mixins import (
     SQLTranslatorMixin,
     SyncPipelinedExecutionMixin,
@@ -15,10 +17,11 @@ from sqlspec.driver.mixins import (
     ToSchemaMixin,
     TypeCoercionMixin,
 )
+from sqlspec.driver.parameters import normalize_parameter_sequence
 from sqlspec.statement.parameters import ParameterStyle
-from sqlspec.statement.result import ArrowResult, DMLResultDict, ScriptResultDict, SelectResultDict, SQLResult
+from sqlspec.statement.result import ArrowResult, SQLResult
 from sqlspec.statement.sql import SQL, SQLConfig
-from sqlspec.typing import ArrowTable, DictRow, ModelDTOT, RowT
+from sqlspec.typing import ArrowTable, DictRow, RowT
 from sqlspec.utils.logging import get_logger
 
 if TYPE_CHECKING:
@@ -81,136 +84,129 @@ class DuckDBDriver(
 
     def _execute_statement(
         self, statement: SQL, connection: Optional["DuckDBConnection"] = None, **kwargs: Any
-    ) -> "Union[SelectResultDict, DMLResultDict, ScriptResultDict]":
+    ) -> SQLResult[RowT]:
         if statement.is_script:
             sql, _ = statement.compile(placeholder_style=ParameterStyle.STATIC)
             return self._execute_script(sql, connection=connection, **kwargs)
 
+        sql, params = statement.compile(placeholder_style=self.default_parameter_style)
+        params = self._process_parameters(params)
+
         if statement.is_many:
-            sql, params = statement.compile(placeholder_style=self.default_parameter_style)
-            params = self._process_parameters(params)
             return self._execute_many(sql, params, connection=connection, **kwargs)
 
-        sql, params = statement.compile(placeholder_style=self.default_parameter_style)
-        params = self._process_parameters(params)
         return self._execute(sql, params, statement, connection=connection, **kwargs)
 
     def _execute(
         self, sql: str, parameters: Any, statement: SQL, connection: Optional["DuckDBConnection"] = None, **kwargs: Any
-    ) -> "Union[SelectResultDict, DMLResultDict]":
-        conn = self._connection(connection)
+    ) -> SQLResult[RowT]:
+        # Use provided connection or driver's default connection
+        conn = connection if connection is not None else self._connection(None)
+
+        with managed_transaction_sync(conn, auto_commit=True) as txn_conn:
+            # Normalize parameters using consolidated utility
+            normalized_params = normalize_parameter_sequence(parameters)
+            final_params = normalized_params or []
+
+            if self.returns_rows(statement.expression):
+                result = txn_conn.execute(sql, final_params)
+                fetched_data = result.fetchall()
+                column_names = [col[0] for col in result.description or []]
+
+                if fetched_data and isinstance(fetched_data[0], tuple):
+                    dict_data = [dict(zip(column_names, row)) for row in fetched_data]
+                else:
+                    dict_data = fetched_data
+
+                return SQLResult[RowT](
+                    statement=statement,
+                    data=dict_data,  # type: ignore[arg-type]
+                    column_names=column_names,
+                    rows_affected=len(dict_data),
+                    operation_type="SELECT",
+                )
 
-        if self.returns_rows(statement.expression):
-            result = conn.execute(sql, parameters or [])
-            fetched_data = result.fetchall()
-            column_names = [col[0] for col in result.description or []]
-            return {"data": fetched_data, "column_names": column_names, "rows_affected": len(fetched_data)}
-
-        with self._get_cursor(conn) as cursor:
-            cursor.execute(sql, parameters or [])
-            # DuckDB returns -1 for rowcount on DML operations
-            # However, fetchone() returns the actual affected row count as (count,)
-            rows_affected = cursor.rowcount
-            if rows_affected < 0:
-                try:
-                    # Get actual affected row count from fetchone()
-                    fetch_result = cursor.fetchone()
-                    if fetch_result and isinstance(fetch_result, (tuple, list)) and len(fetch_result) > 0:
-                        rows_affected = fetch_result[0]
-                    else:
-                        rows_affected = 0
-                except Exception:
-                    # Fallback to 1 if fetchone fails
-                    rows_affected = 1
-            return {"rows_affected": rows_affected}
+            with self._get_cursor(txn_conn) as cursor:
+                cursor.execute(sql, final_params)
+                # DuckDB returns -1 for rowcount on DML operations
+                # However, fetchone() returns the actual affected row count as (count,)
+                rows_affected = cursor.rowcount
+                if rows_affected < 0:
+                    try:
+                        fetch_result = cursor.fetchone()
+                        if fetch_result and isinstance(fetch_result, (tuple, list)) and len(fetch_result) > 0:
+                            rows_affected = fetch_result[0]
+                        else:
+                            rows_affected = 0
+                    except Exception:
+                        rows_affected = 1
+
+                return SQLResult(
+                    statement=statement,
+                    data=[],
+                    rows_affected=rows_affected,
+                    operation_type=self._determine_operation_type(statement),
+                    metadata={"status_message": "OK"},
+                )
 
     def _execute_many(
         self, sql: str, param_list: Any, connection: Optional["DuckDBConnection"] = None, **kwargs: Any
-    ) -> "DMLResultDict":
-        conn = self._connection(connection)
-        param_list = param_list or []
-
-        # DuckDB throws an error if executemany is called with empty parameter list
-        if not param_list:
-            return {"rows_affected": 0}
-        with self._get_cursor(conn) as cursor:
-            cursor.executemany(sql, param_list)
-            # DuckDB returns -1 for rowcount on DML operations
-            # For executemany, fetchone() only returns the count from the last operation,
-            # so use parameter list length as the most accurate estimate
-            rows_affected = cursor.rowcount if cursor.rowcount >= 0 else len(param_list)
-            return {"rows_affected": rows_affected}
+    ) -> SQLResult[RowT]:
+        # Use provided connection or driver's default connection
+        conn = connection if connection is not None else self._connection(None)
+
+        with managed_transaction_sync(conn, auto_commit=True) as txn_conn:
+            # Normalize parameter list using consolidated utility
+            normalized_param_list = normalize_parameter_sequence(param_list)
+            final_param_list = normalized_param_list or []
+
+            # DuckDB throws an error if executemany is called with empty parameter list
+            if not final_param_list:
+                return SQLResult(
+                    statement=SQL(sql, _dialect=self.dialect),
+                    data=[],
+                    rows_affected=0,
+                    operation_type="EXECUTE",
+                    metadata={"status_message": "OK"},
+                )
+
+            with self._get_cursor(txn_conn) as cursor:
+                cursor.executemany(sql, final_param_list)
+                # DuckDB returns -1 for rowcount on DML operations
+                # For executemany, fetchone() only returns the count from the last operation,
+                # so use parameter list length as the most accurate estimate
+                rows_affected = cursor.rowcount if cursor.rowcount >= 0 else len(final_param_list)
+                return SQLResult(
+                    statement=SQL(sql, _dialect=self.dialect),
+                    data=[],
+                    rows_affected=rows_affected,
+                    operation_type="EXECUTE",
+                    metadata={"status_message": "OK"},
+                )
 
     def _execute_script(
         self, script: str, connection: Optional["DuckDBConnection"] = None, **kwargs: Any
-    ) -> "ScriptResultDict":
-        conn = self._connection(connection)
-        with self._get_cursor(conn) as cursor:
-            cursor.execute(script)
-
-        return {
-            "statements_executed": -1,
-            "status_message": "Script executed successfully.",
-            "description": "The script was sent to the database.",
-        }
-
-    def _wrap_select_result(
-        self, statement: SQL, result: "SelectResultDict", schema_type: Optional[type[ModelDTOT]] = None, **kwargs: Any
-    ) -> Union[SQLResult[ModelDTOT], SQLResult[RowT]]:
-        fetched_tuples = result["data"]
-        column_names = result["column_names"]
-        rows_affected = result["rows_affected"]
-
-        rows_as_dicts: list[dict[str, Any]] = [dict(zip(column_names, row)) for row in fetched_tuples]
-
-        logger.debug("Query returned %d rows", len(rows_as_dicts))
-
-        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",
-            )
+    ) -> SQLResult[RowT]:
+        # Use provided connection or driver's default connection
+        conn = connection if connection is not None else self._connection(None)
 
-        return SQLResult[RowT](
-            statement=statement,
-            data=rows_as_dicts,
-            column_names=column_names,
-            rows_affected=rows_affected,
-            operation_type="SELECT",
-        )
+        with managed_transaction_sync(conn, auto_commit=True) as txn_conn:
+            with self._get_cursor(txn_conn) as cursor:
+                cursor.execute(script)
 
-    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()
-
-        if "statements_executed" in result:
-            script_result = cast("ScriptResultDict", result)
-            return SQLResult[RowT](
-                statement=statement,
+            return SQLResult(
+                statement=SQL(script, _dialect=self.dialect).as_script(),
                 data=[],
                 rows_affected=0,
-                operation_type=operation_type or "SCRIPT",
-                metadata={"status_message": script_result.get("status_message", "")},
+                operation_type="SCRIPT",
+                metadata={
+                    "status_message": "Script executed successfully.",
+                    "description": "The script was sent to the database.",
+                },
+                total_statements=-1,
+                successful_statements=-1,
             )
 
-        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},
-        )
-
     # ============================================================================
     # DuckDB Native Arrow Support
     # ============================================================================
@@ -251,7 +247,7 @@ class DuckDBDriver(
                 return True
         return False
 
-    def _export_native(self, query: str, destination_uri: str, format: str, **options: Any) -> int:
+    def _export_native(self, query: str, destination_uri: Union[str, Path], format: str, **options: Any) -> int:
         conn = self._connection(None)
         copy_options: list[str] = []
 
@@ -283,19 +279,21 @@ class DuckDBDriver(
             raise ValueError(msg)
 
         options_str = f"({', '.join(copy_options)})" if copy_options else ""
-        copy_sql = f"COPY ({query}) TO '{destination_uri}' {options_str}"
+        copy_sql = f"COPY ({query}) TO '{destination_uri!s}' {options_str}"
         result_rel = conn.execute(copy_sql)
         result = result_rel.fetchone() if result_rel else None
         return result[0] if result else 0
 
-    def _import_native(self, source_uri: str, table_name: str, format: str, mode: str, **options: Any) -> int:
+    def _import_native(
+        self, source_uri: Union[str, Path], table_name: str, format: str, mode: str, **options: Any
+    ) -> int:
         conn = self._connection(None)
         if format == "parquet":
-            read_func = f"read_parquet('{source_uri}')"
+            read_func = f"read_parquet('{source_uri!s}')"
         elif format == "csv":
-            read_func = f"read_csv_auto('{source_uri}')"
+            read_func = f"read_csv_auto('{source_uri!s}')"
         elif format == "json":
-            read_func = f"read_json_auto('{source_uri}')"
+            read_func = f"read_json_auto('{source_uri!s}')"
         else:
             msg = f"Unsupported format for DuckDB native import: {format}"
             raise ValueError(msg)
@@ -320,16 +318,16 @@ class DuckDBDriver(
         return int(count_result[0]) if count_result else 0
 
     def _read_parquet_native(
-        self, source_uri: str, columns: Optional[list[str]] = None, **options: Any
+        self, source_uri: Union[str, Path], columns: Optional[list[str]] = None, **options: Any
     ) -> "SQLResult[dict[str, Any]]":
         conn = self._connection(None)
         if isinstance(source_uri, list):
             file_list = "[" + ", ".join(f"'{f}'" for f in source_uri) + "]"
             read_func = f"read_parquet({file_list})"
-        elif "*" in source_uri or "?" in source_uri:
-            read_func = f"read_parquet('{source_uri}')"
+        elif "*" in str(source_uri) or "?" in str(source_uri):
+            read_func = f"read_parquet('{source_uri!s}')"
         else:
-            read_func = f"read_parquet('{source_uri}')"
+            read_func = f"read_parquet('{source_uri!s}')"
 
         column_list = ", ".join(columns) if columns else "*"
         query = f"SELECT {column_list} FROM {read_func}"
@@ -350,10 +348,16 @@ class DuckDBDriver(
         rows = [{col: arrow_dict[col][i] for col in column_names} for i in range(num_rows)]
 
         return SQLResult[dict[str, Any]](
-            statement=SQL(query), data=rows, column_names=column_names, rows_affected=num_rows, operation_type="SELECT"
+            statement=SQL(query, _dialect=self.dialect),
+            data=rows,
+            column_names=column_names,
+            rows_affected=num_rows,
+            operation_type="SELECT",
         )
 
-    def _write_parquet_native(self, data: Union[str, "ArrowTable"], destination_uri: str, **options: Any) -> None:
+    def _write_parquet_native(
+        self, data: Union[str, "ArrowTable"], destination_uri: Union[str, Path], **options: Any
+    ) -> None:
         conn = self._connection(None)
         copy_options: list[str] = ["FORMAT PARQUET"]
         if "compression" in options:
@@ -364,18 +368,22 @@ class DuckDBDriver(
         options_str = f"({', '.join(copy_options)})"
 
         if isinstance(data, str):
-            copy_sql = f"COPY ({data}) TO '{destination_uri}' {options_str}"
+            copy_sql = f"COPY ({data}) TO '{destination_uri!s}' {options_str}"
             conn.execute(copy_sql)
         else:
             temp_name = f"_arrow_data_{uuid.uuid4().hex[:8]}"
             conn.register(temp_name, data)
             try:
-                copy_sql = f"COPY {temp_name} TO '{destination_uri}' {options_str}"
+                copy_sql = f"COPY {temp_name} TO '{destination_uri!s}' {options_str}"
                 conn.execute(copy_sql)
             finally:
                 with contextlib.suppress(Exception):
                     conn.unregister(temp_name)
 
+    def _connection(self, connection: Optional["DuckDBConnection"] = None) -> "DuckDBConnection":
+        """Get the connection to use for the operation."""
+        return connection or self.connection
+
     def _ingest_arrow_table(self, table: "ArrowTable", table_name: str, mode: str = "create", **options: Any) -> int:
         """DuckDB-optimized Arrow table ingestion using native registration."""
         self._ensure_pyarrow_installed()
@@ -404,7 +412,7 @@ class DuckDBDriver(
                 msg = f"Unsupported mode: {mode}"
                 raise ValueError(msg)
 
-            result = self.execute(SQL(sql_expr.sql(dialect=self.dialect)))
+            result = self.execute(SQL(sql_expr.sql(dialect=self.dialect), _dialect=self.dialect))
             return result.rows_affected or table.num_rows
         finally:
             with contextlib.suppress(Exception):

sqlspec/adapters/oracledb/config.py

@@ -4,7 +4,6 @@ import contextlib
 import logging
 from collections.abc import AsyncGenerator
 from contextlib import asynccontextmanager
-from dataclasses import replace
 from typing import TYPE_CHECKING, Any, ClassVar, Optional, cast
 
 import oracledb
@@ -293,15 +292,16 @@ class OracleSyncConfig(SyncDatabaseConfig[OracleSyncConnection, "ConnectionPool"
             An OracleSyncDriver instance.
         """
         with self.provide_connection(*args, **kwargs) as conn:
-            # Create statement config with parameter style info if not already set
             statement_config = self.statement_config
+            # Inject parameter style info if not already set
             if statement_config.allowed_parameter_styles is None:
+                from dataclasses import replace
+
                 statement_config = replace(
                     statement_config,
                     allowed_parameter_styles=self.supported_parameter_styles,
                     target_parameter_style=self.preferred_parameter_style,
                 )
-
             driver = self.driver_type(connection=conn, config=statement_config)
             yield driver
 
@@ -602,15 +602,16 @@ class OracleAsyncConfig(AsyncDatabaseConfig[OracleAsyncConnection, "AsyncConnect
             An OracleAsyncDriver instance.
         """
         async with self.provide_connection(*args, **kwargs) as conn:
-            # Create statement config with parameter style info if not already set
             statement_config = self.statement_config
+            # Inject parameter style info if not already set
             if statement_config.allowed_parameter_styles is None:
+                from dataclasses import replace
+
                 statement_config = replace(
                     statement_config,
                     allowed_parameter_styles=self.supported_parameter_styles,
                     target_parameter_style=self.preferred_parameter_style,
                 )
-
             driver = self.driver_type(connection=conn, config=statement_config)
             yield driver