snowflake-cli 3.7.1__py3-none-any.whl → 3.8.0__py3-none-any.whl

snowflake/cli/_plugins/sql/repl.py

@@ -0,0 +1,221 @@
+from logging import getLogger
+from typing import Iterable
+
+from prompt_toolkit import PromptSession
+from prompt_toolkit.filters import Condition, is_done, is_searching
+from prompt_toolkit.history import FileHistory
+from prompt_toolkit.key_binding.key_bindings import KeyBindings
+from prompt_toolkit.keys import Keys
+from prompt_toolkit.lexers import PygmentsLexer
+from snowflake.cli._app.printing import print_result
+from snowflake.cli._plugins.sql.lexer import CliLexer, cli_completer
+from snowflake.cli._plugins.sql.manager import SqlManager
+from snowflake.cli.api.cli_global_context import get_cli_context_manager
+from snowflake.cli.api.console import cli_console
+from snowflake.cli.api.output.types import MultipleResults, QueryResult
+from snowflake.cli.api.secure_path import SecurePath
+from snowflake.connector.config_manager import CONFIG_MANAGER
+from snowflake.connector.cursor import SnowflakeCursor
+
+log = getLogger(__name__)
+
+HISTORY_FILE = SecurePath(
+    CONFIG_MANAGER.file_path.parent / "repl_history"
+).path.expanduser()
+EXIT_KEYWORDS = ("exit", "quit")
+
+log.debug("setting history file to: %s", HISTORY_FILE.as_posix())
+
+
+class Repl:
+    """Basic REPL implementation for the Snowflake CLI."""
+
+    def __init__(
+        self,
+        sql_manager: SqlManager,
+        data: dict | None = None,
+        retain_comments: bool = False,
+    ):
+        """Requires a `SqlManager` instance to execute queries.
+
+        'pass through' variables for SqlManager.execute method:
+        `data` should contain the variables used for template processing,
+        `retain_comments` how to handle comments in queries
+        """
+        super().__init__()
+        setattr(get_cli_context_manager(), "is_repl", True)
+        self._data = data or {}
+        self._retain_comments = retain_comments
+        self._history = FileHistory(HISTORY_FILE)
+        self._lexer = PygmentsLexer(CliLexer)
+        self._completer = cli_completer
+        self._repl_key_bindings = self._setup_key_bindings()
+        self._yes_no_keybindings = self._setup_yn_key_bindings()
+        self._sql_manager = sql_manager
+        self.session = PromptSession(history=self._history)
+
+    def _setup_key_bindings(self) -> KeyBindings:
+        """Key bindings for repl. Helps detecting ; at end of buffer."""
+        kb = KeyBindings()
+
+        @Condition
+        def not_searching():
+            return not is_searching()
+
+        @kb.add(Keys.Enter, filter=not_searching)
+        def _(event):
+            """Handle Enter key press."""
+            buffer = event.app.current_buffer
+            stripped_buffer = buffer.text.strip()
+
+            if stripped_buffer:
+                log.debug("evaluating repl input")
+                cursor_position = buffer.cursor_position
+                ends_with_semicolon = buffer.text.endswith(";")
+
+                if stripped_buffer.lower() in EXIT_KEYWORDS:
+                    log.debug("exit keyword detected %r", stripped_buffer)
+                    buffer.validate_and_handle()
+
+                elif ends_with_semicolon and cursor_position >= len(stripped_buffer):
+                    log.debug("semicolon detected, submitting input")
+                    buffer.validate_and_handle()
+
+                else:
+                    log.debug("adding new line")
+                    buffer.insert_text("\n")
+            else:
+                log.debug("empty input")
+
+        @kb.add(Keys.ControlJ)
+        def _(event):
+            """Control+J (and alias for Control + Enter) always inserts a new line."""
+            event.app.current_buffer.insert_text("\n")
+
+        return kb
+
+    def _setup_yn_key_bindings(self) -> KeyBindings:
+        """Key bindings for easy handling yes/no prompt."""
+        kb = KeyBindings()
+
+        @kb.add(Keys.Enter, filter=~is_done)
+        @kb.add("c-d", filter=~is_done)
+        @kb.add("y", filter=~is_done)
+        def _(event):
+            event.app.exit(result="y")
+
+        @kb.add("n", filter=~is_done)
+        @kb.add("c-c", filter=~is_done)
+        def _(event):
+            event.app.exit(result="n")
+
+        @kb.add(Keys.Any)
+        def _(event):
+            pass
+
+        return kb
+
+    def repl_propmpt(self, msg: str = " > ") -> str:
+        """Regular repl prompt."""
+        return self.session.prompt(
+            msg,
+            lexer=self._lexer,
+            completer=self._completer,
+            multiline=True,
+            wrap_lines=True,
+            key_bindings=self._repl_key_bindings,
+        )
+
+    def yn_prompt(self, msg: str) -> str:
+        """Yes/No prompt."""
+        return self.session.prompt(
+            msg,
+            lexer=None,
+            completer=None,
+            multiline=False,
+            wrap_lines=False,
+            key_bindings=self._yes_no_keybindings,
+        )
+
+    @property
+    def _welcome_banner(self) -> str:
+        return f"Welcome to Snowflake-CLI REPL\nType 'exit' or 'quit' to leave"
+
+    def _initialize_connection(self):
+        """Early connection for possible fast fail."""
+        cursor = self._execute("select current_version();")
+        res = next(iter(cursor))
+        log.debug("REPL: Snowflake version: %s", res.fetchall()[0][0])
+
+    def _execute(self, user_input: str) -> Iterable[SnowflakeCursor]:
+        """Executes a query and returns a list of cursors."""
+        _, cursors = self._sql_manager.execute(
+            query=user_input,
+            files=None,
+            std_in=False,
+            data=self._data,
+            retain_comments=self._retain_comments,
+        )
+        return cursors
+
+    def run(self):
+        try:
+            cli_console.panel(self._welcome_banner)
+            self._initialize_connection()
+            self._repl_loop()
+        except (KeyboardInterrupt, EOFError):
+            cli_console.message("\n[bold orange_red1]Leaving REPL, bye ...")
+
+    def _repl_loop(self):
+        """Main REPL loop. Handles input and query execution.
+
+        Sets up prompt session with history and key bindings.
+        Honors Ctrl-C and Ctrl-D in REPL loop.
+        """
+        while True:
+            try:
+                user_input = self.repl_propmpt().strip()
+
+                if not user_input:
+                    continue
+
+                if user_input.lower() in EXIT_KEYWORDS:
+                    raise EOFError
+
+                try:
+                    log.debug("executing query")
+                    cursors = self._execute(user_input)
+                    print_result(MultipleResults(QueryResult(c) for c in cursors))
+
+                except Exception as e:
+                    log.debug("error occurred: %s", e)
+                    cli_console.warning(f"\nError occurred: {e}")
+
+            except KeyboardInterrupt:  # a.k.a Ctrl-C
+                log.debug("user interrupted with Ctrl-C")
+                continue
+
+            except EOFError:  # a.k.a Ctrl-D
+                log.debug("user interrupted with Ctrl-D")
+                should_exit = self.ask_yn("Do you want to leave?")
+                log.debug("user answered: %r", should_exit)
+                if should_exit:
+                    raise EOFError
+                continue
+
+            except Exception as e:
+                cli_console.warning(f"\nError occurred: {e}")
+
+    def ask_yn(self, question: str) -> bool:
+        """Asks user a Yes/No question."""
+        try:
+            while True:
+                log.debug("asking user: %s", question)
+                answer = self.yn_prompt(f"{question} (y/n): ")
+                log.debug("user answered: %s", answer)
+
+                return answer == "y"
+
+        except KeyboardInterrupt:
+            log.debug("user interrupted with Ctrl-C returning to REPL")
+            return False

snowflake/cli/_plugins/sql/snowsql_commands.py

@@ -0,0 +1,331 @@
+import enum
+import re
+import time
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Any, Dict, Generator, Iterable, List, Tuple
+from urllib.parse import urlencode
+
+from snowflake.cli._app.printing import print_result
+from snowflake.cli.api.output.types import CollectionResult, QueryResult
+from snowflake.connector import SnowflakeConnection
+
+VALID_UUID_RE = re.compile(
+    r"^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$"
+)
+
+
+class CommandType(enum.Enum):
+    QUERIES = "queries"
+    UNKNOWN = "unknown"
+    URL = "url"
+
+
+class SnowSQLCommand:
+    def execute(self, connection: SnowflakeConnection):
+        """Executes command and prints the result."""
+        raise NotImplementedError
+
+    @classmethod
+    def from_args(cls, args, kwargs) -> "CompileCommandResult":
+        """Validates arguments and creates command ready for execution."""
+        raise NotImplementedError
+
+
+def _print_result_to_stdout(headers: Iterable[str], rows: Iterable[Iterable[Any]]):
+    formatted_rows: Generator[Dict[str, Any], None, None] = (
+        {key: value for key, value in zip(headers, row)} for row in rows
+    )
+    print_result(CollectionResult(formatted_rows))
+
+
+@dataclass
+class CompileCommandResult:
+    command: SnowSQLCommand | None = None
+    error_message: str | None = None
+
+
+@dataclass
+class QueriesCommand(SnowSQLCommand):
+    help_mode: bool = False
+    from_current_session: bool = False
+    amount: int = 25
+    user: str | None = None
+    warehouse: str | None = None
+    start_timestamp_ms: int | None = None
+    end_timestamp_ms: int | None = None
+    duration: str | None = None
+    stmt_type: str | None = None
+    status: str | None = None
+
+    def execute(self, connection: SnowflakeConnection) -> None:
+        if self.help_mode:
+            self._execute_help()
+        else:
+            self._execute_queries(connection)
+
+    def _execute_help(self):
+        headers = ["FILTER", "ARGUMENT", "DEFAULT"]
+        filters = [
+            ["amount", "integer", "25"],
+            ["status", "string", "any"],
+            ["warehouse", "string", "any"],
+            ["user", "string", "any"],
+            [
+                "start_date",
+                "datetime in ISO format (for example YYYY-MM-DDTHH:mm:ss.sss)",
+                "any",
+            ],
+            [
+                "end_date",
+                "datetime in ISO format (for example YYYY-MM-DDTHH:mm:ss.sss)",
+                "any",
+            ],
+            ["start", "timestamp in milliseconds (integer)", "any"],
+            ["end", "timestamp in milliseconds (integer)", "any"],
+            ["type", "string", "any"],
+            ["duration", "time in milliseconds", "any"],
+            ["session", "No arguments", "any"],
+        ]
+        _print_result_to_stdout(headers, filters)
+
+    def _execute_queries(self, connection: SnowflakeConnection) -> None:
+        url_parameters = {
+            "_dc": f"{time.time()}",
+            "includeDDL": "false",
+            "max": self.amount,
+        }
+        if self.user:
+            url_parameters["user"] = self.user
+        if self.warehouse:
+            url_parameters["wh"] = self.warehouse
+        if self.start_timestamp_ms:
+            url_parameters["start"] = self.start_timestamp_ms
+        if self.end_timestamp_ms:
+            url_parameters["end"] = self.end_timestamp_ms
+        if self.duration:
+            url_parameters["min_duration"] = self.duration
+        if self.from_current_session:
+            url_parameters["session_id"] = connection.session_id
+        if self.status:
+            url_parameters["subset"] = self.status
+        if self.stmt_type:
+            url_parameters["stmt_type"] = self.stmt_type
+
+        url = "/monitoring/queries?" + urlencode(url_parameters)
+        ret = connection.rest.request(url=url, method="get", client="rest")
+        if ret.get("data") and ret["data"].get("queries"):
+            _result: Generator[Tuple[str, str, str, str], None, None] = (
+                (
+                    query["id"],
+                    query["sqlText"],
+                    query["state"],
+                    query["totalDuration"],
+                )
+                for query in ret["data"]["queries"]
+            )
+            _print_result_to_stdout(
+                ["QUERY ID", "SQL TEXT", "STATUS", "DURATION_MS"], _result
+            )
+
+    @classmethod
+    def from_args(cls, args: List[str], kwargs: Dict[str, Any]) -> CompileCommandResult:
+        if "help" in args:
+            return CompileCommandResult(command=cls(help_mode=True))
+
+        # "session" is set by default if no other arguments are provided
+        from_current_session = "session" in args or not kwargs
+        amount = kwargs.pop("amount", "25")
+        if not amount.isdigit():
+            return CompileCommandResult(
+                error_message=f"Invalid argument passed to 'amount' filter: {amount}"
+            )
+        user = kwargs.pop("user", None)
+        warehouse = kwargs.pop("warehouse", None)
+        duration = kwargs.pop("duration", None)
+
+        start_timestamp_ms = kwargs.pop("start", None)
+        if start_timestamp_ms:
+            try:
+                start_timestamp_ms = int(start_timestamp_ms)
+            except ValueError:
+                return CompileCommandResult(
+                    error_message=f"Invalid argument passed to 'start' filter: {start_timestamp_ms}"
+                )
+        end_timestamp_ms = kwargs.pop("end", None)
+        if end_timestamp_ms:
+            try:
+                end_timestamp_ms = int(end_timestamp_ms)
+            except ValueError:
+                return CompileCommandResult(
+                    error_message=f"Invalid argument passed to 'end' filter: {end_timestamp_ms}"
+                )
+
+        start_date = kwargs.pop("start_date", None)
+        if start_date:
+            if start_timestamp_ms:
+                return CompileCommandResult(
+                    error_message=f"'start_date' filter cannot be used with 'start' filter"
+                )
+            try:
+                seconds = datetime.fromisoformat(start_date).timestamp()
+                start_timestamp_ms = int(seconds * 1000)  # convert to milliseconds
+            except ValueError:
+                return CompileCommandResult(
+                    error_message=f"Invalid date format passed to 'start_date' filter: {start_date}"
+                )
+        end_date = kwargs.pop("end_date", None)
+        if end_date:
+            if end_timestamp_ms:
+                return CompileCommandResult(
+                    error_message=f"'end_date' filter cannot be used with 'end' filter"
+                )
+            try:
+                seconds = datetime.fromisoformat(end_date).timestamp()
+                end_timestamp_ms = int(seconds * 1000)  # convert to milliseconds
+            except ValueError:
+                return CompileCommandResult(
+                    error_message=f"Invalid date format passed to 'end_date' filter: {end_date}"
+                )
+
+        stmt_type = kwargs.pop("type", None)
+        if stmt_type:
+            stmt_type = stmt_type.upper()
+            if stmt_type not in [
+                "ANY",
+                "SELECT",
+                "INSERT",
+                "UPDATE",
+                "DELETE",
+                "MERGE",
+                "MULTI_TABLE_INSERT",
+                "COPY",
+                "COMMIT",
+                "ROLLBACK",
+                "BEGIN_TRANSACTION",
+                "SHOW",
+                "GRANT",
+                "CREATE",
+                "ALTER",
+            ]:
+                return CompileCommandResult(
+                    error_message=f"Invalid argument passed to 'type' filter: {stmt_type}"
+                )
+
+        status = kwargs.pop("status", None)
+        if status:
+            status = status.upper()
+            if status not in [
+                "RUNNING",
+                "SUCCEEDED",
+                "FAILED",
+                "BLOCKED",
+                "QUEUED",
+                "ABORTED",
+            ]:
+                return CompileCommandResult(
+                    error_message=f"Invalid argument passed to 'status' filter: {status}"
+                )
+
+        for arg in args:
+            if arg.lower() not in ["session", "help"]:
+                return CompileCommandResult(
+                    error_message=f"Invalid argument passed to 'queries' command: {arg}"
+                )
+        if kwargs:
+            key, value = kwargs.popitem()
+            return CompileCommandResult(
+                error_message=f"Invalid argument passed to 'queries' command: {key}={value}"
+            )
+
+        return CompileCommandResult(
+            command=cls(
+                help_mode=False,
+                from_current_session=from_current_session,
+                amount=int(amount),
+                user=user,
+                warehouse=warehouse,
+                start_timestamp_ms=start_timestamp_ms,
+                end_timestamp_ms=end_timestamp_ms,
+                duration=duration,
+                stmt_type=stmt_type,
+                status=status,
+            )
+        )
+
+
+def _validate_only_arg_is_query_id(
+    command_name: str, args: List[str], kwargs: Dict[str, Any]
+) -> str | None:
+    if kwargs:
+        key, value = kwargs.popitem()
+        return f"Invalid argument passed to '{command_name}' command: {key}={value}"
+    if len(args) != 1:
+        amount = "Too many" if args else "No"
+        return f"{amount} arguments passed to '{command_name}' command. Usage: `!{command_name} <query id>`"
+
+    qid = args[0]
+    if not VALID_UUID_RE.match(qid):
+        return f"Invalid query ID passed to '{command_name}' command: {qid}"
+
+    return None
+
+
+@dataclass
+class ResultCommand(SnowSQLCommand):
+    query_id: str
+
+    def execute(self, connection: SnowflakeConnection):
+        cursor = connection.cursor()
+        cursor.query_result(self.query_id)
+        print_result(QueryResult(cursor=cursor))
+
+    @classmethod
+    def from_args(cls, args, kwargs) -> CompileCommandResult:
+        error_msg = _validate_only_arg_is_query_id("result", args, kwargs)
+        if error_msg:
+            return CompileCommandResult(error_message=error_msg)
+        return CompileCommandResult(command=cls(args[0]))
+
+
+@dataclass
+class AbortCommand(SnowSQLCommand):
+    query_id: str
+
+    def execute(self, connection: SnowflakeConnection):
+        cursor = connection.cursor()
+        cursor.execute(f"SELECT SYSTEM$CANCEL_QUERY('{self.query_id}')")
+        print_result(QueryResult(cursor=cursor))
+
+    @classmethod
+    def from_args(cls, args, kwargs) -> CompileCommandResult:
+        error_msg = _validate_only_arg_is_query_id("abort", args, kwargs)
+        if error_msg:
+            return CompileCommandResult(error_message=error_msg)
+        return CompileCommandResult(command=cls(args[0]))
+
+
+def compile_snowsql_command(command: str, cmd_args: List[str]):
+    """Parses command into SQL query"""
+    args = []
+    kwargs = {}
+    for cmd_arg in cmd_args:
+        if "=" not in cmd_arg:
+            args.append(cmd_arg)
+        else:
+            key, val = cmd_arg.split("=", maxsplit=1)
+            if key in kwargs:
+                return CompileCommandResult(
+                    error_message=f"duplicated argument '{key}' for command '{command}'",
+                )
+            kwargs[key] = val
+
+    match command.lower():
+        case "!queries":
+            return QueriesCommand.from_args(args, kwargs)
+        case "!result":
+            return ResultCommand.from_args(args, kwargs)
+        case "!abort":
+            return AbortCommand.from_args(args, kwargs)
+        case _:
+            return CompileCommandResult(error_message=f"Unknown command '{command}'")