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

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/commands.py

@@ -39,6 +39,7 @@ from snowflake.cli.api.identifiers import FQN
 from snowflake.cli.api.output.types import (
     CollectionResult,
     MessageResult,
+    QueryResult,
     SingleQueryResult,
 )
 from snowflake.cli.api.project.util import is_valid_object_name
@@ -99,44 +100,10 @@ def list_images(
     **options,
 ) -> CollectionResult:
     """Lists images in the given repository."""
-    repository_manager = ImageRepositoryManager()
-    database = repository_manager.get_database()
-    schema = repository_manager.get_schema()
-    url = repository_manager.get_repository_url(name.identifier)
-    api_url = repository_manager.get_repository_api_url(url)
-    bearer_login = RegistryManager().login_to_registry(api_url)
-    repos = []
-    query: Optional[str] = f"{api_url}/_catalog?n=10"
-
-    while query:
-        # Make paginated catalog requests
-        response = requests.get(
-            query, headers={"Authorization": f"Bearer {bearer_login}"}
-        )
-
-        if response.status_code != 200:
-            raise ClickException(f"Call to the registry failed {response.text}")
-
-        data = json.loads(response.text)
-        if "repositories" in data:
-            repos.extend(data["repositories"])
-
-        if "Link" in response.headers:
-            # There are more results
-            query = f"{api_url}/_catalog?n=10&last={repos[-1]}"
-        else:
-            query = None
-
-    images = []
-    for repo in repos:
-        prefix = f"/{database}/{schema}/{name}/"
-        repo = repo.replace("baserepo/", prefix)
-        images.append({"image": repo})
-
-    return CollectionResult(images)
+    return QueryResult(ImageRepositoryManager().list_images(name.identifier))
 
 
-@app.command("list-tags", requires_connection=True)
+@app.command("list-tags", requires_connection=True, deprecated=True)
 def list_tags(
     name: FQN = REPO_NAME_ARGUMENT,
     image_name: str = typer.Option(
@@ -149,7 +116,7 @@ def list_tags(
     ),
     **options,
 ) -> CollectionResult:
-    """Lists tags for the given image in a repository."""
+    """Lists tags for the given image in a repository. This command is deprecated and will be removed in a future release. Use `list-images` instead."""
 
     repository_manager = ImageRepositoryManager()
     url = repository_manager.get_repository_url(name.identifier)

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

@@ -18,6 +18,7 @@ from snowflake.cli._plugins.spcs.common import handle_object_already_exists
 from snowflake.cli.api.constants import ObjectType
 from snowflake.cli.api.identifiers import FQN
 from snowflake.cli.api.sql_execution import SqlExecutionMixin
+from snowflake.connector.cursor import SnowflakeCursor
 from snowflake.connector.errors import ProgrammingError
 
 
@@ -32,7 +33,6 @@ class ImageRepositoryManager(SqlExecutionMixin):
         return self._conn.role
 
     def get_repository_url(self, repo_name: str, with_scheme: bool = True):
-
         repo_row = self.show_specific_object(
             "image repositories", repo_name, check_schema=True
         )
@@ -82,3 +82,6 @@ class ImageRepositoryManager(SqlExecutionMixin):
             handle_object_already_exists(
                 e, ObjectType.IMAGE_REPOSITORY, name, replace_available=True
             )
+
+    def list_images(self, repo_name: str) -> SnowflakeCursor:
+        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(
@@ -131,7 +134,7 @@ def create(
     external_access_integrations: Optional[List[str]] = typer.Option(
         None,
         "--eai-name",
-        help="Identifies External Access Integrations(EAI) that the service can access. This option may be specified multiple times for multiple EAIs.",
+        help="Identifies external access integrations (EAI) that the service can access. This option may be specified multiple times for multiple EAIs.",
     ),
     query_warehouse: Optional[str] = QueryWarehouseOption(),
     tags: Optional[List[Tag]] = TagOption(help="Tag for the service."),
@@ -174,7 +177,7 @@ def execute_job(
     external_access_integrations: Optional[List[str]] = typer.Option(
         None,
         "--eai-name",
-        help="Identifies External Access Integrations(EAI) that the job service can access. This option may be specified multiple times for multiple EAIs.",
+        help="Identifies external access integrations (EAI) that the job service can access. This option may be specified multiple times for multiple EAIs.",
     ),
     query_warehouse: Optional[str] = QueryWarehouseOption(),
     comment: Optional[str] = CommentOption(help=_COMMENT_HELP),
@@ -194,10 +197,10 @@ def execute_job(
     return SingleQueryResult(cursor)
 
 
-@app.command(requires_connection=True)
+@app.command(requires_connection=True, deprecated=True)
 def status(name: FQN = ServiceNameArgument, **options) -> CommandResult:
     """
-    Retrieves the status of a service.
+    Retrieves the status of a service. This command is deprecated and will be removed in a future release. Use `describe` instead to get service status and use `list-instances` and `list-containers` to get more detailed information about service instances and containers.
     """
     cursor = ServiceManager().status(service_name=name.identifier)
     return QueryJsonValueResult(cursor)
@@ -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)
@@ -259,6 +310,32 @@ def list_endpoints(name: FQN = ServiceNameArgument, **options):
     return QueryResult(ServiceManager().list_endpoints(service_name=name.identifier))
 
 
+@app.command("list-instances", requires_connection=True)
+def list_service_instances(name: FQN = ServiceNameArgument, **options) -> CommandResult:
+    """
+    Lists all service instances in a service.
+    """
+    return QueryResult(ServiceManager().list_instances(service_name=name.identifier))
+
+
+@app.command("list-containers", requires_connection=True)
+def list_service_containers(
+    name: FQN = ServiceNameArgument, **options
+) -> CommandResult:
+    """
+    Lists all service containers in a service.
+    """
+    return QueryResult(ServiceManager().list_containers(service_name=name.identifier))
+
+
+@app.command("list-roles", requires_connection=True)
+def list_service_roles(name: FQN = ServiceNameArgument, **options) -> CommandResult:
+    """
+    Lists all service roles in a service.
+    """
+    return QueryResult(ServiceManager().list_roles(service_name=name.identifier))
+
+
 @app.command(requires_connection=True)
 def suspend(name: FQN = ServiceNameArgument, **options) -> CommandResult:
     """
@@ -282,6 +359,11 @@ def set_property(
     max_instances: Optional[int] = MaxInstancesOption(show_default=False),
     query_warehouse: Optional[str] = QueryWarehouseOption(show_default=False),
     auto_resume: Optional[bool] = AutoResumeOption(default=None, show_default=False),
+    external_access_integrations: Optional[List[str]] = typer.Option(
+        None,
+        "--eai-name",
+        help="Identifies external access integrations (EAI) that the service can access. This option may be specified multiple times for multiple EAIs.",
+    ),
     comment: Optional[str] = CommentOption(help=_COMMENT_HELP, show_default=False),
     **options,
 ):
@@ -294,6 +376,7 @@ def set_property(
         max_instances=max_instances,
         query_warehouse=query_warehouse,
         auto_resume=auto_resume,
+        external_access_integrations=external_access_integrations,
         comment=comment,
     )
     return SingleQueryResult(cursor)

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,28 +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}")
+
+    def list_containers(self, service_name: str) -> SnowflakeCursor:
+        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}")
 
     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,
@@ -160,6 +229,7 @@ class ServiceManager(SqlExecutionMixin):
         max_instances: Optional[int],
         query_warehouse: Optional[str],
         auto_resume: Optional[bool],
+        external_access_integrations: Optional[List[str]],
         comment: Optional[str],
     ):
         property_pairs = [
@@ -167,6 +237,7 @@ class ServiceManager(SqlExecutionMixin):
             ("max_instances", max_instances),
             ("query_warehouse", query_warehouse),
             ("auto_resume", auto_resume),
+            ("external_access_integrations", external_access_integrations),
             ("comment", comment),
         ]
 
@@ -175,11 +246,32 @@ class ServiceManager(SqlExecutionMixin):
             raise NoPropertiesProvidedError(
                 f"No properties specified for service '{service_name}'. Please provide at least one property to set."
             )
-        query: List[str] = [f"alter service {service_name} set"]
-        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))
+        query: List[str] = [f"alter service {service_name} set "]
+
+        if min_instances is not None:
+            query.append(f" min_instances = {min_instances}")
+
+        if max_instances is not None:
+            query.append(f" max_instances = {max_instances}")
+
+        if query_warehouse is not None:
+            query.append(f" query_warehouse = {query_warehouse}")
+
+        if auto_resume is not None:
+            query.append(f" auto_resume = {auto_resume}")
+
+        if external_access_integrations is not None:
+            external_access_integration_list = ",".join(
+                f"{e}" for e in external_access_integrations
+            )
+            query.append(
+                f"external_access_integrations = ({external_access_integration_list})"
+            )
+
+        if comment is not None:
+            query.append(f" comment = {comment}")
+
+        return self.execute_query(strip_empty_lines(query))
 
     def unset_property(
         self,
@@ -205,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