snowflake-cli-labs 2.8.0rc1__py3-none-any.whl → 2.8.2__py3-none-any.whl

snowflake/cli/plugins/spcs/common.py

@@ -1,99 +0,0 @@
-# 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 typing import TextIO
-
-from click import ClickException
-from snowflake.cli.api.constants import ObjectType
-from snowflake.cli.api.exceptions import ObjectAlreadyExistsError
-from snowflake.cli.api.project.util import unquote_identifier
-from snowflake.connector.errors import ProgrammingError
-
-if not sys.stdout.closed and sys.stdout.isatty():
-    GREEN = "\033[32m"
-    BLUE = "\033[34m"
-    ORANGE = "\033[38:2:238:76:44m"
-    GRAY = "\033[2m"
-    ENDC = "\033[0m"
-else:
-    GREEN = ""
-    ORANGE = ""
-    BLUE = ""
-    GRAY = ""
-    ENDC = ""
-
-
-def _prefix_line(prefix: str, line: str) -> str:
-    """
-    _prefix_line ensure the prefix is still present even when dealing with return characters
-    """
-    if "\r" in line:
-        line = line.replace("\r", f"\r{prefix}")
-    if "\n" in line[:-1]:
-        line = line[:-1].replace("\n", f"\n{prefix}") + line[-1:]
-    if not line.startswith("\r"):
-        line = f"{prefix}{line}"
-    return line
-
-
-def print_log_lines(file: TextIO, name, identifier, logs):
-    prefix = f"{GREEN}{name}/{identifier}{ENDC} "
-    logs = logs[0:-1]
-    for log in logs:
-        print(_prefix_line(prefix, log + "\n"), file=file, end="", flush=True)
-
-
-def strip_empty_lines(lines: list[str]) -> str:
-    return "\n".join(stripped for l in lines if (stripped := l.strip()))
-
-
-def validate_and_set_instances(min_instances, max_instances, instance_name):
-    """
-    Used to validate that min_instances is positive and that max_instances is not less than min_instances. In the
-    case that max_instances is none, sets it equal to min_instances by default. Used like `max_instances =
-    validate_and_set_instances(min_instances, max_instances, "name")`.
-    """
-    if min_instances < 1:
-        raise ClickException(f"min_{instance_name} must be positive")
-
-    if max_instances is None:
-        max_instances = min_instances
-    elif max_instances < min_instances:
-        raise ClickException(
-            f"max_{instance_name} must be greater or equal to min_{instance_name}"
-        )
-    return max_instances
-
-
-def handle_object_already_exists(
-    error: ProgrammingError,
-    object_type: ObjectType,
-    object_name: str,
-    replace_available: bool = False,
-):
-    if error.errno == 2002:
-        raise ObjectAlreadyExistsError(
-            object_type=object_type,
-            name=unquote_identifier(object_name),
-            replace_available=replace_available,
-        )
-    else:
-        raise error
-
-
-class NoPropertiesProvidedError(ClickException):
-    pass

snowflake/cli/plugins/spcs/compute_pool/__init__.py

@@ -1,13 +0,0 @@
-# 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/spcs/compute_pool/commands.py

@@ -1,240 +0,0 @@
-# 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 typing import Optional
-
-import typer
-from click import ClickException
-from snowflake.cli.api.commands.flags import (
-    IfNotExistsOption,
-    OverrideableOption,
-    like_option,
-)
-from snowflake.cli.api.commands.snow_typer import SnowTyperFactory
-from snowflake.cli.api.constants import ObjectType
-from snowflake.cli.api.output.types import CommandResult, SingleQueryResult
-from snowflake.cli.api.project.util import is_valid_object_name
-from snowflake.cli.plugins.object.command_aliases import (
-    add_object_command_aliases,
-)
-from snowflake.cli.plugins.object.common import CommentOption
-from snowflake.cli.plugins.spcs.common import (
-    validate_and_set_instances,
-)
-from snowflake.cli.plugins.spcs.compute_pool.manager import ComputePoolManager
-
-app = SnowTyperFactory(
-    name="compute-pool",
-    help="Manages Snowpark Container Services compute pools.",
-    short_help="Manages compute pools.",
-)
-
-
-def _compute_pool_name_callback(name: str) -> str:
-    """
-    Verifies that compute pool name is a single valid identifier.
-    """
-    if not is_valid_object_name(name, max_depth=0, allow_quoted=False):
-        raise ClickException(
-            f"'{name}' is not a valid compute pool name. Note that compute pool names must be unquoted identifiers."
-        )
-    return name
-
-
-ComputePoolNameArgument = typer.Argument(
-    ...,
-    help="Name of the compute pool.",
-    callback=_compute_pool_name_callback,
-    show_default=False,
-)
-
-
-MinNodesOption = OverrideableOption(
-    1,
-    "--min-nodes",
-    help="Minimum number of nodes for the compute pool.",
-    min=1,
-)
-MaxNodesOption = OverrideableOption(
-    None,
-    "--max-nodes",
-    help="Maximum number of nodes for the compute pool.",
-    min=1,
-)
-
-_AUTO_RESUME_HELP = "The compute pool will automatically resume when a service or job is submitted to it."
-
-AutoResumeOption = OverrideableOption(
-    True,
-    "--auto-resume/--no-auto-resume",
-    help=_AUTO_RESUME_HELP,
-)
-
-_AUTO_SUSPEND_SECS_HELP = "Number of seconds of inactivity after which you want Snowflake to automatically suspend the compute pool."
-AutoSuspendSecsOption = OverrideableOption(
-    3600,
-    "--auto-suspend-secs",
-    help=_AUTO_SUSPEND_SECS_HELP,
-    min=1,
-)
-
-_COMMENT_HELP = "Comment for the compute pool."
-
-add_object_command_aliases(
-    app=app,
-    object_type=ObjectType.COMPUTE_POOL,
-    name_argument=ComputePoolNameArgument,
-    like_option=like_option(
-        help_example='`list --like "my%"` lists all compute pools that begin with “my”.'
-    ),
-    scope_option=None,
-)
-
-
-@app.command(requires_connection=True)
-def create(
-    name: str = ComputePoolNameArgument,
-    instance_family: str = typer.Option(
-        ...,
-        "--family",
-        help="Name of the instance family. For more information about instance families, refer to the SQL CREATE COMPUTE POOL command.",
-        show_default=False,
-    ),
-    min_nodes: int = MinNodesOption(),
-    max_nodes: Optional[int] = MaxNodesOption(),
-    auto_resume: bool = AutoResumeOption(),
-    initially_suspended: bool = typer.Option(
-        False,
-        "--init-suspend/--no-init-suspend",
-        help="Starts the compute pool in a suspended state.",
-    ),
-    auto_suspend_secs: int = AutoSuspendSecsOption(),
-    comment: Optional[str] = CommentOption(help=_COMMENT_HELP),
-    if_not_exists: bool = IfNotExistsOption(),
-    **options,
-) -> CommandResult:
-    """
-    Creates a new compute pool.
-    """
-    max_nodes = validate_and_set_instances(min_nodes, max_nodes, "nodes")
-    cursor = ComputePoolManager().create(
-        pool_name=name,
-        min_nodes=min_nodes,
-        max_nodes=max_nodes,
-        instance_family=instance_family,
-        auto_resume=auto_resume,
-        initially_suspended=initially_suspended,
-        auto_suspend_secs=auto_suspend_secs,
-        comment=comment,
-        if_not_exists=if_not_exists,
-    )
-    return SingleQueryResult(cursor)
-
-
-@app.command("stop-all", requires_connection=True)
-def stop_all(name: str = ComputePoolNameArgument, **options) -> CommandResult:
-    """
-    Deletes all services running on the compute pool.
-    """
-    cursor = ComputePoolManager().stop(pool_name=name)
-    return SingleQueryResult(cursor)
-
-
-@app.command(requires_connection=True)
-def suspend(name: str = ComputePoolNameArgument, **options) -> CommandResult:
-    """
-    Suspends the compute pool by suspending all currently running services and then releasing compute pool nodes.
-    """
-    return SingleQueryResult(ComputePoolManager().suspend(name))
-
-
-@app.command(requires_connection=True)
-def resume(name: str = ComputePoolNameArgument, **options) -> CommandResult:
-    """
-    Resumes the compute pool from a SUSPENDED state.
-    """
-    return SingleQueryResult(ComputePoolManager().resume(name))
-
-
-@app.command("set", requires_connection=True)
-def set_property(
-    name: str = ComputePoolNameArgument,
-    min_nodes: Optional[int] = MinNodesOption(default=None, show_default=False),
-    max_nodes: Optional[int] = MaxNodesOption(show_default=False),
-    auto_resume: Optional[bool] = AutoResumeOption(default=None, show_default=False),
-    auto_suspend_secs: Optional[int] = AutoSuspendSecsOption(
-        default=None, show_default=False
-    ),
-    comment: Optional[str] = CommentOption(
-        help="Comment for the compute pool.", show_default=False
-    ),
-    **options,
-) -> CommandResult:
-    """
-    Sets one or more properties for the compute pool.
-    """
-    cursor = ComputePoolManager().set_property(
-        pool_name=name,
-        min_nodes=min_nodes,
-        max_nodes=max_nodes,
-        auto_resume=auto_resume,
-        auto_suspend_secs=auto_suspend_secs,
-        comment=comment,
-    )
-    return SingleQueryResult(cursor)
-
-
-@app.command("unset", requires_connection=True)
-def unset_property(
-    name: str = ComputePoolNameArgument,
-    auto_resume: bool = AutoResumeOption(
-        default=False,
-        param_decls=["--auto-resume"],
-        help=f"Reset the AUTO_RESUME property - {_AUTO_RESUME_HELP}",
-        show_default=False,
-    ),
-    auto_suspend_secs: bool = AutoSuspendSecsOption(
-        default=False,
-        help=f"Reset the AUTO_SUSPEND_SECS property - {_AUTO_SUSPEND_SECS_HELP}",
-        show_default=False,
-    ),
-    comment: bool = CommentOption(
-        default=False,
-        help=f"Reset the COMMENT property - {_COMMENT_HELP}",
-        callback=None,
-        show_default=False,
-    ),
-    **options,
-) -> CommandResult:
-    """
-    Resets one or more properties for the compute pool to their default value(s).
-    """
-    cursor = ComputePoolManager().unset_property(
-        pool_name=name,
-        auto_resume=auto_resume,
-        auto_suspend_secs=auto_suspend_secs,
-        comment=comment,
-    )
-    return SingleQueryResult(cursor)
-
-
-@app.command(requires_connection=True)
-def status(pool_name: str = ComputePoolNameArgument, **options) -> CommandResult:
-    """
-    Retrieves the status of a compute pool along with a relevant message, if one exists.
-    """
-    cursor = ComputePoolManager().status(pool_name=pool_name)
-    return SingleQueryResult(cursor)

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

@@ -1,121 +0,0 @@
-# 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 typing import List, Optional
-
-from snowflake.cli.api.constants import ObjectType
-from snowflake.cli.api.sql_execution import SqlExecutionMixin
-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 ComputePoolManager(SqlExecutionMixin):
-    def create(
-        self,
-        pool_name: str,
-        min_nodes: int,
-        max_nodes: int,
-        instance_family: str,
-        auto_resume: bool,
-        initially_suspended: bool,
-        auto_suspend_secs: int,
-        comment: Optional[str],
-        if_not_exists: bool,
-    ) -> SnowflakeCursor:
-        create_statement = "CREATE COMPUTE POOL"
-        if if_not_exists:
-            create_statement = f"{create_statement} IF NOT EXISTS"
-        query = f"""\
-            {create_statement} {pool_name}
-            MIN_NODES = {min_nodes}
-            MAX_NODES = {max_nodes}
-            INSTANCE_FAMILY = {instance_family}
-            AUTO_RESUME = {auto_resume}
-            INITIALLY_SUSPENDED = {initially_suspended}
-            AUTO_SUSPEND_SECS = {auto_suspend_secs}
-            """.splitlines()
-        if comment:
-            query.append(f"COMMENT = {comment}")
-
-        try:
-            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")
-
-    def suspend(self, pool_name: str) -> SnowflakeCursor:
-        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")
-
-    def set_property(
-        self,
-        pool_name: str,
-        min_nodes: Optional[int],
-        max_nodes: Optional[int],
-        auto_resume: Optional[bool],
-        auto_suspend_secs: Optional[int],
-        comment: Optional[str],
-    ) -> SnowflakeCursor:
-        property_pairs = [
-            ("min_nodes", min_nodes),
-            ("max_nodes", max_nodes),
-            ("auto_resume", auto_resume),
-            ("auto_suspend_secs", auto_suspend_secs),
-            ("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 compute pool '{pool_name}'. Please provide at least one property to set."
-            )
-        query: List[str] = [f"alter compute pool {pool_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, pool_name: str, auto_resume: bool, auto_suspend_secs: bool, comment: bool
-    ):
-        property_pairs = [
-            ("auto_resume", auto_resume),
-            ("auto_suspend_secs", auto_suspend_secs),
-            ("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 compute pool '{pool_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 compute pool {pool_name} unset {','.join(unset_list)}"
-        return self._execute_query(query)
-
-    def status(self, pool_name: str):
-        return self._execute_query(
-            f"call system$get_compute_pool_status('{pool_name}')"
-        )

snowflake/cli/plugins/spcs/image_registry/__init__.py

@@ -1,13 +0,0 @@
-# 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/spcs/image_registry/commands.py

@@ -1,65 +0,0 @@
-# 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.commands.snow_typer import SnowTyperFactory
-from snowflake.cli.api.output.types import MessageResult, ObjectResult
-from snowflake.cli.plugins.spcs.image_registry.manager import (
-    RegistryManager,
-)
-
-app = SnowTyperFactory(
-    name="image-registry",
-    help="Manages Snowpark Container Services image registries.",
-    short_help="Manages image registries.",
-)
-
-
-_TOKEN_EPILOG = """\
-Usage Example: snow spcs image-registry token --format JSON | docker login $(snow spcs image-registry url) -u 0sessiontoken --password-stdin
-
-See the following for how to use this command to authenticate with SSO:
-https://community.snowflake.com/s/article/Authenticating-with-Snowpark-Container-Services-Image-Repository-via-SSO-and-Token-with-Snow-CLI\
-"""
-
-
-@app.command(requires_connection=True, epilog=_TOKEN_EPILOG)
-def token(**options) -> ObjectResult:
-    """
-    Retrieves a registry authentication token based on your current connection.
-
-    Note that this token is specific to your current user and will not grant access to any repositories that your current user cannot access.
-
-
-    """
-    return ObjectResult(RegistryManager().get_token())
-
-
-@app.command(requires_connection=True)
-def url(**options) -> MessageResult:
-    """
-    Gets the image registry URL for the current account.
-
-    Must be called from a role that can view at least one image repository in the image registry.
-    """
-    return MessageResult(RegistryManager().get_registry_url())
-
-
-@app.command(requires_connection=True)
-def login(**options) -> MessageResult:
-    """
-    Logs in to the account image registry with the current user's credentials through Docker.
-
-    Must be called from a role that can view at least one image repository in the image registry.
-    """
-    return MessageResult(RegistryManager().docker_registry_login().strip())

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

@@ -1,105 +0,0 @@
-# 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.
-
-import base64
-import json
-import re
-import subprocess
-from urllib.parse import urlparse
-
-import requests
-from click import ClickException
-from snowflake.cli.api.sql_execution import SqlExecutionMixin
-from snowflake.connector.cursor import DictCursor
-
-
-class NoImageRepositoriesFoundError(ClickException):
-    def __init__(self):
-        super().__init__(
-            f"No image repository found. To run this command, please switch to a role with read access to at least one image repository or create a new image repository first."
-        )
-
-
-class RegistryManager(SqlExecutionMixin):
-    def get_token(self):
-        """
-        Get token to authenticate with registry.
-        """
-        self._execute_query(
-            "alter session set PYTHON_CONNECTOR_QUERY_RESULT_FORMAT = 'json'"
-        ).fetchall()
-        # disable session deletion
-        self._conn._all_async_queries_finished = lambda: False  # noqa: SLF001
-        if self._conn._rest is None:  # noqa: SLF001
-            raise Exception("Failed to connect to Snowflake to retrieve token.")
-        # obtain and create the token
-        token_data = self._conn._rest._token_request("ISSUE")  # noqa: SLF001
-
-        return {
-            "token": token_data["data"]["sessionToken"],
-            "expires_in": token_data["data"]["validityInSecondsST"],
-        }
-
-    def login_to_registry(self, repo_url):
-        """
-        Logs in to the registry using basic authentication and generates a bearer authentication token.
-        """
-        token = json.dumps(self.get_token())
-        parsed_url = urlparse(repo_url)
-
-        scheme = parsed_url.scheme
-        host = parsed_url.netloc
-
-        login_url = f"{scheme}://{host}/login"
-        creds = base64.b64encode(f"0sessiontoken:{token}".encode("utf-8"))
-        creds = creds.decode("utf-8")
-        resp = requests.get(login_url, headers={"Authorization": f"Basic {creds}"})
-
-        if resp.status_code != 200:
-            raise ClickException(f"Failed to login to the repository {resp.text}")
-        return json.loads(resp.text)["token"]
-
-    def _has_url_scheme(self, url: str):
-        return re.fullmatch(r"^.*//.+", url) is not None
-
-    def get_registry_url(self) -> str:
-        repositories_query = "show image repositories in account"
-        result_set = self._execute_query(repositories_query, cursor_class=DictCursor)
-        results = result_set.fetchall()
-        if len(results) == 0:
-            raise NoImageRepositoriesFoundError()
-        sample_repository_url = results[0]["repository_url"]
-        if not self._has_url_scheme(sample_repository_url):
-            sample_repository_url = f"//{sample_repository_url}"
-        return urlparse(sample_repository_url).netloc
-
-    def docker_registry_login(self) -> str:
-        registry_url = self.get_registry_url()
-        token = self.get_token()
-        command = [
-            "docker",
-            "login",
-            "--username",
-            "0sessiontoken",
-            "--password-stdin",
-            registry_url,
-        ]
-        try:
-            return subprocess.check_output(
-                command, input=json.dumps(token), text=True, stderr=subprocess.PIPE
-            )
-        except subprocess.CalledProcessError as e:
-            raise ClickException(f"Login Failed: {e.stderr}".strip())
-        except FileNotFoundError:
-            raise ClickException("Docker is not installed.")

snowflake/cli/plugins/spcs/image_repository/__init__.py

@@ -1,13 +0,0 @@
-# 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.