feat: fix g+r/g+t keybindings and add multi-statement support to atomic execution

- Add missing action_g_execute_query and action_g_execute_query_atomic
  handlers so g+r and g+t work from the g menu
- Modify atomic_execute to split and execute statements individually,
  collecting results into MultiStatementResult
- Update _run_query_atomic_async to display multiple result tables
- All statements run in same transaction; rollback on any failure

Author: Maxteabag

sqlit/domains/query/app/transaction.py

@@ -13,6 +13,7 @@
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any
 
+from .multi_statement import MultiStatementResult
 from .query_service import KeywordQueryAnalyzer, NonQueryResult, QueryKind, QueryResult
 
 if TYPE_CHECKING:
@@ -254,25 +255,35 @@ def _execute_on_connection(
             rows_affected = self.provider.query_executor.execute_non_query(conn, sql)
             return NonQueryResult(rows_affected=rows_affected)
 
-    def atomic_execute(self, sql: str, max_rows: int | None = None) -> QueryResult | NonQueryResult:
+    def atomic_execute(
+        self, sql: str, max_rows: int | None = None
+    ) -> QueryResult | NonQueryResult | MultiStatementResult:
         """Execute SQL atomically (all-or-nothing).
 
         Wraps the SQL in BEGIN/COMMIT and rolls back on any error.
+        Supports multiple statements, returning results for each.
 
         Args:
             sql: SQL statement(s) to execute atomically.
             max_rows: Maximum rows to fetch for SELECT queries.
 
         Returns:
-            Result of the last statement.
+            For single statement: QueryResult or NonQueryResult.
+            For multiple statements: MultiStatementResult with all results.
 
         Raises:
             Exception: If any statement fails (after rollback).
         """
-        from .multi_statement import normalize_for_execution
+        from .multi_statement import (
+            MultiStatementResult,
+            StatementResult,
+            normalize_for_execution,
+            split_statements,
+        )
 
         # Normalize SQL: convert blank-line-separated to semicolon-separated
         sql = normalize_for_execution(sql)
+        statements = split_statements(sql)
 
         # Create a dedicated connection for this atomic operation
         conn = self.provider.connection_factory.connect(self.config)
@@ -285,16 +296,56 @@ def atomic_execute(self, sql: str, max_rows: int | None = None) -> QueryResult |
             # Start transaction
             self.provider.query_executor.execute_non_query(conn, "BEGIN")
 
-            # Execute the SQL
-            result = self._execute_on_connection(conn, sql, max_rows)
+            # Single statement - return simple result for backwards compatibility
+            if len(statements) <= 1:
+                result = self._execute_on_connection(conn, sql, max_rows)
+                self.provider.query_executor.execute_non_query(conn, "COMMIT")
+                return result
 
-            # Commit
+            # Multiple statements - execute each and collect results
+            results: list[StatementResult] = []
+            for i, statement in enumerate(statements):
+                try:
+                    result = self._execute_on_connection(conn, statement, max_rows)
+                    results.append(
+                        StatementResult(
+                            statement=statement,
+                            result=result,
+                            success=True,
+                            error=None,
+                        )
+                    )
+                except Exception as e:
+                    # Record the error
+                    results.append(
+                        StatementResult(
+                            statement=statement,
+                            result=None,
+                            success=False,
+                            error=str(e),
+                        )
+                    )
+                    # Rollback and return partial results
+                    try:
+                        self.provider.query_executor.execute_non_query(conn, "ROLLBACK")
+                    except Exception:
+                        pass
+                    return MultiStatementResult(
+                        results=results,
+                        completed=False,
+                        error_index=i,
+                    )
+
+            # All succeeded - commit
             self.provider.query_executor.execute_non_query(conn, "COMMIT")
-
-            return result
+            return MultiStatementResult(
+                results=results,
+                completed=True,
+                error_index=None,
+            )
 
         except Exception:
-            # Rollback on any error
+            # Rollback on any error (e.g., BEGIN or COMMIT failed)
             try:
                 self.provider.query_executor.execute_non_query(conn, "ROLLBACK")
             except Exception:

sqlit/domains/query/ui/mixins/query_editing_cursor.py

@@ -45,6 +45,16 @@ def action_g_execute_single_statement(self: QueryMixinHost) -> None:
         self._clear_leader_pending()
         self.action_execute_single_statement()
 
+    def action_g_execute_query(self: QueryMixinHost) -> None:
+        """Execute query via g menu (gr)."""
+        self._clear_leader_pending()
+        self.action_execute_query()
+
+    def action_g_execute_query_atomic(self: QueryMixinHost) -> None:
+        """Execute query as transaction via g menu (gt)."""
+        self._clear_leader_pending()
+        self.action_execute_query_atomic()
+
     def action_cursor_left(self: QueryMixinHost) -> None:
         """Move cursor left (h in normal mode)."""
         row, col = self.query_input.cursor_location

sqlit/domains/query/ui/mixins/query_execution.py

@@ -439,6 +439,7 @@ async def _run_query_atomic_async(self: QueryMixinHost, query: str) -> None:
         import asyncio
         import time
 
+        from sqlit.domains.query.app.multi_statement import MultiStatementResult
         from sqlit.domains.query.app.query_service import QueryResult
         from sqlit.domains.query.app.transaction import TransactionExecutor
 
@@ -476,14 +477,21 @@ async def _run_query_atomic_async(self: QueryMixinHost, query: str) -> None:
             except Exception:
                 pass
 
-            if isinstance(result, QueryResult):
+            if isinstance(result, MultiStatementResult):
+                # Multi-statement atomic execution
+                self._display_multi_statement_results(result, elapsed_ms)
+                if result.has_error:
+                    self.notify("Transaction rolled back (error in statement)", severity="error")
+                else:
+                    self.notify("Query executed atomically (committed)", severity="information")
+            elif isinstance(result, QueryResult):
                 await self._display_query_results(
                     result.columns, result.rows, result.row_count, result.truncated, elapsed_ms
                 )
+                self.notify("Query executed atomically (committed)", severity="information")
             else:
                 self._display_non_query_result(result.rows_affected, elapsed_ms)
-
-            self.notify("Query executed atomically (committed)", severity="information")
+                self.notify("Query executed atomically (committed)", severity="information")
 
         except Exception as e:
             self._display_query_error(f"Transaction rolled back: {e}")

sqlit/shared/ui/protocols/query.py

@@ -43,7 +43,10 @@ class QueryActionsProtocol(Protocol):
     def action_execute_query(self) -> None:
         ...
 
-    def action_execute_single_statement(self) -> None: 
+    def action_execute_query_atomic(self) -> None:
+        ...
+
+    def action_execute_single_statement(self) -> None:
         ...
 
     def _get_history_store(self) -> Any: