sqlspec 0.12.2__py3-none-any.whl → 0.13.1__py3-none-any.whl

sqlspec/statement/result.py

@@ -1,68 +1,23 @@
 """SQL statement result classes for handling different types of SQL operations."""
 
 from abc import ABC, abstractmethod
-
-# Import Mapping for type checking in __post_init__
 from collections.abc import Mapping, Sequence
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any, Generic, Optional, Union, cast
+from typing import TYPE_CHECKING, Any, Generic, Literal, Optional, Union
 
-from typing_extensions import TypedDict, TypeVar
+from typing_extensions import TypeVar
 
 from sqlspec.typing import ArrowTable, RowT
 
 if TYPE_CHECKING:
     from sqlspec.statement.sql import SQL
 
-__all__ = ("ArrowResult", "DMLResultDict", "SQLResult", "ScriptResultDict", "SelectResultDict", "StatementResult")
+__all__ = ("ArrowResult", "SQLResult", "StatementResult")
 
 
 T = TypeVar("T")
 
-
-class SelectResultDict(TypedDict):
-    """TypedDict for SELECT/RETURNING query results.
-
-    This structure is returned by drivers when executing SELECT queries
-    or DML queries with RETURNING clauses.
-    """
-
-    data: "list[Any]"
-    """List of rows returned by the query."""
-    column_names: "list[str]"
-    """List of column names in the result set."""
-    rows_affected: int
-    """Number of rows affected (-1 when unsupported)."""
-
-
-class DMLResultDict(TypedDict, total=False):
-    """TypedDict for DML (INSERT/UPDATE/DELETE) results without RETURNING.
-
-    This structure is returned by drivers when executing DML operations
-    that don't return data (no RETURNING clause).
-    """
-
-    rows_affected: int
-    """Number of rows affected by the operation."""
-    status_message: str
-    """Status message from the database (-1 when unsupported)."""
-    description: str
-    """Optional description of the operation."""
-
-
-class ScriptResultDict(TypedDict, total=False):
-    """TypedDict for script execution results.
-
-    This structure is returned by drivers when executing multi-statement
-    SQL scripts.
-    """
-
-    statements_executed: int
-    """Number of statements that were executed."""
-    status_message: str
-    """Overall status message from the script execution."""
-    description: str
-    """Optional description of the script execution."""
+OperationType = Literal["SELECT", "INSERT", "UPDATE", "DELETE", "EXECUTE", "SCRIPT"]
 
 
 @dataclass
@@ -133,9 +88,6 @@ class StatementResult(ABC, Generic[RowT]):
         self.metadata[key] = value
 
 
-# RowT is introduced for clarity within SQLResult, representing the type of a single row.
-
-
 @dataclass
 class SQLResult(StatementResult[RowT], Generic[RowT]):
     """Unified result class for SQL operations that return a list of rows
@@ -147,22 +99,16 @@ class SQLResult(StatementResult[RowT], Generic[RowT]):
     For script execution, this class also tracks multiple statement results and errors.
     """
 
+    data: "list[RowT]" = field(default_factory=list)
     error: Optional[Exception] = None
+    operation_type: OperationType = "SELECT"
     operation_index: Optional[int] = None
     pipeline_sql: Optional["SQL"] = None
     parameters: Optional[Any] = None
-
-    # Attributes primarily for SELECT-like results or results with column structure
     column_names: "list[str]" = field(default_factory=list)
-    total_count: Optional[int] = None  # Total rows if pagination/limit was involved
-    has_more: bool = False  # For pagination
-
-    # Attributes primarily for DML-like results
-    operation_type: str = "SELECT"  # Default, override for DML
+    total_count: Optional[int] = None
+    has_more: bool = False
     inserted_ids: "list[Union[int, str]]" = field(default_factory=list)
-    # rows_affected and last_inserted_id are inherited from StatementResult
-
-    # Attributes for script execution
     statement_results: "list[SQLResult[Any]]" = field(default_factory=list)
     """Individual statement results when executing scripts."""
     errors: "list[str]" = field(default_factory=list)
@@ -176,40 +122,35 @@ class SQLResult(StatementResult[RowT], Generic[RowT]):
         """Post-initialization to infer column names and total count if not provided."""
         if not self.column_names and self.data and isinstance(self.data[0], Mapping):
             self.column_names = list(self.data[0].keys())
-
         if self.total_count is None:
             self.total_count = len(self.data) if self.data is not None else 0
 
-        # If data is populated for a DML, it implies returning data.
-        # No separate returning_data field needed; self.data serves this purpose.
-
     def is_success(self) -> bool:
         """Check if the operation was successful.
+
         - For SELECT: True if data is not None and rows_affected is not negative.
         - For DML (INSERT, UPDATE, DELETE, EXECUTE): True if rows_affected is >= 0.
         - For SCRIPT: True if no errors and all statements succeeded.
         """
-        op_type_upper = self.operation_type.upper()
-
-        # For script execution, check if there are no errors and all statements succeeded
-        if op_type_upper == "SCRIPT" or self.statement_results:
-            return len(self.errors) == 0 and self.total_statements == self.successful_statements
-
-        if op_type_upper == "SELECT":
-            # For SELECT, success means we got some data container and rows_affected is not negative
-            data_success = self.data is not None
-            rows_success = self.rows_affected is None or self.rows_affected >= 0
-            return data_success and rows_success
-        if op_type_upper in {"INSERT", "UPDATE", "DELETE", "EXECUTE"}:
+        op_type = self.operation_type.upper()
+
+        if op_type == "SCRIPT" or self.statement_results:
+            return not self.errors and self.total_statements == self.successful_statements
+
+        if op_type == "SELECT":
+            return self.data is not None and (self.rows_affected is None or self.rows_affected >= 0)
+
+        if op_type in {"INSERT", "UPDATE", "DELETE", "EXECUTE"}:
             return self.rows_affected is not None and self.rows_affected >= 0
-        return False  # Should not happen if operation_type is one of the above
+
+        return False
 
     def get_data(self) -> "Union[list[RowT], dict[str, Any]]":
         """Get the data from the result.
+
         For regular operations, returns the list of rows.
         For script operations, returns a summary dictionary.
         """
-        # For script execution, return summary data
         if self.operation_type.upper() == "SCRIPT" or self.statement_results:
             return {
                 "total_statements": self.total_statements,
@@ -219,11 +160,7 @@ class SQLResult(StatementResult[RowT], Generic[RowT]):
                 "statement_results": self.statement_results,
                 "total_rows_affected": self.get_total_rows_affected(),
             }
-
-        # For regular operations, return the data as usual
-        return cast("list[RowT]", self.data)
-
-    # --- Script execution methods ---
+        return self.data
 
     def add_statement_result(self, result: "SQLResult[Any]") -> None:
         """Add a statement result to the script execution results."""
@@ -245,15 +182,10 @@ class SQLResult(StatementResult[RowT], Generic[RowT]):
     def get_total_rows_affected(self) -> int:
         """Get the total number of rows affected across all statements."""
         if self.statement_results:
-            # For script execution, sum up rows affected from all statements
-            total = 0
-            for stmt_result in self.statement_results:
-                if stmt_result.rows_affected is not None and stmt_result.rows_affected >= 0:
-                    # Only count non-negative values, -1 indicates failure
-                    total += stmt_result.rows_affected
-            return total
-        # For single statement execution
-        return max(self.rows_affected or 0, 0)  # Treat negative values as 0
+            return sum(
+                stmt.rows_affected for stmt in self.statement_results if stmt.rows_affected and stmt.rows_affected > 0
+            )
+        return self.rows_affected if self.rows_affected and self.rows_affected > 0 else 0
 
     @property
     def num_rows(self) -> int:
@@ -272,8 +204,6 @@ class SQLResult(StatementResult[RowT], Generic[RowT]):
         """Check if there are any errors from script execution."""
         return len(self.errors) > 0
 
-    # --- Existing methods for regular operations ---
-
     def get_first(self) -> "Optional[RowT]":
         """Get the first row from the result, if any."""
         return self.data[0] if self.data else None
@@ -286,25 +216,10 @@ class SQLResult(StatementResult[RowT], Generic[RowT]):
         """Check if the result set (self.data) is empty."""
         return not self.data
 
-    # --- Methods related to DML operations ---
     def get_affected_count(self) -> int:
         """Get the number of rows affected by a DML operation."""
         return self.rows_affected or 0
 
-    def get_inserted_id(self) -> "Optional[Union[int, str]]":
-        """Get the last inserted ID (typically for single row inserts)."""
-        return self.last_inserted_id
-
-    def get_inserted_ids(self) -> "list[Union[int, str]]":
-        """Get all inserted IDs (useful for batch inserts)."""
-        return self.inserted_ids
-
-    def get_returning_data(self) -> "list[RowT]":
-        """Get data returned by RETURNING clauses.
-        This is effectively self.data for this unified class.
-        """
-        return cast("list[RowT]", self.data)
-
     def was_inserted(self) -> bool:
         """Check if this was an INSERT operation."""
         return self.operation_type.upper() == "INSERT"
@@ -340,9 +255,7 @@ class SQLResult(StatementResult[RowT], Generic[RowT]):
         if self.data is None:
             msg = "No data available"
             raise TypeError(msg)
-        return cast("RowT", self.data[index])
-
-    # --- SQLAlchemy-style convenience methods ---
+        return self.data[index]
 
     def all(self) -> "list[RowT]":
         """Return all rows as a list.
@@ -352,7 +265,7 @@ class SQLResult(StatementResult[RowT], Generic[RowT]):
         """
         if self.data is None:
             return []
-        return cast("list[RowT]", self.data)
+        return self.data
 
     def one(self) -> "RowT":
         """Return exactly one row.
@@ -369,7 +282,7 @@ class SQLResult(StatementResult[RowT], Generic[RowT]):
         if len(self.data) > 1:
             msg = f"Multiple results found ({len(self.data)}), exactly one row expected"
             raise ValueError(msg)
-        return cast("RowT", self.data[0])
+        return self.data[0]
 
     def one_or_none(self) -> "Optional[RowT]":
         """Return at most one row.
@@ -385,7 +298,7 @@ class SQLResult(StatementResult[RowT], Generic[RowT]):
         if len(self.data) > 1:
             msg = f"Multiple results found ({len(self.data)}), at most one row expected"
             raise ValueError(msg)
-        return cast("RowT", self.data[0])
+        return self.data[0]
 
     def scalar(self) -> Any:
         """Return the first column of the first row.
@@ -398,19 +311,16 @@ class SQLResult(StatementResult[RowT], Generic[RowT]):
         """
         row = self.one()
         if isinstance(row, Mapping):
-            # For dict-like rows, get the first column value
             if not row:
                 msg = "Row has no columns"
                 raise ValueError(msg)
-            first_key = cast("str", next(iter(row.keys())))
-            return cast("Any", row[first_key])
+            first_key = next(iter(row.keys()))
+            return row[first_key]
         if isinstance(row, Sequence) and not isinstance(row, (str, bytes)):
-            # For tuple/list-like rows
             if len(row) == 0:
                 msg = "Row has no columns"
                 raise ValueError(msg)
-            return cast("Any", row[0])
-        # For scalar values returned directly
+            return row[0]
         return row
 
     def scalar_or_none(self) -> Any:
@@ -429,11 +339,9 @@ class SQLResult(StatementResult[RowT], Generic[RowT]):
             first_key = next(iter(row.keys()))
             return row[first_key]
         if isinstance(row, Sequence) and not isinstance(row, (str, bytes)):
-            # For tuple/list-like rows
             if len(row) == 0:
                 return None
-            return cast("Any", row[0])
-        # For scalar values returned directly
+            return row[0]
         return row
 
 
@@ -457,12 +365,12 @@ class ArrowResult(StatementResult[ArrowTable]):
     """The result data from the operation."""
 
     def is_success(self) -> bool:
-        """Check if the Arrow operation was successful.
+        """Check if the operation was successful.
 
         Returns:
-            True if the operation completed successfully and has valid Arrow data.
+            True if Arrow table data is available, False otherwise.
         """
-        return bool(self.data)
+        return self.data is not None
 
     def get_data(self) -> "ArrowTable":
         """Get the Apache Arrow Table from the result.

sqlspec/statement/splitter.py

@@ -115,7 +115,6 @@ class DialectConfig(ABC):
         # 3. Dynamically build and insert keyword/separator patterns
         all_keywords = self.block_starters | self.block_enders | self.batch_separators
         if all_keywords:
-            # Sort by length descending to match longer keywords first
             sorted_keywords = sorted(all_keywords, key=len, reverse=True)
             patterns.append((TokenType.KEYWORD, r"\b(" + "|".join(re.escape(kw) for kw in sorted_keywords) + r")\b"))
 
@@ -337,7 +336,6 @@ class PostgreSQLDialectConfig(DialectConfig):
         tag = start_match.group(0)  # The full opening tag, e.g., "$tag$"
         content_start = position + len(tag)
 
-        # Find the corresponding closing tag
         try:
             content_end = text.index(tag, content_start)
             full_value = text[position : content_end + len(tag)]
@@ -502,7 +500,6 @@ class StatementSplitter:
                     column = pos - line_start + 1
                     token = pattern(sql, pos, line, column)
                     if token:
-                        # Update position tracking
                         newlines = token.value.count("\n")
                         if newlines > 0:
                             line += newlines
@@ -520,7 +517,6 @@ class StatementSplitter:
                         value = match.group(0)
                         column = pos - line_start + 1
 
-                        # Update line tracking
                         newlines = value.count("\n")
                         if newlines > 0:
                             line += newlines
@@ -544,7 +540,6 @@ class StatementSplitter:
         current_statement_chars = []
         block_stack = []
 
-        # Convert token generator to list so we can look ahead
         all_tokens = list(self._tokenize(sql))
 
         for token_idx, token in enumerate(all_tokens):
@@ -559,7 +554,6 @@ class StatementSplitter:
             current_statement_tokens.append(token)
             token_upper = token.value.upper()
 
-            # Update block nesting
             if token.type == TokenType.KEYWORD:
                 if token_upper in self.dialect.block_starters:
                     block_stack.append(token_upper)
@@ -567,7 +561,6 @@ class StatementSplitter:
                         msg = f"Maximum nesting depth ({self.dialect.max_nesting_depth}) exceeded"
                         raise ValueError(msg)
                 elif token_upper in self.dialect.block_enders:
-                    # Check if this is actually a block ender (not END IF, END LOOP, etc.)
                     if block_stack and self.dialect.is_real_block_ender(all_tokens, token_idx):
                         block_stack.pop()
 
@@ -576,7 +569,6 @@ class StatementSplitter:
             if not block_stack:  # Only terminate when not inside a block
                 if token.type == TokenType.TERMINATOR:
                     if token.value in self.dialect.statement_terminators:
-                        # Check if we should delay this termination (e.g., for Oracle END; /)
                         should_delay = self.dialect.should_delay_semicolon_termination(all_tokens, token_idx)
 
                         # Also check if there's a batch separator coming up (for T-SQL GO)
@@ -601,7 +593,6 @@ class StatementSplitter:
                 # Save the statement
                 statement = "".join(current_statement_chars).strip()
 
-                # Determine if this is a PL/SQL block
                 is_plsql_block = self._is_plsql_block(current_statement_tokens)
 
                 # Optionally strip the trailing terminator
@@ -619,7 +610,6 @@ class StatementSplitter:
                 current_statement_tokens = []
                 current_statement_chars = []
 
-        # Handle any remaining content
         if current_statement_chars:
             statement = "".join(current_statement_chars).strip()
             if statement and self._contains_executable_content(statement):
@@ -637,7 +627,6 @@ class StatementSplitter:
         Returns:
             True if this is a PL/SQL block (BEGIN...END or DECLARE...END)
         """
-        # Find the first meaningful keyword token (skip whitespace and comments)
         for token in tokens:
             if token.type == TokenType.KEYWORD:
                 return token.value.upper() in {"BEGIN", "DECLARE"}
@@ -655,7 +644,6 @@ class StatementSplitter:
         # Tokenize the statement to check its content
         tokens = list(self._tokenize(statement))
 
-        # Check if there are any non-comment, non-whitespace tokens
         for token in tokens:
             if token.type not in {TokenType.WHITESPACE, TokenType.COMMENT_LINE, TokenType.COMMENT_BLOCK}:
                 return True