Add snow logs command (#2130)

* Solution

* Fixes

* Test fix

* Test fix

* Test fix

Author: sfc-gh-jsikorski

RELEASE-NOTES.md

@@ -21,6 +21,7 @@
 ## New additions
 * Added `--prune` flag to `deploy` commands, which removes files that exist in the stage,
   but not in the local filesystem.
+* Added `snow logs` command for retrieving and streaming logs from the server.
 
 ## Fixes and improvements
 

src/snowflake/cli/_app/commands_registration/builtin_plugins.py

@@ -18,6 +18,7 @@
 from snowflake.cli._plugins.git import plugin_spec as git_plugin_spec
 from snowflake.cli._plugins.helpers import plugin_spec as migrate_plugin_spec
 from snowflake.cli._plugins.init import plugin_spec as init_plugin_spec
+from snowflake.cli._plugins.logs import plugin_spec as logs_plugin_spec
 from snowflake.cli._plugins.nativeapp import plugin_spec as nativeapp_plugin_spec
 from snowflake.cli._plugins.notebook import plugin_spec as notebook_plugin_spec
 from snowflake.cli._plugins.object import plugin_spec as object_plugin_spec
@@ -51,6 +52,7 @@ def get_builtin_plugin_name_to_plugin_spec():
         "init": init_plugin_spec,
         "workspace": workspace_plugin_spec,
         "plugin": plugin_plugin_spec,
+        "logs": logs_plugin_spec,
     }
 
     return plugin_specs

src/snowflake/cli/_plugins/logs/__init__.py

@@ -0,0 +1,83 @@
+import itertools
+from datetime import datetime
+from typing import Generator, Iterable, Optional, cast
+
+import typer
+from click import ClickException
+from snowflake.cli._plugins.logs.manager import LogsManager, LogsQueryRow
+from snowflake.cli._plugins.object.commands import NameArgument, ObjectArgument
+from snowflake.cli.api.commands.snow_typer import SnowTyperFactory
+from snowflake.cli.api.identifiers import FQN
+from snowflake.cli.api.output.types import (
+    CommandResult,
+    MessageResult,
+    StreamResult,
+)
+
+app = SnowTyperFactory()
+
+
+@app.command(name="logs", requires_connection=True)
+def get_logs(
+    object_type: str = ObjectArgument,
+    object_name: FQN = NameArgument,
+    from_: Optional[str] = typer.Option(
+        None,
+        "--from",
+        help="The start time of the logs to retrieve. Accepts all ISO8061 formats",
+    ),
+    to: Optional[str] = typer.Option(
+        None,
+        "--to",
+        help="The end time of the logs to retrieve. Accepts all ISO8061 formats",
+    ),
+    refresh_time: int = typer.Option(
+        None,
+        "--refresh",
+        help="If set, the logs will be streamed with the given refresh time in seconds",
+    ),
+    **options,
+):
+    """
+    Retrieves logs for a given object.
+    """
+    if refresh_time and to:
+        raise ClickException(
+            "You cannot set both --refresh and --to parameters. Please check the values"
+        )
+
+    from_time = get_datetime_from_string(from_, "--from") if from_ else None
+    to_time = get_datetime_from_string(to, "--to") if to else None
+
+    if refresh_time:
+        logs_stream: Iterable[LogsQueryRow] = LogsManager().stream_logs(
+            object_type=object_type,
+            object_name=object_name,
+            from_time=from_time,
+            refresh_time=refresh_time,
+        )
+        logs = itertools.chain(
+            (MessageResult(log.log_message) for logs in logs_stream for log in logs)
+        )
+    else:
+        logs_iterable: Iterable[LogsQueryRow] = LogsManager().get_logs(
+            object_type=object_type,
+            object_name=object_name,
+            from_time=from_time,
+            to_time=to_time,
+        )
+        logs = (MessageResult(log.log_message) for log in logs_iterable)  # type: ignore
+
+    return StreamResult(cast(Generator[CommandResult, None, None], logs))
+
+
+def get_datetime_from_string(
+    date_str: str,
+    name: Optional[str] = None,
+) -> datetime:
+    try:
+        return datetime.fromisoformat(date_str)
+    except ValueError:
+        raise ClickException(
+            f"Incorrect format for '{name}'. Please use one of approved ISO formats."
+        )

src/snowflake/cli/_plugins/logs/commands.py

@@ -0,0 +1,151 @@
+import functools
+import time
+from datetime import datetime
+from textwrap import dedent
+from typing import Iterable, List, NamedTuple, Optional, Tuple
+
+from click import ClickException
+from snowflake.cli._plugins.object.commands import NameArgument, ObjectArgument
+from snowflake.cli.api.identifiers import FQN
+from snowflake.cli.api.sql_execution import SqlExecutionMixin
+from snowflake.connector.cursor import SnowflakeCursor
+
+LogsTableQueryResult = NamedTuple(
+    "LogsTableQueryResult",
+    [
+        ("key", str),
+        ("table_name", str),
+        ("default", str),
+        ("level", str),
+        ("description", str),
+        ("type", str),
+    ],
+)
+
+LogsQueryRow = NamedTuple(
+    "LogsQueryRow",
+    [
+        ("timestamp", datetime),
+        ("database_name", str),
+        ("schema_name", str),
+        ("object_name", str),
+        ("log_level", str),
+        ("log_message", str),
+    ],
+)
+
+
+class LogsManager(SqlExecutionMixin):
+    def stream_logs(
+        self,
+        refresh_time: int,
+        object_type: str = ObjectArgument,
+        object_name: FQN = NameArgument,
+        from_time: Optional[datetime] = None,
+    ) -> Iterable[List[LogsQueryRow]]:
+        try:
+            previous_end = from_time
+
+            while True:
+                raw_logs = self.get_raw_logs(
+                    object_type, object_name, previous_end, None
+                ).fetchall()
+
+                if raw_logs:
+                    result = self.sanitize_logs(raw_logs)
+                    yield result
+                    if result:
+                        previous_end = result[-1].timestamp
+                time.sleep(refresh_time)
+
+        except KeyboardInterrupt:
+            return
+
+    def get_logs(
+        self,
+        object_type: str = ObjectArgument,
+        object_name: FQN = NameArgument,
+        from_time: Optional[datetime] = None,
+        to_time: Optional[datetime] = None,
+    ) -> Iterable[LogsQueryRow]:
+        """
+        Basic function to get a single batch of logs from the server
+        """
+
+        logs = self.get_raw_logs(object_type, object_name, from_time, to_time)
+
+        return self.sanitize_logs(logs)
+
+    def get_raw_logs(
+        self,
+        object_type: str = ObjectArgument,
+        object_name: FQN = NameArgument,
+        from_time: Optional[datetime] = None,
+        to_time: Optional[datetime] = None,
+    ) -> SnowflakeCursor:
+        query = dedent(
+            f"""
+            SELECT
+                timestamp,
+                resource_attributes:"snow.database.name"::string as database_name,
+                resource_attributes:"snow.schema.name"::string as schema_name,
+                resource_attributes:"snow.{object_type}.name"::string as object_name,
+                record:severity_text::string as log_level,
+                value::string as log_message
+            FROM {self.logs_table}
+            WHERE record_type = 'LOG'
+            AND (record:severity_text = 'INFO' or record:severity_text is NULL )
+            AND object_name = '{object_name}'
+            {self._get_timestamp_query(from_time, to_time)}
+            ORDER BY timestamp;
+"""
+        ).strip()
+
+        result = self.execute_query(query)
+
+        return result
+
+    @functools.cached_property
+    def logs_table(self) -> str:
+        """
+        Get the table where logs are."""
+        query_result = self.execute_query(
+            f"SHOW PARAMETERS LIKE 'event_table' IN ACCOUNT;"
+        ).fetchone()
+
+        try:
+            logs_table_query_result = LogsTableQueryResult(*query_result)
+        except TypeError:
+            raise ClickException(
+                "Encountered error while querying for logs table. Please check if your account has an event_table"
+            )
+        return logs_table_query_result.table_name
+
+    def _get_timestamp_query(
+        self, from_time: Optional[datetime], to_time: Optional[datetime]
+    ):
+        if from_time and to_time and from_time > to_time:
+            raise ClickException(
+                "From_time cannot be later than to_time. Please check the values"
+            )
+        query = []
+
+        if from_time is not None:
+            query.append(
+                f"AND timestamp >= TO_TIMESTAMP_LTZ('{from_time.isoformat()}')\n"
+            )
+
+        if to_time is not None:
+            query.append(
+                f"AND timestamp <= TO_TIMESTAMP_LTZ('{to_time.isoformat()}')\n"
+            )
+
+        return "".join(query)
+
+    def sanitize_logs(self, logs: SnowflakeCursor | List[Tuple]) -> List[LogsQueryRow]:
+        try:
+            return [LogsQueryRow(*log) for log in logs]
+        except TypeError:
+            raise ClickException(
+                "Logs table has incorrect format. Please check the logs_table in your database"
+            )

src/snowflake/cli/_plugins/logs/manager.py

@@ -0,0 +1,16 @@
+from snowflake.cli._plugins.logs import commands
+from snowflake.cli.api.plugins.command import (
+    SNOWCLI_ROOT_COMMAND_PATH,
+    CommandSpec,
+    CommandType,
+    plugin_hook_impl,
+)
+
+
+@plugin_hook_impl
+def command_spec():
+    return CommandSpec(
+        parent_command_path=SNOWCLI_ROOT_COMMAND_PATH,
+        command_type=CommandType.SINGLE_COMMAND,
+        typer_instance=commands.app.create_instance(),
+    )

src/snowflake/cli/_plugins/logs/plugin_spec.py

@@ -89,7 +89,7 @@ def execute_string(self, query: str, **kwargs) -> Iterable[SnowflakeCursor]:
         """Executes a single SQL query and returns the results"""
         return self._execute_string(query, **kwargs)
 
-    def execute_query(self, query: str, **kwargs):
+    def execute_query(self, query: str, **kwargs) -> SnowflakeCursor:
         """Executes a single SQL query and returns the last result"""
         *_, last_result = list(self.execute_string(dedent(query), **kwargs))
         return last_result