snowflake-cli 2.8.2__py3-none-any.whl

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

@@ -0,0 +1,312 @@
+# Copyright (c) 2024 Snowflake Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+from typing import List, Optional
+
+import typer
+from click import ClickException
+from snowflake.cli.api.commands.flags import (
+    IfNotExistsOption,
+    OverrideableOption,
+    identifier_argument,
+    like_option,
+)
+from snowflake.cli.api.commands.snow_typer import SnowTyperFactory
+from snowflake.cli.api.constants import ObjectType
+from snowflake.cli.api.identifiers import FQN
+from snowflake.cli.api.output.types import (
+    CommandResult,
+    QueryJsonValueResult,
+    QueryResult,
+    SingleQueryResult,
+)
+from snowflake.cli.api.project.util import is_valid_object_name
+from snowflake.cli.plugins.object.command_aliases import (
+    add_object_command_aliases,
+    scope_option,
+)
+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
+
+app = SnowTyperFactory(
+    name="service",
+    help="Manages Snowpark Container Services services.",
+    short_help="Manages services.",
+)
+
+
+def _service_name_callback(name: FQN) -> FQN:
+    if not is_valid_object_name(name.identifier, max_depth=2, allow_quoted=False):
+        raise ClickException(
+            f"'{name}' is not a valid service name. Note service names must be unquoted identifiers. The same constraint also applies to database and schema names where you create a service."
+        )
+    return name
+
+
+ServiceNameArgument = identifier_argument(
+    sf_object="service pool",
+    example="my_service",
+    callback=_service_name_callback,
+)
+
+SpecPathOption = typer.Option(
+    ...,
+    "--spec-path",
+    help="Path to service specification file.",
+    file_okay=True,
+    dir_okay=False,
+    exists=True,
+    show_default=False,
+)
+
+_MIN_INSTANCES_HELP = "Minimum number of service instances to run."
+MinInstancesOption = OverrideableOption(
+    1, "--min-instances", help=_MIN_INSTANCES_HELP, min=1
+)
+
+_MAX_INSTANCES_HELP = "Maximum number of service instances to run."
+MaxInstancesOption = OverrideableOption(
+    None, "--max-instances", help=_MAX_INSTANCES_HELP, min=1
+)
+
+_QUERY_WAREHOUSE_HELP = "Warehouse to use if a service container connects to Snowflake to execute a query without explicitly specifying a warehouse to use."
+QueryWarehouseOption = OverrideableOption(
+    None,
+    "--query-warehouse",
+    help=_QUERY_WAREHOUSE_HELP,
+)
+
+_AUTO_RESUME_HELP = "The service will automatically resume when a service function or ingress is called."
+AutoResumeOption = OverrideableOption(
+    True,
+    "--auto-resume/--no-auto-resume",
+    help=_AUTO_RESUME_HELP,
+)
+
+_COMMENT_HELP = "Comment for the service."
+
+add_object_command_aliases(
+    app=app,
+    object_type=ObjectType.SERVICE,
+    name_argument=ServiceNameArgument,
+    like_option=like_option(
+        help_example='`list --like "my%"` lists all services that begin with “my”.'
+    ),
+    scope_option=scope_option(help_example="`list --in compute-pool my_pool`"),
+)
+
+
+@app.command(requires_connection=True)
+def create(
+    name: FQN = ServiceNameArgument,
+    compute_pool: str = typer.Option(
+        ...,
+        "--compute-pool",
+        help="Compute pool to run the service on.",
+        show_default=False,
+    ),
+    spec_path: Path = SpecPathOption,
+    min_instances: int = MinInstancesOption(),
+    max_instances: Optional[int] = MaxInstancesOption(),
+    auto_resume: bool = AutoResumeOption(),
+    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.",
+    ),
+    query_warehouse: Optional[str] = QueryWarehouseOption(),
+    tags: Optional[List[Tag]] = TagOption(help="Tag for the service."),
+    comment: Optional[str] = CommentOption(help=_COMMENT_HELP),
+    if_not_exists: bool = IfNotExistsOption(),
+    **options,
+) -> CommandResult:
+    """
+    Creates a new service in the current schema.
+    """
+    max_instances = validate_and_set_instances(
+        min_instances, max_instances, "instances"
+    )
+    cursor = ServiceManager().create(
+        service_name=name.identifier,
+        min_instances=min_instances,
+        max_instances=max_instances,
+        compute_pool=compute_pool,
+        spec_path=spec_path,
+        external_access_integrations=external_access_integrations,
+        auto_resume=auto_resume,
+        query_warehouse=query_warehouse,
+        tags=tags,
+        comment=comment,
+        if_not_exists=if_not_exists,
+    )
+    return SingleQueryResult(cursor)
+
+
+@app.command(requires_connection=True)
+def status(name: FQN = ServiceNameArgument, **options) -> CommandResult:
+    """
+    Retrieves the status of a service.
+    """
+    cursor = ServiceManager().status(service_name=name.identifier)
+    return QueryJsonValueResult(cursor)
+
+
+@app.command(requires_connection=True)
+def logs(
+    name: FQN = ServiceNameArgument,
+    container_name: str = typer.Option(
+        ...,
+        "--container-name",
+        help="Name of the container.",
+        show_default=False,
+    ),
+    instance_id: str = typer.Option(
+        ...,
+        "--instance-id",
+        help="ID of the service instance, starting with 0.",
+        show_default=False,
+    ),
+    num_lines: int = typer.Option(
+        500, "--num-lines", help="Number of lines to retrieve."
+    ),
+    **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)
+
+
+@app.command(requires_connection=True)
+def upgrade(
+    name: FQN = ServiceNameArgument,
+    spec_path: Path = SpecPathOption,
+    **options,
+):
+    """
+    Updates an existing service with a new specification file.
+    """
+    return SingleQueryResult(
+        ServiceManager().upgrade_spec(service_name=name.identifier, spec_path=spec_path)
+    )
+
+
+@app.command("list-endpoints", requires_connection=True)
+def list_endpoints(name: FQN = ServiceNameArgument, **options):
+    """
+    Lists the endpoints in a service.
+    """
+    return QueryResult(ServiceManager().list_endpoints(service_name=name.identifier))
+
+
+@app.command(requires_connection=True)
+def suspend(name: FQN = ServiceNameArgument, **options) -> CommandResult:
+    """
+    Suspends the service, shutting down and deleting all its containers.
+    """
+    return SingleQueryResult(ServiceManager().suspend(name))
+
+
+@app.command(requires_connection=True)
+def resume(name: FQN = ServiceNameArgument, **options) -> CommandResult:
+    """
+    Resumes the service from a SUSPENDED state.
+    """
+    return SingleQueryResult(ServiceManager().resume(name))
+
+
+@app.command("set", requires_connection=True)
+def set_property(
+    name: FQN = ServiceNameArgument,
+    min_instances: Optional[int] = MinInstancesOption(default=None, show_default=False),
+    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),
+    comment: Optional[str] = CommentOption(help=_COMMENT_HELP, show_default=False),
+    **options,
+):
+    """
+    Sets one or more properties for the service.
+    """
+    cursor = ServiceManager().set_property(
+        service_name=name.identifier,
+        min_instances=min_instances,
+        max_instances=max_instances,
+        query_warehouse=query_warehouse,
+        auto_resume=auto_resume,
+        comment=comment,
+    )
+    return SingleQueryResult(cursor)
+
+
+@app.command("unset", requires_connection=True)
+def unset_property(
+    name: FQN = ServiceNameArgument,
+    min_instances: bool = MinInstancesOption(
+        default=False,
+        help=f"Reset the MIN_INSTANCES property - {_MIN_INSTANCES_HELP}",
+        show_default=False,
+    ),
+    max_instances: bool = MaxInstancesOption(
+        default=False,
+        help=f"Reset the MAX_INSTANCES property - {_MAX_INSTANCES_HELP}",
+        show_default=False,
+    ),
+    query_warehouse: bool = QueryWarehouseOption(
+        default=False,
+        help=f"Reset the QUERY_WAREHOUSE property - {_QUERY_WAREHOUSE_HELP}",
+        show_default=False,
+    ),
+    auto_resume: bool = AutoResumeOption(
+        default=False,
+        param_decls=["--auto-resume"],
+        help=f"Reset the AUTO_RESUME property - {_AUTO_RESUME_HELP}",
+        show_default=False,
+    ),
+    comment: bool = CommentOption(
+        default=False,
+        help=f"Reset the COMMENT property - {_COMMENT_HELP}",
+        callback=None,
+        show_default=False,
+    ),
+    **options,
+):
+    """
+    Resets one or more properties for the service to their default value(s).
+    """
+    cursor = ServiceManager().unset_property(
+        service_name=name.identifier,
+        min_instances=min_instances,
+        max_instances=max_instances,
+        query_warehouse=query_warehouse,
+        auto_resume=auto_resume,
+        comment=comment,
+    )
+    return SingleQueryResult(cursor)

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

@@ -0,0 +1,170 @@
+# Copyright (c) 2024 Snowflake Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import List, Optional
+
+import yaml
+from snowflake.cli.api.constants import DEFAULT_SIZE_LIMIT_MB, ObjectType
+from snowflake.cli.api.secure_path import SecurePath
+from snowflake.cli.api.sql_execution import SqlExecutionMixin
+from snowflake.cli.plugins.object.common import Tag
+from snowflake.cli.plugins.spcs.common import (
+    NoPropertiesProvidedError,
+    handle_object_already_exists,
+    strip_empty_lines,
+)
+from snowflake.connector.cursor import SnowflakeCursor
+from snowflake.connector.errors import ProgrammingError
+
+
+class ServiceManager(SqlExecutionMixin):
+    def create(
+        self,
+        service_name: str,
+        compute_pool: str,
+        spec_path: Path,
+        min_instances: int,
+        max_instances: int,
+        auto_resume: bool,
+        external_access_integrations: Optional[List[str]],
+        query_warehouse: Optional[str],
+        tags: Optional[List[Tag]],
+        comment: Optional[str],
+        if_not_exists: bool,
+    ) -> SnowflakeCursor:
+        spec = self._read_yaml(spec_path)
+        create_statement = "CREATE SERVICE"
+        if if_not_exists:
+            create_statement = f"{create_statement} IF NOT EXISTS"
+        query = f"""\
+            {create_statement} {service_name}
+            IN COMPUTE POOL {compute_pool}
+            FROM SPECIFICATION $$
+            {spec}
+            $$
+            MIN_INSTANCES = {min_instances}
+            MAX_INSTANCES = {max_instances}
+            AUTO_RESUME = {auto_resume}
+            """.splitlines()
+
+        if external_access_integrations:
+            external_access_integration_list = ",".join(
+                f"{e}" for e in external_access_integrations
+            )
+            query.append(
+                f"EXTERNAL_ACCESS_INTEGRATIONS = ({external_access_integration_list})"
+            )
+
+        if query_warehouse:
+            query.append(f"QUERY_WAREHOUSE = {query_warehouse}")
+
+        if comment:
+            query.append(f"COMMENT = {comment}")
+
+        if tags:
+            tag_list = ",".join(f"{t.name}={t.value_string_literal()}" for t in tags)
+            query.append(f"WITH TAG ({tag_list})")
+
+        try:
+            return self._execute_query(strip_empty_lines(query))
+        except ProgrammingError as e:
+            handle_object_already_exists(e, ObjectType.SERVICE, service_name)
+
+    def _read_yaml(self, path: Path) -> str:
+        # TODO(aivanou): Add validation towards schema
+        with SecurePath(path).open("r", read_file_limit_mb=DEFAULT_SIZE_LIMIT_MB) as fh:
+            data = yaml.safe_load(fh)
+        return json.dumps(data)
+
+    def status(self, service_name: str) -> SnowflakeCursor:
+        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
+    ):
+        return self._execute_query(
+            f"call SYSTEM$GET_SERVICE_LOGS('{service_name}', '{instance_id}', '{container_name}', {num_lines});"
+        )
+
+    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)
+
+    def list_endpoints(self, service_name: str) -> SnowflakeCursor:
+        return self._execute_query(f"show endpoints in service {service_name}")
+
+    def suspend(self, service_name: str):
+        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")
+
+    def set_property(
+        self,
+        service_name: str,
+        min_instances: Optional[int],
+        max_instances: Optional[int],
+        query_warehouse: Optional[str],
+        auto_resume: Optional[bool],
+        comment: Optional[str],
+    ):
+        property_pairs = [
+            ("min_instances", min_instances),
+            ("max_instances", max_instances),
+            ("query_warehouse", query_warehouse),
+            ("auto_resume", auto_resume),
+            ("comment", comment),
+        ]
+
+        # Check if all provided properties are set to None (no properties are being set)
+        if all([value is None for property_name, value in property_pairs]):
+            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))
+
+    def unset_property(
+        self,
+        service_name: str,
+        min_instances: bool,
+        max_instances: bool,
+        query_warehouse: bool,
+        auto_resume: bool,
+        comment: bool,
+    ):
+        property_pairs = [
+            ("min_instances", min_instances),
+            ("max_instances", max_instances),
+            ("query_warehouse", query_warehouse),
+            ("auto_resume", auto_resume),
+            ("comment", comment),
+        ]
+
+        # Check if all properties provided are False (no properties are being unset)
+        if not any([value for property_name, value in property_pairs]):
+            raise NoPropertiesProvidedError(
+                f"No properties specified for service '{service_name}'. Please provide at least one property to reset to its default value."
+            )
+        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)

snowflake/cli/plugins/sql/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2024 Snowflake Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

snowflake/cli/plugins/sql/commands.py

@@ -0,0 +1,83 @@
+# Copyright (c) 2024 Snowflake Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import List, Optional
+
+import typer
+from snowflake.cli.api.commands.decorators import with_project_definition
+from snowflake.cli.api.commands.flags import (
+    parse_key_value_variables,
+    variables_option,
+)
+from snowflake.cli.api.commands.snow_typer import SnowTyperFactory
+from snowflake.cli.api.output.types import CommandResult, MultipleResults, QueryResult
+from snowflake.cli.plugins.sql.manager import SqlManager
+
+# simple Typer with defaults because it won't become a command group as it contains only one command
+app = SnowTyperFactory()
+
+
+@app.command(name="sql", requires_connection=True, no_args_is_help=True)
+@with_project_definition(is_optional=True)
+def execute_sql(
+    query: Optional[str] = typer.Option(
+        None,
+        "--query",
+        "-q",
+        help="Query to execute.",
+    ),
+    files: Optional[List[Path]] = typer.Option(
+        None,
+        "--filename",
+        "-f",
+        exists=True,
+        file_okay=True,
+        dir_okay=False,
+        readable=True,
+        help="File to execute.",
+    ),
+    std_in: Optional[bool] = typer.Option(
+        False,
+        "--stdin",
+        "-i",
+        help="Read the query from standard input. Use it when piping input to this command.",
+    ),
+    data_override: List[str] = variables_option(
+        "String in format of key=value. If provided the SQL content will "
+        "be treated as template and rendered using provided data.",
+    ),
+    **options,
+) -> CommandResult:
+    """
+    Executes Snowflake query.
+
+    Use either query, filename or input option.
+
+    Query to execute can be specified using query option, filename option (all queries from file will be executed)
+    or via stdin by piping output from other command. For example `cat my.sql | snow sql -i`.
+
+    The command supports variable substitution that happens on client-side. Both &VARIABLE or &{ VARIABLE }
+    syntax are supported.
+    """
+    data = {}
+    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)
+    if single_statement:
+        return QueryResult(next(cursors))
+    return MultipleResults((QueryResult(c) for c in cursors))

snowflake/cli/plugins/sql/manager.py

@@ -0,0 +1,92 @@
+# Copyright (c) 2024 Snowflake Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import sys
+from io import StringIO
+from itertools import chain
+from pathlib import Path
+from typing import Dict, Iterable, List, Tuple
+
+from click import ClickException, UsageError
+from jinja2 import UndefinedError
+from snowflake.cli.api.rendering.sql_templates import snowflake_sql_jinja_render
+from snowflake.cli.api.secure_path import UNLIMITED, SecurePath
+from snowflake.cli.api.sql_execution import SqlExecutionMixin, VerboseCursor
+from snowflake.cli.plugins.sql.snowsql_templating import transpile_snowsql_templates
+from snowflake.connector.cursor import SnowflakeCursor
+from snowflake.connector.util_text import split_statements
+
+IsSingleStatement = bool
+
+
+class SqlManager(SqlExecutionMixin):
+    def execute(
+        self,
+        query: str | None,
+        files: List[Path] | None,
+        std_in: bool,
+        data: Dict | None = None,
+    ) -> Tuple[IsSingleStatement, Iterable[SnowflakeCursor]]:
+        inputs = [query, files, std_in]
+        # Check if any two inputs were provided simultaneously
+        if len([i for i in inputs if i]) > 1:
+            raise UsageError(
+                "Multiple input sources specified. Please specify only one."
+            )
+
+        if std_in:
+            query = sys.stdin.read()
+        if query:
+            return self._execute_single_query(query=query, data=data)
+
+        if files:
+            # Multiple files
+            results = []
+            single_statement = False
+            for file in files:
+                query_from_file = SecurePath(file).read_text(
+                    file_size_limit_mb=UNLIMITED
+                )
+                single_statement, result = self._execute_single_query(
+                    query=query_from_file, data=data
+                )
+                results.append(result)
+
+            # Use single_statement if there's only one, otherwise this is multi statement result
+            single_statement = len(files) == 1 and single_statement
+            return single_statement, chain.from_iterable(results)
+
+        # At that point, no stdin, query or files were provided
+        raise UsageError("Use either query, filename or input option.")
+
+    def _execute_single_query(
+        self, query: str, data: Dict | None = None
+    ) -> Tuple[IsSingleStatement, Iterable[SnowflakeCursor]]:
+        try:
+            query = transpile_snowsql_templates(query)
+            query = snowflake_sql_jinja_render(content=query, data=data)
+        except UndefinedError as err:
+            raise ClickException(f"SQL template rendering error: {err}")
+
+        statements = tuple(
+            statement
+            for statement, _ in split_statements(StringIO(query), remove_comments=True)
+        )
+        single_statement = len(statements) == 1
+
+        return single_statement, self._execute_string(
+            "\n".join(statements), cursor_class=VerboseCursor
+        )

snowflake/cli/plugins/sql/plugin_spec.py

@@ -0,0 +1,30 @@
+# Copyright (c) 2024 Snowflake Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from snowflake.cli.api.plugins.command import (
+    SNOWCLI_ROOT_COMMAND_PATH,
+    CommandSpec,
+    CommandType,
+    plugin_hook_impl,
+)
+from snowflake.cli.plugins.sql import commands
+
+
+@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(),
+    )