snowflake-cli-labs 3.0.0rc4__py3-none-any.whl → 3.0.1__py3-none-any.whl

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

@@ -1,208 +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 json
-from pathlib import Path
-from typing import List, Optional
-
-import yaml
-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.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.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 execute_job(
-        self,
-        job_service_name: str,
-        compute_pool: str,
-        spec_path: Path,
-        external_access_integrations: Optional[List[str]],
-        query_warehouse: Optional[str],
-        comment: Optional[str],
-    ) -> SnowflakeCursor:
-        spec = self._read_yaml(spec_path)
-        query = f"""\
-                EXECUTE JOB SERVICE
-                IN COMPUTE POOL {compute_pool}
-                FROM SPECIFICATION $$
-                {spec}
-                $$
-                NAME = {job_service_name}
-                """.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}")
-
-        try:
-            return self._execute_query(strip_empty_lines(query))
-        except ProgrammingError as e:
-            handle_object_already_exists(e, ObjectType.SERVICE, job_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

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

@@ -1,86 +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 pathlib import Path
-from typing import List, Optional
-
-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 (
-    variables_option,
-)
-from snowflake.cli.api.commands.overrideable_parameter import OverrideableOption
-from snowflake.cli.api.commands.snow_typer import SnowTyperFactory
-from snowflake.cli.api.commands.utils import parse_key_value_variables
-from snowflake.cli.api.output.types import CommandResult, MultipleResults, QueryResult
-
-# simple Typer with defaults because it won't become a command group as it contains only one command
-app = SnowTyperFactory()
-
-SOURCE_EXCLUSIVE_OPTIONS_NAMES = ["query", "files", "std_in"]
-
-SourceOption = OverrideableOption(
-    mutually_exclusive=SOURCE_EXCLUSIVE_OPTIONS_NAMES, show_default=False
-)
-
-
-@app.command(name="sql", requires_connection=True, no_args_is_help=True)
-@with_project_definition(is_optional=True)
-def execute_sql(
-    query: Optional[str] = SourceOption(
-        default=None,
-        param_decls=["--query", "-q"],
-        help="Query to execute.",
-    ),
-    files: Optional[List[Path]] = SourceOption(
-        default=[],
-        param_decls=["--filename", "-f"],
-        exists=True,
-        file_okay=True,
-        dir_okay=False,
-        readable=True,
-        help="File to execute.",
-    ),
-    std_in: Optional[bool] = SourceOption(
-        default=False,
-        param_decls=["--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.
-    """
-
-    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

@@ -1,92 +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 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._plugins.sql.snowsql_templating import transpile_snowsql_templates
-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.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

@@ -1,30 +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._plugins.sql import commands
-from snowflake.cli.api.plugins.command import (
-    SNOWCLI_ROOT_COMMAND_PATH,
-    CommandSpec,
-    CommandType,
-    plugin_hook_impl,
-)
-
-
-@plugin_hook_impl
-def command_spec():
-    return CommandSpec(
-        parent_command_path=SNOWCLI_ROOT_COMMAND_PATH,
-        command_type=CommandType.SINGLE_COMMAND,
-        typer_instance=commands.app.create_instance(),
-    )

snowflake/cli/_plugins/sql/snowsql_templating.py

@@ -1,28 +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 string
-
-
-class _SnowSQLTemplate(string.Template):
-    delimiter = "&"
-
-
-class _Mapper:
-    def __getitem__(self, item):
-        return "&{ " + item + " }"
-
-
-def transpile_snowsql_templates(text: str) -> str:
-    return _SnowSQLTemplate(text).safe_substitute(_Mapper())  # type: ignore[arg-type]

snowflake/cli/_plugins/stage/__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.