snowflake-cli 3.1.0__py3-none-any.whl → 3.2.0__py3-none-any.whl

snowflake/cli/_plugins/object/commands.py

@@ -21,6 +21,8 @@ from click import ClickException
 from snowflake.cli._plugins.object.manager import ObjectManager
 from snowflake.cli.api.commands.flags import (
     IdentifierType,
+    IfNotExistsOption,
+    ReplaceOption,
     like_option,
 )
 from snowflake.cli.api.commands.snow_typer import SnowTyperFactory
@@ -145,6 +147,8 @@ def create(
     object_type: str = ObjectArgument,
     object_attributes: Optional[List[str]] = ObjectAttributesArgument,
     object_json: str = ObjectDefinitionJsonOption,
+    if_not_exists: bool = IfNotExistsOption(),
+    replace: bool = ReplaceOption(),
     **options,
 ):
     """
@@ -176,5 +180,10 @@ def create(
             "Provide either list of object attributes, or object definition in JSON format"
         )
 
-    result = ObjectManager().create(object_type=object_type, object_data=object_data)
+    result = ObjectManager().create(
+        object_type=object_type,
+        object_data=object_data,
+        if_not_exists=if_not_exists,
+        replace=replace,
+    )
     return MessageResult(result)

snowflake/cli/_plugins/object/manager.py

@@ -52,11 +52,11 @@ class ObjectManager(SqlExecutionMixin):
             query += f" like '{like}'"
         if scope[0] is not None:
             query += f" in {scope[0].replace('-', ' ')} {scope[1]}"
-        return self._execute_query(query, **kwargs)
+        return self.execute_query(query, **kwargs)
 
     def drop(self, *, object_type: str, fqn: FQN) -> SnowflakeCursor:
         object_name = _get_object_names(object_type).sf_name
-        return self._execute_query(f"drop {object_name} {fqn.sql_identifier}")
+        return self.execute_query(f"drop {object_name} {fqn.sql_identifier}")
 
     def describe(self, *, object_type: str, fqn: FQN):
         # Image repository is the only supported object that does not have a DESCRIBE command.
@@ -65,7 +65,7 @@ class ObjectManager(SqlExecutionMixin):
                 f"Describe is currently not supported for object of type image-repository"
             )
         object_name = _get_object_names(object_type).sf_name
-        return self._execute_query(f"describe {object_name} {fqn.sql_identifier}")
+        return self.execute_query(f"describe {object_name} {fqn.sql_identifier}")
 
     def object_exists(self, *, object_type: str, fqn: FQN):
         try:
@@ -74,9 +74,17 @@ class ObjectManager(SqlExecutionMixin):
         except ProgrammingError:
             return False
 
-    def create(self, object_type: str, object_data: Dict[str, Any]) -> str:
+    def create(
+        self,
+        object_type: str,
+        object_data: Dict[str, Any],
+        replace: bool = False,
+        if_not_exists: bool = False,
+    ) -> str:
         rest = RestApi(self._conn)
-        url = rest.determine_url_for_create_query(object_type=object_type)
+        url = rest.determine_url_for_create_query(
+            object_type=object_type, replace=replace, if_not_exists=if_not_exists
+        )
         try:
             response = rest.send_rest_request(url=url, method="post", data=object_data)
         except Exception as err:

snowflake/cli/_plugins/snowpark/common.py

@@ -58,9 +58,9 @@ class SnowparkObjectManager(SqlExecutionMixin):
         self, execution_identifier: str, object_type: SnowparkObject
     ) -> SnowflakeCursor:
         if object_type == SnowparkObject.FUNCTION:
-            return self._execute_query(f"select {execution_identifier}")
+            return self.execute_query(f"select {execution_identifier}")
         if object_type == SnowparkObject.PROCEDURE:
-            return self._execute_query(f"call {execution_identifier}")
+            return self.execute_query(f"call {execution_identifier}")
         raise UsageError(f"Unknown object type: {object_type}.")
 
     def create_or_replace(
@@ -95,7 +95,7 @@ class SnowparkObjectManager(SqlExecutionMixin):
         if isinstance(entity, ProcedureEntityModel) and entity.execute_as_caller:
             query.append("execute as caller")
 
-        return self._execute_query("\n".join(query))
+        return self.execute_query("\n".join(query))
 
     def deploy_entity(
         self,

snowflake/cli/_plugins/snowpark/package/anaconda_packages.py

@@ -177,7 +177,7 @@ class AnacondaPackagesManager(SqlExecutionMixin):
         return AnacondaPackages(packages)
 
     def _query_snowflake_for_available_packages(self) -> dict[str, AvailablePackage]:
-        cursor = self._execute_query(
+        cursor = self.execute_query(
             "select package_name, version from snowflake.information_schema.packages where language = 'python'",
             cursor_class=DictCursor,
         )

snowflake/cli/_plugins/spcs/common.py

@@ -95,5 +95,34 @@ def handle_object_already_exists(
         raise error
 
 
+def filter_log_timestamp(log: str, include_timestamps: bool) -> str:
+    if include_timestamps:
+        return log
+    else:
+        return log.split(" ", 1)[1] if " " in log else log
+
+
+def new_logs_only(prev_log_records: list[str], new_log_records: list[str]) -> list[str]:
+    # Sort the log records, we get time-ordered logs
+    # due to ISO 8601 timestamp format in the log content
+    # eg: 2024-10-22T01:12:29.873896187Z Count: 1
+    new_log_records_sorted = sorted(new_log_records)
+
+    # Get the first new log record to establish the overlap point
+    first_new_log_record = new_log_records_sorted[0]
+
+    # Traverse previous logs in reverse and remove duplicates from new logs
+    for prev_log in reversed(prev_log_records):
+        # Stop if the previous log is earlier than the first new log
+        if prev_log < first_new_log_record:
+            break
+
+        # Remove matching previous logs from the new logs list
+        if prev_log in new_log_records_sorted:
+            new_log_records_sorted.remove(prev_log)
+
+    return new_log_records_sorted
+
+
 class NoPropertiesProvidedError(ClickException):
     pass

snowflake/cli/_plugins/spcs/compute_pool/manager.py

@@ -56,18 +56,18 @@ class ComputePoolManager(SqlExecutionMixin):
             query.append(f"COMMENT = {comment}")
 
         try:
-            return self._execute_query(strip_empty_lines(query))
+            return self.execute_query(strip_empty_lines(query))
         except ProgrammingError as e:
             handle_object_already_exists(e, ObjectType.COMPUTE_POOL, pool_name)
 
     def stop(self, pool_name: str) -> SnowflakeCursor:
-        return self._execute_query(f"alter compute pool {pool_name} stop all")
+        return self.execute_query(f"alter compute pool {pool_name} stop all")
 
     def suspend(self, pool_name: str) -> SnowflakeCursor:
-        return self._execute_query(f"alter compute pool {pool_name} suspend")
+        return self.execute_query(f"alter compute pool {pool_name} suspend")
 
     def resume(self, pool_name: str) -> SnowflakeCursor:
-        return self._execute_query(f"alter compute pool {pool_name} resume")
+        return self.execute_query(f"alter compute pool {pool_name} resume")
 
     def set_property(
         self,
@@ -95,7 +95,7 @@ class ComputePoolManager(SqlExecutionMixin):
         for property_name, value in property_pairs:
             if value is not None:
                 query.append(f"{property_name} = {value}")
-        return self._execute_query(strip_empty_lines(query))
+        return self.execute_query(strip_empty_lines(query))
 
     def unset_property(
         self, pool_name: str, auto_resume: bool, auto_suspend_secs: bool, comment: bool
@@ -113,9 +113,7 @@ class ComputePoolManager(SqlExecutionMixin):
             )
         unset_list = [property_name for property_name, value in property_pairs if value]
         query = f"alter compute pool {pool_name} unset {','.join(unset_list)}"
-        return self._execute_query(query)
+        return self.execute_query(query)
 
     def status(self, pool_name: str):
-        return self._execute_query(
-            f"call system$get_compute_pool_status('{pool_name}')"
-        )
+        return self.execute_query(f"call system$get_compute_pool_status('{pool_name}')")

snowflake/cli/_plugins/spcs/image_registry/manager.py

@@ -36,7 +36,7 @@ class RegistryManager(SqlExecutionMixin):
         """
         Get token to authenticate with registry.
         """
-        self._execute_query(
+        self.execute_query(
             "alter session set PYTHON_CONNECTOR_QUERY_RESULT_FORMAT = 'json'"
         ).fetchall()
         # disable session deletion
@@ -75,7 +75,7 @@ class RegistryManager(SqlExecutionMixin):
 
     def get_registry_url(self) -> str:
         repositories_query = "show image repositories in account"
-        result_set = self._execute_query(repositories_query, cursor_class=DictCursor)
+        result_set = self.execute_query(repositories_query, cursor_class=DictCursor)
         results = result_set.fetchall()
         if len(results) == 0:
             raise NoImageRepositoriesFoundError()

snowflake/cli/_plugins/spcs/image_repository/manager.py

@@ -84,4 +84,4 @@ class ImageRepositoryManager(SqlExecutionMixin):
             )
 
     def list_images(self, repo_name: str) -> SnowflakeCursor:
-        return self._execute_query(f"show images in image repository {repo_name}")
+        return self.execute_query(f"show images in image repository {repo_name}")

snowflake/cli/_plugins/spcs/services/commands.py

@@ -14,9 +14,9 @@
 
 from __future__ import annotations
 
-import sys
+import itertools
 from pathlib import Path
-from typing import List, Optional
+from typing import Generator, Iterable, List, Optional, cast
 
 import typer
 from click import ClickException
@@ -26,7 +26,6 @@ from snowflake.cli._plugins.object.command_aliases import (
 )
 from snowflake.cli._plugins.object.common import CommentOption, Tag, TagOption
 from snowflake.cli._plugins.spcs.common import (
-    print_log_lines,
     validate_and_set_instances,
 )
 from snowflake.cli._plugins.spcs.services.manager import ServiceManager
@@ -38,12 +37,15 @@ from snowflake.cli.api.commands.flags import (
 )
 from snowflake.cli.api.commands.snow_typer import SnowTyperFactory
 from snowflake.cli.api.constants import ObjectType
+from snowflake.cli.api.exceptions import IncompatibleParametersError
 from snowflake.cli.api.identifiers import FQN
 from snowflake.cli.api.output.types import (
     CommandResult,
+    MessageResult,
     QueryJsonValueResult,
     QueryResult,
     SingleQueryResult,
+    StreamResult,
 )
 from snowflake.cli.api.project.util import is_valid_object_name
 
@@ -77,6 +79,7 @@ SpecPathOption = typer.Option(
     exists=True,
     show_default=False,
 )
+DEFAULT_NUM_LINES = 500
 
 _MIN_INSTANCES_HELP = "Minimum number of service instances to run."
 MinInstancesOption = OverrideableOption(
@@ -219,22 +222,70 @@ def logs(
         show_default=False,
     ),
     num_lines: int = typer.Option(
-        500, "--num-lines", help="Number of lines to retrieve."
+        DEFAULT_NUM_LINES, "--num-lines", help="Number of lines to retrieve."
+    ),
+    previous_logs: bool = typer.Option(
+        False,
+        "--previous-logs",
+        help="Retrieve logs from the last terminated container.",
+        is_flag=True,
+    ),
+    since_timestamp: Optional[str] = typer.Option(
+        "", "--since", help="Start log retrieval from a specified UTC timestamp."
+    ),
+    include_timestamps: bool = typer.Option(
+        False, "--include-timestamps", help="Include timestamps in logs.", is_flag=True
+    ),
+    follow: bool = typer.Option(
+        False, "--follow", help="Stream logs in real-time.", is_flag=True
+    ),
+    follow_interval: int = typer.Option(
+        2,
+        "--follow-interval",
+        help="Set custom polling intervals for log streaming (--follow flag) in seconds.",
     ),
     **options,
 ):
     """
     Retrieves local logs from a service container.
     """
-    results = ServiceManager().logs(
-        service_name=name.identifier,
-        instance_id=instance_id,
-        container_name=container_name,
-        num_lines=num_lines,
-    )
-    cursor = results.fetchone()
-    logs = next(iter(cursor)).split("\n")
-    print_log_lines(sys.stdout, name, "0", logs)
+    if follow:
+        if num_lines != DEFAULT_NUM_LINES:
+            raise IncompatibleParametersError(["--follow", "--num-lines"])
+        if previous_logs:
+            raise IncompatibleParametersError(["--follow", "--previous-logs"])
+
+    manager = ServiceManager()
+
+    if follow:
+        stream: Iterable[CommandResult] = (
+            MessageResult(log_batch)
+            for log_batch in manager.stream_logs(
+                service_name=name.identifier,
+                container_name=container_name,
+                instance_id=instance_id,
+                num_lines=num_lines,
+                since_timestamp=since_timestamp,
+                include_timestamps=include_timestamps,
+                interval_seconds=follow_interval,
+            )
+        )
+        stream = itertools.chain(stream, [MessageResult("")])
+    else:
+        stream = (
+            MessageResult(log)
+            for log in manager.logs(
+                service_name=name.identifier,
+                container_name=container_name,
+                instance_id=instance_id,
+                num_lines=num_lines,
+                previous_logs=previous_logs,
+                since_timestamp=since_timestamp,
+                include_timestamps=include_timestamps,
+            )
+        )
+
+    return StreamResult(cast(Generator[CommandResult, None, None], stream))
 
 
 @app.command(requires_connection=True)

snowflake/cli/_plugins/spcs/services/manager.py

@@ -15,6 +15,7 @@
 from __future__ import annotations
 
 import json
+import time
 from pathlib import Path
 from typing import List, Optional
 
@@ -22,7 +23,9 @@ import yaml
 from snowflake.cli._plugins.object.common import Tag
 from snowflake.cli._plugins.spcs.common import (
     NoPropertiesProvidedError,
+    filter_log_timestamp,
     handle_object_already_exists,
+    new_logs_only,
     strip_empty_lines,
 )
 from snowflake.cli.api.constants import DEFAULT_SIZE_LIMIT_MB, ObjectType
@@ -81,7 +84,7 @@ class ServiceManager(SqlExecutionMixin):
             query.append(f"WITH TAG ({tag_list})")
 
         try:
-            return self._execute_query(strip_empty_lines(query))
+            return self.execute_query(strip_empty_lines(query))
         except ProgrammingError as e:
             handle_object_already_exists(e, ObjectType.SERVICE, service_name)
 
@@ -119,7 +122,7 @@ class ServiceManager(SqlExecutionMixin):
             query.append(f"COMMENT = {comment}")
 
         try:
-            return self._execute_query(strip_empty_lines(query))
+            return self.execute_query(strip_empty_lines(query))
         except ProgrammingError as e:
             handle_object_already_exists(e, ObjectType.SERVICE, job_service_name)
 
@@ -130,37 +133,94 @@ class ServiceManager(SqlExecutionMixin):
         return json.dumps(data)
 
     def status(self, service_name: str) -> SnowflakeCursor:
-        return self._execute_query(f"CALL SYSTEM$GET_SERVICE_STATUS('{service_name}')")
+        return self.execute_query(f"CALL SYSTEM$GET_SERVICE_STATUS('{service_name}')")
 
     def logs(
-        self, service_name: str, instance_id: str, container_name: str, num_lines: int
+        self,
+        service_name: str,
+        instance_id: str,
+        container_name: str,
+        num_lines: int,
+        previous_logs: bool = False,
+        since_timestamp: str = "",
+        include_timestamps: bool = False,
     ):
-        return self._execute_query(
-            f"call SYSTEM$GET_SERVICE_LOGS('{service_name}', '{instance_id}', '{container_name}', {num_lines});"
+        cursor = self.execute_query(
+            f"call SYSTEM$GET_SERVICE_LOGS('{service_name}', '{instance_id}', '{container_name}', "
+            f"{num_lines}, {previous_logs}, '{since_timestamp}', {include_timestamps});"
         )
 
+        for log in cursor.fetchall():
+            yield log[0] if isinstance(log, tuple) else log
+
+    def stream_logs(
+        self,
+        service_name: str,
+        instance_id: str,
+        container_name: str,
+        num_lines: int,
+        since_timestamp: str,
+        include_timestamps: bool,
+        interval_seconds: int,
+    ):
+        try:
+            prev_timestamp = since_timestamp
+            prev_log_records: List[str] = []
+
+            while True:
+                raw_log_blocks = [
+                    log
+                    for log in self.logs(
+                        service_name=service_name,
+                        instance_id=instance_id,
+                        container_name=container_name,
+                        num_lines=num_lines,
+                        since_timestamp=prev_timestamp,
+                        include_timestamps=True,
+                    )
+                ]
+
+                new_log_records = []
+                for block in raw_log_blocks:
+                    new_log_records.extend(block.split("\n"))
+
+                new_log_records = [line for line in new_log_records if line.strip()]
+
+                if new_log_records:
+                    dedup_log_records = new_logs_only(prev_log_records, new_log_records)
+                    for log in dedup_log_records:
+                        yield filter_log_timestamp(log, include_timestamps)
+
+                    prev_timestamp = dedup_log_records[-1].split(" ", 1)[0]
+                    prev_log_records = dedup_log_records
+
+                time.sleep(interval_seconds)
+
+        except KeyboardInterrupt:
+            return
+
     def upgrade_spec(self, service_name: str, spec_path: Path):
         spec = self._read_yaml(spec_path)
         query = f"alter service {service_name} from specification $$ {spec} $$"
-        return self._execute_query(query)
+        return self.execute_query(query)
 
     def list_endpoints(self, service_name: str) -> SnowflakeCursor:
-        return self._execute_query(f"show endpoints in service {service_name}")
+        return self.execute_query(f"show endpoints in service {service_name}")
 
     def list_instances(self, service_name: str) -> SnowflakeCursor:
-        return self._execute_query(f"show service instances in service {service_name}")
+        return self.execute_query(f"show service instances in service {service_name}")
 
     def list_containers(self, service_name: str) -> SnowflakeCursor:
-        return self._execute_query(f"show service containers in service {service_name}")
+        return self.execute_query(f"show service containers in service {service_name}")
 
     def list_roles(self, service_name: str) -> SnowflakeCursor:
-        return self._execute_query(f"show roles in service {service_name}")
+        return self.execute_query(f"show roles in service {service_name}")
 
     def suspend(self, service_name: str):
-        return self._execute_query(f"alter service {service_name} suspend")
+        return self.execute_query(f"alter service {service_name} suspend")
 
     def resume(self, service_name: str):
-        return self._execute_query(f"alter service {service_name} resume")
+        return self.execute_query(f"alter service {service_name} resume")
 
     def set_property(
         self,
@@ -211,7 +271,7 @@ class ServiceManager(SqlExecutionMixin):
         if comment is not None:
             query.append(f" comment = {comment}")
 
-        return self._execute_query(strip_empty_lines(query))
+        return self.execute_query(strip_empty_lines(query))
 
     def unset_property(
         self,
@@ -237,4 +297,4 @@ class ServiceManager(SqlExecutionMixin):
             )
         unset_list = [property_name for property_name, value in property_pairs if value]
         query = f"alter service {service_name} unset {','.join(unset_list)}"
-        return self._execute_query(query)
+        return self.execute_query(query)

snowflake/cli/_plugins/sql/commands.py

@@ -17,6 +17,7 @@ from __future__ import annotations
 from pathlib import Path
 from typing import List, Optional
 
+import typer
 from snowflake.cli._plugins.sql.manager import SqlManager
 from snowflake.cli.api.commands.decorators import with_project_definition
 from snowflake.cli.api.commands.flags import (
@@ -63,6 +64,11 @@ def execute_sql(
         "String in format of key=value. If provided the SQL content will "
         "be treated as template and rendered using provided data.",
     ),
+    retain_comments: Optional[bool] = typer.Option(
+        False,
+        "--retain-comments",
+        help="Retains comments in queries passed to Snowflake",
+    ),
     **options,
 ) -> CommandResult:
     """
@@ -80,7 +86,9 @@ def execute_sql(
     if data_override:
         data = {v.key: v.value for v in parse_key_value_variables(data_override)}
 
-    single_statement, cursors = SqlManager().execute(query, files, std_in, data=data)
+    single_statement, cursors = SqlManager().execute(
+        query, files, std_in, data=data, retain_comments=retain_comments
+    )
     if single_statement:
         return QueryResult(next(cursors))
     return MultipleResults((QueryResult(c) for c in cursors))

snowflake/cli/_plugins/sql/manager.py

@@ -39,6 +39,7 @@ class SqlManager(SqlExecutionMixin):
         files: List[Path] | None,
         std_in: bool,
         data: Dict | None = None,
+        retain_comments: bool = False,
     ) -> Tuple[IsSingleStatement, Iterable[SnowflakeCursor]]:
         inputs = [query, files, std_in]
         # Check if any two inputs were provided simultaneously
@@ -50,7 +51,9 @@ class SqlManager(SqlExecutionMixin):
         if std_in:
             query = sys.stdin.read()
         if query:
-            return self._execute_single_query(query=query, data=data)
+            return self._execute_single_query(
+                query=query, data=data, retain_comments=retain_comments
+            )
 
         if files:
             # Multiple files
@@ -61,7 +64,7 @@ class SqlManager(SqlExecutionMixin):
                     file_size_limit_mb=UNLIMITED
                 )
                 single_statement, result = self._execute_single_query(
-                    query=query_from_file, data=data
+                    query=query_from_file, data=data, retain_comments=retain_comments
                 )
                 results.append(result)
 
@@ -73,7 +76,7 @@ class SqlManager(SqlExecutionMixin):
         raise UsageError("Use either query, filename or input option.")
 
     def _execute_single_query(
-        self, query: str, data: Dict | None = None
+        self, query: str, data: Dict | None = None, retain_comments: bool = False
     ) -> Tuple[IsSingleStatement, Iterable[SnowflakeCursor]]:
         try:
             query = transpile_snowsql_templates(query)
@@ -83,7 +86,9 @@ class SqlManager(SqlExecutionMixin):
 
         statements = tuple(
             statement
-            for statement, _ in split_statements(StringIO(query), remove_comments=True)
+            for statement, _ in split_statements(
+                StringIO(query), remove_comments=not retain_comments
+            )
         )
         single_statement = len(statements) == 1
 

snowflake/cli/_plugins/stage/commands.py

@@ -141,7 +141,7 @@ def copy(
         )
     return _put(
         recursive=recursive,
-        source_path=source_path,
+        source_path=Path(source_path),
         destination_path=destination_path,
         parallel=parallel,
         overwrite=overwrite,
@@ -247,23 +247,27 @@ def get(recursive: bool, source_path: str, destination_path: str, parallel: int)
 
 def _put(
     recursive: bool,
-    source_path: str,
+    source_path: Path,
     destination_path: str,
     parallel: int,
     overwrite: bool,
     auto_compress: bool,
 ):
-    if recursive:
-        raise click.ClickException("Recursive flag for upload is not supported.")
-
-    source = Path(source_path).resolve()
-    local_path = str(source) + "/*" if source.is_dir() else str(source)
-
-    cursor = StageManager().put(
-        local_path=local_path,
-        stage_path=destination_path,
-        overwrite=overwrite,
-        parallel=parallel,
-        auto_compress=auto_compress,
-    )
-    return QueryResult(cursor)
+    if recursive and not source_path.is_file():
+        cursor_generator = StageManager().put_recursive(
+            local_path=source_path,
+            stage_path=destination_path,
+            overwrite=overwrite,
+            parallel=parallel,
+            auto_compress=auto_compress,
+        )
+        return CollectionResult(cursor_generator)
+    else:
+        cursor = StageManager().put(
+            local_path=source_path.resolve(),
+            stage_path=destination_path,
+            overwrite=overwrite,
+            parallel=parallel,
+            auto_compress=auto_compress,
+        )
+        return QueryResult(cursor)

snowflake/cli/_plugins/stage/diff.py

@@ -221,7 +221,7 @@ def put_files_on_stage(
 
 
 def sync_local_diff_with_stage(
-    role: str, deploy_root_path: Path, diff_result: DiffResult, stage_fqn: str
+    role: str | None, deploy_root_path: Path, diff_result: DiffResult, stage_fqn: str
 ):
     """
     Syncs a given local directory's contents with a Snowflake stage, including removing old files, and re-uploading modified and new files.