"""SQL result classes for query execution results.
This module provides result classes for handling SQL query execution results
including regular results and Apache Arrow format results.
Classes:
StatementResult: Abstract base class for SQL results.
SQLResult: Standard implementation for regular results.
ArrowResult: Apache Arrow format results for data interchange.
"""
from abc import abstractmethod
from typing import TYPE_CHECKING, Any, Final, Literal, cast, overload
from mypy_extensions import mypyc_attr
from typing_extensions import TypeVar
from sqlspec.core.statement import SQL
from sqlspec.exceptions import MultipleResultsFoundError
from sqlspec.storage import (
AsyncStoragePipeline,
StorageDestination,
StorageFormat,
StorageTelemetry,
SyncStoragePipeline,
)
from sqlspec.utils.arrow_helpers import (
arrow_reader_to_return_format,
arrow_table_column_names,
arrow_table_num_columns,
arrow_table_num_rows,
arrow_table_to_pandas,
arrow_table_to_polars,
arrow_table_to_pylist,
arrow_table_to_return_format,
cast_arrow_table_schema,
convert_dict_to_arrow,
ensure_arrow_table,
)
from sqlspec.utils.module_loader import ensure_pandas, ensure_polars
from sqlspec.utils.schema import to_schema
if TYPE_CHECKING:
from collections.abc import Iterator, Sequence
from sqlspec.core.compiler import OperationType
from sqlspec.typing import ArrowReturnFormat, ArrowTable, PandasDataFrame, PolarsDataFrame, SchemaT
__all__ = ("ArrowResult", "DMLResult", "EmptyResult", "SQLResult", "StackResult", "StatementResult")
T = TypeVar("T")
RowFormat = Literal["dict", "tuple", "record"]
_EMPTY_RESULT_STATEMENT: Final = SQL("-- empty stack result --")
_EMPTY_RESULT_DATA: Final[tuple[Any, ...]] = ()
_DEFAULT_DML_METADATA: Final[dict[str, Any]] = {}
_TWO_COLUMN_THRESHOLD: Final[int] = 2
@mypyc_attr(allow_interpreted_subclasses=False)
class StatementResult:
"""Abstract base class for SQL statement execution results.
Provides a common interface for handling different types of SQL operation
results. Subclasses implement specific behavior for SELECT, INSERT, UPDATE,
DELETE, and script operations.
Attributes:
statement: The original SQL statement that was executed.
data: The result data from the operation.
rows_affected: Number of rows affected by the operation.
last_inserted_id: Last inserted ID from INSERT operations.
execution_time: Time taken to execute the statement in seconds.
metadata: Additional metadata about the operation.
"""
__slots__ = (
"_operation_type",
"data",
"execution_time",
"last_inserted_id",
"metadata",
"rows_affected",
"statement",
)
_operation_type: "OperationType"
def __init__(
self,
statement: "SQL",
data: Any = None,
rows_affected: int = 0,
last_inserted_id: int | str | None = None,
execution_time: float | None = None,
metadata: "dict[str, Any] | None" = None,
) -> None:
"""Initialize statement result.
Args:
statement: The original SQL statement that was executed.
data: The result data from the operation.
rows_affected: Number of rows affected by the operation.
last_inserted_id: Last inserted ID from the operation.
execution_time: Time taken to execute the statement in seconds.
metadata: Additional metadata about the operation.
"""
self.statement = statement
self._operation_type = statement.operation_type
self.data = data
self.rows_affected = rows_affected
self.last_inserted_id = last_inserted_id
self.execution_time = execution_time
self.metadata = metadata if metadata is not None else {}