Added the Spool Command in REPL

RELEASE-NOTES.md

@@ -19,6 +19,7 @@
 ## Deprecations
 
 ## New additions
+* Added `!spool` command to SQL REPL for writing query output to a file (`!spool <filename>` to start, `!spool off` to stop).
 
 ## Fixes and improvements
 * Fix git repository path parsing to allow quotes around both repo and branch names (e.g., `@"example-repo"/branches/"feature/branch"/*`).

src/snowflake/cli/_plugins/sql/repl.py

@@ -1,6 +1,9 @@
+import sys
 from contextlib import contextmanager
+from io import StringIO
 from logging import getLogger
-from typing import Iterable
+from pathlib import Path
+from typing import IO, Iterable, TextIO
 
 from prompt_toolkit import PromptSession
 from prompt_toolkit.filters import Condition, is_done, is_searching
@@ -45,6 +48,21 @@ def repl_context(repl_instance):
         context_manager.repl_instance = None
 
 
+class _TeeWriter:
+    """Write to both original stream and a capture buffer."""
+
+    def __init__(self, original: TextIO, capture: StringIO):
+        self.original = original
+        self.capture = capture
+
+    def write(self, text: str) -> int:
+        self.original.write(text)
+        return self.capture.write(text)
+
+    def flush(self) -> None:
+        self.original.flush()
+
+
 class Repl:
     """Basic REPL implementation for the Snowflake CLI."""
 
@@ -73,6 +91,8 @@ def __init__(
         self._sql_manager = sql_manager
         self.session = PromptSession(history=self._history)
         self._next_input: str | None = None
+        self._spool_file: IO[str] | None = None
+        self._spool_path: Path | None = None
 
     def _setup_key_bindings(self) -> KeyBindings:
         """Key bindings for repl. Helps detecting ; at end of buffer."""
@@ -220,6 +240,22 @@ def _execute(self, user_input: str) -> Iterable[SnowflakeCursor]:
         )
         return cursors
 
+    def _print_and_spool(self, result: MultipleResults, user_input: str) -> None:
+        """Print results to console and write to spool file if active."""
+        if self.is_spooling:
+            captured = StringIO()
+            tee = _TeeWriter(sys.stdout, captured)
+            sys.stdout = tee
+            try:
+                print_result(result)
+            finally:
+                sys.stdout = tee.original
+
+            self.write_to_spool(f"\n-- Query: {user_input}\n")
+            self.write_to_spool(captured.getvalue())
+        else:
+            print_result(result)
+
     def run(self):
         with repl_context(self):
             try:
@@ -228,6 +264,8 @@ def run(self):
                 self._repl_loop()
             except (KeyboardInterrupt, EOFError):
                 cli_console.message("\n[bold orange_red1]Leaving REPL, bye ...")
+            finally:
+                self.stop_spool()
 
     def _repl_loop(self):
         """Main REPL loop. Handles input and query execution.
@@ -248,7 +286,8 @@ def _repl_loop(self):
                 try:
                     log.debug("executing query")
                     cursors = self._execute(user_input)
-                    print_result(MultipleResults(QueryResult(c) for c in cursors))
+                    result = MultipleResults(QueryResult(c) for c in cursors)
+                    self._print_and_spool(result, user_input)
 
                 except Exception as e:
                     log.debug("error occurred: %s", e)
@@ -284,6 +323,42 @@ def history(self) -> FileHistory:
         """Get the FileHistory instance used by the REPL."""
         return self._history
 
+    @property
+    def spool_path(self) -> Path | None:
+        """Get the current spool file path, or None if not spooling."""
+        return self._spool_path
+
+    @property
+    def is_spooling(self) -> bool:
+        """Check if output is currently being spooled to a file."""
+        return self._spool_file is not None
+
+    def start_spool(self, path: Path) -> None:
+        """Start spooling output to the specified file.
+
+        If already spooling, closes the current file first.
+        """
+        if self._spool_file:
+            self.stop_spool()
+
+        self._spool_file = open(path, "w", encoding="utf-8")
+        self._spool_path = path
+        log.debug("Started spooling to: %s", path)
+
+    def stop_spool(self) -> None:
+        """Stop spooling output and close the spool file."""
+        if self._spool_file:
+            self._spool_file.close()
+            log.debug("Stopped spooling to: %s", self._spool_path)
+            self._spool_file = None
+            self._spool_path = None
+
+    def write_to_spool(self, text: str) -> None:
+        """Write text to the spool file if spooling is active."""
+        if self._spool_file:
+            self._spool_file.write(text)
+            self._spool_file.flush()
+
     def ask_yn(self, question: str) -> bool:
         """Asks user a Yes/No question."""
         try:

src/snowflake/cli/_plugins/sql/repl_commands.py

@@ -5,6 +5,7 @@
 import time
 from dataclasses import dataclass
 from datetime import datetime
+from pathlib import Path
 from typing import Any, Dict, Generator, Iterable, List, Tuple, Type
 from urllib.parse import urlencode
 
@@ -546,6 +547,85 @@ def _from_parsed_args(cls, args, kwargs) -> CompileCommandResult:
         return CompileCommandResult(command=cls(sql_content=sql_content))
 
 
+@register_command("!spool")
+@dataclass
+class SpoolCommand(ReplCommand):
+    """Command to spool (write) query output to a file.
+
+    Usage:
+        !spool filename.txt  - Start writing output to file
+        !spool off           - Stop writing output to file
+    """
+
+    action: str  # Either a filename to start spooling, or "off" to stop
+
+    def execute(self, connection: SnowflakeConnection):
+        """Execute the spool command.
+
+        If action is "off", stops spooling.
+        Otherwise, starts spooling to the specified file.
+        """
+        if not get_cli_context().is_repl:
+            raise CliError("The spool command can only be used in interactive mode.")
+
+        repl = get_cli_context().repl
+        if not repl:
+            raise CliError("REPL instance not found.")
+
+        if self.action.lower() == "off":
+            if repl.is_spooling:
+                spool_path = repl.spool_path
+                repl.stop_spool()
+                cli_console.message(
+                    f"[green]Spooling stopped. Output saved to: {spool_path}[/green]"
+                )
+            else:
+                cli_console.message(
+                    "[yellow]Spooling is not currently active.[/yellow]"
+                )
+        else:
+            spool_path = Path(self.action).expanduser().resolve()
+            try:
+                repl.start_spool(spool_path)
+            except OSError as e:
+                raise CliError(f"Cannot open spool file '{spool_path}': {e.strerror}")
+            cli_console.message(
+                f"[green]Spooling started. Output will be written to: {spool_path}[/green]"
+            )
+
+    @classmethod
+    def from_args(cls, raw_args, kwargs=None) -> CompileCommandResult:
+        """Parse arguments and create SpoolCommand instance.
+
+        Supports:
+        - !spool filename.txt - start spooling to file
+        - !spool off - stop spooling
+        """
+        if isinstance(raw_args, str):
+            try:
+                args, kwargs = cls._parse_args(raw_args)
+            except ValueError as e:
+                return CompileCommandResult(error_message=str(e))
+        else:
+            args, kwargs = raw_args, kwargs or {}
+
+        return cls._from_parsed_args(args, kwargs)
+
+    @classmethod
+    def _from_parsed_args(cls, args, kwargs) -> CompileCommandResult:
+        kwargs_error = _validate_kwargs_empty("spool", kwargs)
+        if kwargs_error:
+            return CompileCommandResult(error_message=kwargs_error)
+
+        if len(args) != 1:
+            amount = "Too many" if len(args) > 1 else "No"
+            return CompileCommandResult(
+                error_message=f"{amount} arguments passed to 'spool' command. Usage: `!spool <filename>` or `!spool off`"
+            )
+
+        return CompileCommandResult(command=cls(action=args[0]))
+
+
 def detect_command(input_text: str) -> tuple[str, str] | None:
     """Detect if input text matches a command pattern.
 

tests/sql/test_snowsql_commands.py

@@ -7,6 +7,7 @@
     QueriesCommand,
     ReplCommand,
     ResultCommand,
+    SpoolCommand,
     UnknownCommandError,
     compile_repl_command,
 )
@@ -286,6 +287,8 @@ def test_queries_execute_help(mock_print, mock_ctx):
         ("!abort", [_FAKE_QID], AbortCommand(_FAKE_QID)),
         ("!queries", ["amount=3", "user=jdoe"], QueriesCommand(amount=3, user="jdoe")),
         ("!QuERies", ["session"], QueriesCommand(from_current_session=True)),
+        ("!spool", ["output.txt"], SpoolCommand(action="output.txt")),
+        ("!SPOOL", ["off"], SpoolCommand(action="off")),
         (
             "!ResUlT",
             [],
@@ -296,6 +299,11 @@ def test_queries_execute_help(mock_print, mock_ctx):
             ["incorrect_id"],
             "Invalid query ID passed to 'abort' command: incorrect_id",
         ),
+        (
+            "!spool",
+            [],
+            "No arguments passed to 'spool' command. Usage: `!spool <filename>` or `!spool off`",
+        ),
         ("!unknown", [], "Unknown command '!unknown'"),
     ],
 )