snowflake-cli-labs 2.6.0rc0__py3-none-any.whl → 2.7.0rc0__py3-none-any.whl

snowflake/cli/plugins/connection/util.py

@@ -12,14 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
+import json
 import logging
 
 from click.exceptions import ClickException
 from snowflake.connector import SnowflakeConnection
 from snowflake.connector.cursor import DictCursor
 
-LOCAL_DEPLOYMENT: str = "us-west-2"
-
 log = logging.getLogger(__name__)
 
 REGIONLESS_QUERY = """
@@ -29,11 +30,23 @@ REGIONLESS_QUERY = """
     )) where value['name'] = 'UI_SNOWSIGHT_ENABLE_REGIONLESS_REDIRECT';
 """
 
+ALLOWLIST_QUERY = "SELECT SYSTEM$ALLOWLIST()"
+SNOWFLAKE_DEPLOYMENT = "SNOWFLAKE_DEPLOYMENT"
+LOCAL_DEPLOYMENT_REGION: str = "us-west-2"
 
-class MissingConnectionHostError(ClickException):
+
+class MissingConnectionAccountError(ClickException):
     def __init__(self, conn: SnowflakeConnection):
         super().__init__(
-            f"The connection host ({conn.host}) was missing or not in "
+            "Could not determine account by system call, configured account name, or configured host. Connection: "
+            + repr(conn)
+        )
+
+
+class MissingConnectionRegionError(ClickException):
+    def __init__(self, host: str | None):
+        super().__init__(
+            f"The connection host ({host}) was missing or not in "
             "the expected format "
             "(<account>.<deployment>.snowflakecomputing.com)"
         )
@@ -50,11 +63,60 @@ def is_regionless_redirect(conn: SnowflakeConnection) -> bool:
         *_, cursor = conn.execute_string(REGIONLESS_QUERY, cursor_class=DictCursor)
         return cursor.fetchone()["REGIONLESS"].lower() == "true"
     except:
-        # by default, assume that
-        log.warning("Cannot determine regionless redirect; assuming True.")
+        log.warning(
+            "Cannot determine regionless redirect; assuming True.", exc_info=True
+        )
         return True
 
 
+def get_host_region(host: str) -> str | None:
+    """
+    Looks for hosts of form
+    <account>.[x.y.z].snowflakecomputing.com
+    Returns the three-part [region identifier] or None.
+    """
+    host_parts = host.split(".")
+    if host_parts[-1] == "local":
+        return LOCAL_DEPLOYMENT_REGION
+    elif len(host_parts) == 6:
+        return ".".join(host_parts[1:4])
+    return None
+
+
+def guess_regioned_host_from_allowlist(conn: SnowflakeConnection) -> str | None:
+    """
+    Use SYSTEM$ALLOWLIST to find a regioned host (<account>.x.y.z.snowflakecomputing.com)
+    that corresponds to the given Snowflake connection object.
+    """
+    try:
+        *_, cursor = conn.execute_string(ALLOWLIST_QUERY, cursor_class=DictCursor)
+        allowlist_tuples = json.loads(cursor.fetchone()["SYSTEM$ALLOWLIST()"])
+        for t in allowlist_tuples:
+            if t["type"] == SNOWFLAKE_DEPLOYMENT:
+                if get_host_region(t["host"]) is not None:
+                    return t["host"]
+    except:
+        log.warning(
+            "Could not call SYSTEM$ALLOWLIST; returning an empty guess.", exc_info=True
+        )
+    return None
+
+
+def get_region(conn: SnowflakeConnection) -> str:
+    """
+    Get the region of the given connection, or raise MissingConnectionRegionError.
+    """
+    if conn.host:
+        if region := get_host_region(conn.host):
+            return region
+
+    if host := guess_regioned_host_from_allowlist(conn):
+        if region := get_host_region(host):
+            return region
+
+    raise MissingConnectionRegionError(host or conn.host)
+
+
 def get_context(conn: SnowflakeConnection) -> str:
     """
     Determines the first part of the path in a Snowsight URL.
@@ -67,14 +129,7 @@ def get_context(conn: SnowflakeConnection) -> str:
         )
         return cursor.fetchone()["SYSTEM$RETURN_CURRENT_ORG_NAME()"]
 
-    host_parts = conn.host.split(".")
-    if host_parts[-1] == "local":
-        return LOCAL_DEPLOYMENT
-
-    if len(host_parts) == 6:
-        return ".".join(host_parts[1:4])
-
-    raise MissingConnectionHostError(conn)
+    return get_region(conn)
 
 
 def get_account(conn: SnowflakeConnection) -> str:
@@ -91,11 +146,11 @@ def get_account(conn: SnowflakeConnection) -> str:
         if conn.account:
             return conn.account
 
-        if not conn.host:
-            raise MissingConnectionHostError(conn)
+        if conn.host:
+            host_parts = conn.host.split(".")
+            return host_parts[0]
 
-        host_parts = conn.host.split(".")
-        return host_parts[0]
+        raise MissingConnectionAccountError(conn)
 
 
 def get_snowsight_host(conn: SnowflakeConnection) -> str:

snowflake/cli/plugins/cortex/commands.py

@@ -24,6 +24,7 @@ from click import UsageError
 from snowflake.cli.api.cli_global_context import cli_context
 from snowflake.cli.api.commands.flags import readable_file_option
 from snowflake.cli.api.commands.snow_typer import SnowTyperFactory
+from snowflake.cli.api.constants import PYTHON_3_12
 from snowflake.cli.api.output.types import (
     CollectionResult,
     CommandResult,
@@ -45,7 +46,7 @@ app = SnowTyperFactory(
     help="Provides access to Snowflake Cortex.",
 )
 
-SEARCH_COMMAND_ENABLED = sys.version_info < (3, 12)
+SEARCH_COMMAND_ENABLED = sys.version_info < PYTHON_3_12
 
 
 @app.command(

snowflake/cli/plugins/git/commands.py

@@ -14,15 +14,18 @@
 
 from __future__ import annotations
 
+import itertools
 import logging
+from os import path
+from pathlib import Path
 from typing import List, Optional
 
 import typer
 from click import ClickException
 from snowflake.cli.api.commands.flags import (
+    ExecuteVariablesOption,
     OnErrorOption,
     PatternOption,
-    VariablesOption,
     identifier_argument,
     like_option,
 )
@@ -37,7 +40,6 @@ from snowflake.cli.plugins.object.command_aliases import (
     scope_option,
 )
 from snowflake.cli.plugins.object.manager import ObjectManager
-from snowflake.cli.plugins.stage.commands import get
 from snowflake.cli.plugins.stage.manager import OnErrorType
 
 app = SnowTyperFactory(
@@ -264,7 +266,6 @@ def copy(
             )
         )
     return get(
-        recursive=True,
         source_path=repository_path,
         destination_path=destination_path,
         parallel=parallel,
@@ -275,7 +276,7 @@ def copy(
 def execute(
     repository_path: str = RepoPathArgument,
     on_error: OnErrorType = OnErrorOption,
-    variables: Optional[List[str]] = VariablesOption,
+    variables: Optional[List[str]] = ExecuteVariablesOption,
     **options,
 ):
     """
@@ -287,3 +288,18 @@ def execute(
         stage_path=repository_path, on_error=on_error, variables=variables
     )
     return CollectionResult(results)
+
+
+def get(source_path: str, destination_path: str, parallel: int):
+    target = Path(destination_path).resolve()
+
+    cursors = GitManager().get_recursive(
+        stage_path=source_path, dest_path=target, parallel=parallel
+    )
+    results = [list(QueryResult(c).result) for c in cursors]
+    flattened_results = list(itertools.chain.from_iterable(results))
+    sorted_results = sorted(
+        flattened_results,
+        key=lambda e: (path.dirname(e["file"]), path.basename(e["file"])),
+    )
+    return CollectionResult(sorted_results)

snowflake/cli/plugins/git/manager.py

@@ -12,13 +12,50 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 from pathlib import Path
 from textwrap import dedent
+from typing import List
 
-from snowflake.cli.plugins.stage.manager import StageManager, StagePathParts
+from snowflake.cli.plugins.stage.manager import (
+    USER_STAGE_PREFIX,
+    StageManager,
+    StagePathParts,
+    UserStagePathParts,
+)
 from snowflake.connector.cursor import SnowflakeCursor
 
 
+class GitStagePathParts(StagePathParts):
+    def __init__(self, stage_path: str):
+        self.stage = GitManager.get_stage_from_path(stage_path)
+        stage_path_parts = Path(stage_path).parts
+        git_repo_name = stage_path_parts[0].split(".")[-1]
+        if git_repo_name.startswith("@"):
+            git_repo_name = git_repo_name[1:]
+        self.stage_name = "/".join([git_repo_name, *stage_path_parts[1:3], ""])
+        self.directory = "/".join(stage_path_parts[3:])
+        self.is_directory = True if stage_path.endswith("/") else False
+
+    @property
+    def path(self) -> str:
+        return (
+            f"{self.stage_name}{self.directory}"
+            if self.stage_name.endswith("/")
+            else f"{self.stage_name}/{self.directory}"
+        )
+
+    def add_stage_prefix(self, file_path: str) -> str:
+        stage = Path(self.stage).parts[0]
+        file_path_without_prefix = Path(file_path).parts[1:]
+        return f"{stage}/{'/'.join(file_path_without_prefix)}"
+
+    def get_directory_from_file_path(self, file_path: str) -> List[str]:
+        stage_path_length = len(Path(self.directory).parts)
+        return list(Path(file_path).parts[3 + stage_path_length : -1])
+
+
 class GitManager(StageManager):
     def show_branches(self, repo_name: str, like: str) -> SnowflakeCursor:
         return self._execute_query(f"show git branches like '{like}' in {repo_name}")
@@ -51,22 +88,9 @@ class GitManager(StageManager):
         """
         return f"{'/'.join(Path(path).parts[0:3])}/"
 
-    def _split_stage_path(self, stage_path: str) -> StagePathParts:
-        """
-        Splits Git repository path `@repo/branch/main/dir`
-            stage -> @repo/branch/main/
-            stage_name -> repo/branch/main/
-            directory -> dir
-        For Git repository with fully qualified name `@db.schema.repo/branch/main/dir`
-            stage -> @db.schema.repo/branch/main/
-            stage_name -> repo/branch/main/
-            directory -> dir
-        """
-        stage = self.get_stage_from_path(stage_path)
-        stage_path_parts = Path(stage_path).parts
-        git_repo_name = stage_path_parts[0].split(".")[-1]
-        if git_repo_name.startswith("@"):
-            git_repo_name = git_repo_name[1:]
-        stage_name = "/".join([git_repo_name, *stage_path_parts[1:3], ""])
-        directory = "/".join(stage_path_parts[3:])
-        return StagePathParts(stage, stage_name, directory)
+    @staticmethod
+    def _stage_path_part_factory(stage_path: str) -> StagePathParts:
+        stage_path = StageManager.get_standard_stage_prefix(stage_path)
+        if stage_path.startswith(USER_STAGE_PREFIX):
+            return UserStagePathParts(stage_path)
+        return GitStagePathParts(stage_path)

snowflake/cli/plugins/init/__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/init/commands.py

@@ -0,0 +1,242 @@
+# 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 logging
+from typing import Any, Dict, List, Optional
+
+import typer
+import yaml
+from click import ClickException
+from snowflake.cli.api.commands.flags import (
+    NoInteractiveOption,
+    parse_key_value_variables,
+    variables_option,
+)
+from snowflake.cli.api.commands.snow_typer import SnowTyperFactory
+from snowflake.cli.api.constants import DEFAULT_SIZE_LIMIT_MB
+from snowflake.cli.api.exceptions import InvalidTemplate
+from snowflake.cli.api.output.types import (
+    CommandResult,
+    MessageResult,
+)
+from snowflake.cli.api.project.schemas.template import Template, TemplateVariable
+from snowflake.cli.api.rendering.project_templates import render_template_files
+from snowflake.cli.api.secure_path import SecurePath
+
+# simple Typer with defaults because it won't become a command group as it contains only one command
+app = SnowTyperFactory()
+
+
+DEFAULT_SOURCE = "https://github.com/snowflakedb/snowflake-cli-templates"
+
+log = logging.getLogger(__name__)
+
+
+def _path_argument_callback(path: str) -> str:
+    if SecurePath(path).exists():
+        raise ClickException(
+            f"The directory {path} already exists. Please specify a different path for the project."
+        )
+    return path
+
+
+PathArgument = typer.Argument(
+    ...,
+    help="Directory to be initialized with the project. This directory must not already exist",
+    show_default=False,
+    callback=_path_argument_callback,
+)
+TemplateOption = typer.Option(
+    None,
+    "--template",
+    help="which template (subdirectory of --template-source) should be used. If not provided,"
+    " whole source will be used as the template.",
+    show_default=False,
+)
+SourceOption = typer.Option(
+    default=DEFAULT_SOURCE,
+    help=f"local path to template directory or URL to git repository with templates.",
+)
+VariablesOption = variables_option(
+    "String in `key=value` format. Provided variables will not be prompted for."
+)
+
+TEMPLATE_METADATA_FILE_NAME = "template.yml"
+
+
+def _fetch_local_template(
+    template_source: SecurePath, path: Optional[str], destination: SecurePath
+) -> SecurePath:
+    """Copies local template to [dest] and returns path to the template root.
+    Ends with an error of the template does not exist."""
+
+    template_source.assert_exists()
+    template_origin = template_source / path if path else template_source
+    log.info("Copying local template from %s", template_origin.path)
+    if not template_origin.exists():
+        raise ClickException(
+            f"Template '{path}' cannot be found under {template_source}"
+        )
+
+    template_origin.copy(destination.path)
+    return destination / template_origin.name
+
+
+def _fetch_remote_template(
+    url: str, path: Optional[str], destination: SecurePath
+) -> SecurePath:
+    """Downloads remote repository template to [dest],
+    and returns path to the template root.
+    Ends with an error of the template does not exist."""
+    from git import GitCommandError
+    from git import rmtree as git_rmtree
+
+    # TODO: during nativeapp refactor get rid of this dependency
+    from snowflake.cli.plugins.nativeapp.utils import shallow_git_clone
+
+    log.info("Downloading remote template from %s", url)
+    try:
+        shallow_git_clone(url, to_path=destination.path)
+    except GitCommandError as err:
+        import re
+
+        if re.search("fatal: repository '.*' not found", err.stderr):
+            raise ClickException(f"Repository '{url}' does not exist")
+        raise
+
+    if path:
+        # template is a subdirectoruy of the repository
+        template_root = destination / path
+    else:
+        # template is a whole repository
+        # removing .git directory not to copy it to the template
+        template_root = destination
+        git_rmtree((template_root / ".git").path)
+    if not template_root.exists():
+        raise ClickException(f"Template '{path}' cannot be found under {url}")
+
+    return template_root
+
+
+def _read_template_metadata(template_root: SecurePath) -> Template:
+    """Parse template.yml file."""
+    template_metadata_path = template_root / TEMPLATE_METADATA_FILE_NAME
+    log.debug("Reading template metadata from %s", template_metadata_path.path)
+    if not template_metadata_path.exists():
+        raise InvalidTemplate(
+            f"Template does not have {TEMPLATE_METADATA_FILE_NAME} file."
+        )
+    with template_metadata_path.open(read_file_limit_mb=DEFAULT_SIZE_LIMIT_MB) as fd:
+        yaml_contents = yaml.safe_load(fd) or {}
+    return Template(template_root, **yaml_contents)
+
+
+def _remove_template_metadata_file(template_root: SecurePath) -> None:
+    (template_root / TEMPLATE_METADATA_FILE_NAME).unlink()
+
+
+def _determine_variable_values(
+    variables_metadata: List[TemplateVariable],
+    variables_from_flags: Dict[str, Any],
+    no_interactive: bool,
+) -> Dict[str, Any]:
+    """
+    Prompt user for values not provided in [variables_from_flags].
+    If [no_interactive] is True, fill not provided variables with their default values.
+    """
+    result = {}
+
+    log.debug(
+        "Resolving values of variables: %s",
+        ", ".join(v.name for v in variables_metadata),
+    )
+    for variable in variables_metadata:
+        if variable.name in variables_from_flags:
+            value = variable.python_type(variables_from_flags[variable.name])
+        else:
+            value = variable.prompt_user_for_value(no_interactive)
+
+        result[variable.name] = value
+
+    return result
+
+
+def _validate_cli_version(required_version: str) -> None:
+    from packaging.version import parse
+    from snowflake.cli.__about__ import VERSION
+
+    if parse(required_version) > parse(VERSION):
+        raise ClickException(
+            f"Snowflake CLI version ({VERSION}) is too low - minimum version required"
+            f" by template is {required_version}. Please upgrade before continuing."
+        )
+
+
+@app.command(no_args_is_help=True)
+def init(
+    path: str = PathArgument,
+    template: Optional[str] = TemplateOption,
+    template_source: Optional[str] = SourceOption,
+    variables: Optional[List[str]] = VariablesOption,
+    no_interactive: bool = NoInteractiveOption,
+    **options,
+) -> CommandResult:
+    """
+    Creates project directory from template.
+    """
+    variables_from_flags = {
+        v.key: v.value for v in parse_key_value_variables(variables)
+    }
+    is_remote = any(
+        template_source.startswith(prefix) for prefix in ["git@", "http://", "https://"]  # type: ignore
+    )
+
+    # copy/download template into tmpdir, so it is going to be removed in case command ends with an error
+    with SecurePath.temporary_directory() as tmpdir:
+        if is_remote:
+            template_root = _fetch_remote_template(
+                url=template_source, path=template, destination=tmpdir  # type: ignore
+            )
+        else:
+            template_root = _fetch_local_template(
+                template_source=SecurePath(template_source),
+                path=template,
+                destination=tmpdir,
+            )
+
+        template_metadata = _read_template_metadata(template_root)
+        if template_metadata.minimum_cli_version:
+            _validate_cli_version(template_metadata.minimum_cli_version)
+
+        variable_values = _determine_variable_values(
+            variables_metadata=template_metadata.variables,
+            variables_from_flags=variables_from_flags,
+            no_interactive=no_interactive,
+        )
+        variable_values["project_dir_name"] = SecurePath(path).name
+        log.debug(
+            "Rendering template files: %s", ", ".join(template_metadata.files_to_render)
+        )
+        render_template_files(
+            template_root=template_root,
+            files_to_render=template_metadata.files_to_render,
+            data=variable_values,
+        )
+        _remove_template_metadata_file(template_root)
+        SecurePath(path).parent.mkdir(exist_ok=True, parents=True)
+        template_root.copy(path)
+
+    return MessageResult(f"Initialized the new project in {path}")

snowflake/cli/plugins/init/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.init 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(),
+    )

snowflake/cli/plugins/nativeapp/codegen/artifact_processor.py

@@ -15,6 +15,7 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
+from pathlib import Path
 from typing import Optional
 
 from click import ClickException
@@ -34,6 +35,42 @@ class UnsupportedArtifactProcessorError(ClickException):
         )
 
 
+def is_python_file_artifact(src: Path, _: Path):
+    """Determines whether the provided source path is an existing python file."""
+    return src.is_file() and src.suffix == ".py"
+
+
+class ProjectFileContextManager:
+    """
+    A context manager that encapsulates the logic required to update a project file
+    in processor logic. The processor can use this manager to gain access to the contents
+    of a file, and optionally provide replacement contents. If it does, the file is
+    correctly modified in the deploy root directory to reflect the new contents.
+    """
+
+    def __init__(self, path: Path):
+        self.path = path
+        self._contents = None
+        self.edited_contents = None
+
+    @property
+    def contents(self):
+        return self._contents
+
+    def __enter__(self):
+        self._contents = self.path.read_text(encoding="utf-8")
+
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if self.edited_contents is not None:
+            if self.path.is_symlink():
+                # if the file is a symlink, make sure we don't overwrite the original
+                self.path.unlink()
+
+            self.path.write_text(self.edited_contents, encoding="utf-8")
+
+
 class ArtifactProcessor(ABC):
     def __init__(
         self,
@@ -49,3 +86,6 @@ class ArtifactProcessor(ABC):
         **kwargs,
     ) -> None:
         pass
+
+    def edit_file(self, path: Path):
+        return ProjectFileContextManager(path)