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

snowflake/cli/_plugins/stage/commands.py

@@ -1,264 +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 itertools
-from os import path
-from pathlib import Path
-from typing import List, Optional
-
-import click
-import typer
-from snowflake.cli._plugins.object.command_aliases import (
-    add_object_command_aliases,
-    scope_option,
-)
-from snowflake.cli._plugins.stage.diff import (
-    DiffResult,
-    compute_stage_diff,
-)
-from snowflake.cli._plugins.stage.manager import StageManager
-from snowflake.cli._plugins.stage.utils import print_diff_to_console
-from snowflake.cli.api.cli_global_context import get_cli_context
-from snowflake.cli.api.commands.common import OnErrorType
-from snowflake.cli.api.commands.flags import (
-    ExecuteVariablesOption,
-    OnErrorOption,
-    PatternOption,
-    identifier_stage_argument,
-    like_option,
-)
-from snowflake.cli.api.commands.snow_typer import SnowTyperFactory
-from snowflake.cli.api.console import cli_console
-from snowflake.cli.api.constants import ObjectType
-from snowflake.cli.api.identifiers import FQN
-from snowflake.cli.api.output.formats import OutputFormat
-from snowflake.cli.api.output.types import (
-    CollectionResult,
-    CommandResult,
-    ObjectResult,
-    QueryResult,
-    SingleQueryResult,
-)
-from snowflake.cli.api.utils.path_utils import is_stage_path
-
-app = SnowTyperFactory(
-    name="stage",
-    help="Manages stages.",
-)
-
-StageNameArgument = identifier_stage_argument(sf_object="stage", example="@my_stage")
-
-add_object_command_aliases(
-    app=app,
-    object_type=ObjectType.STAGE,
-    name_argument=StageNameArgument,
-    like_option=like_option(
-        help_example='`list --like "my%"` lists all stages that begin with “my”',
-    ),
-    scope_option=scope_option(help_example="`list --in database my_db`"),
-)
-
-
-@app.command("list-files", requires_connection=True)
-def stage_list_files(
-    stage_name: str = StageNameArgument, pattern=PatternOption, **options
-) -> CommandResult:
-    """
-    Lists the stage contents.
-    """
-    cursor = StageManager().list_files(stage_name=stage_name, pattern=pattern)
-    return QueryResult(cursor)
-
-
-@app.command("copy", requires_connection=True)
-def copy(
-    source_path: str = typer.Argument(
-        help="Source path for copy operation. Can be either stage path or local. You can use a glob pattern for local files but the pattern has to be enclosed in quotes.",
-        show_default=False,
-    ),
-    destination_path: str = typer.Argument(
-        help="Target directory path for copy operation. Should be stage if source is local or local if source is stage.",
-        show_default=False,
-    ),
-    overwrite: bool = typer.Option(
-        False,
-        help="Overwrites existing files in the target path.",
-    ),
-    parallel: int = typer.Option(
-        4,
-        help="Number of parallel threads to use when uploading files.",
-    ),
-    recursive: bool = typer.Option(
-        False,
-        help="Copy files recursively with directory structure.",
-    ),
-    auto_compress: bool = typer.Option(
-        default=False,
-        help="Specifies whether Snowflake uses gzip to compress files during upload. Ignored when downloading.",
-    ),
-    **options,
-) -> CommandResult:
-    """
-    Copies all files from target path to target directory. This works for both uploading
-    to and downloading files from the stage.
-    """
-    is_get = is_stage_path(source_path)
-    is_put = is_stage_path(destination_path)
-
-    if is_get and is_put:
-        raise click.ClickException(
-            "Both source and target path are remote. This operation is not supported."
-        )
-    if not is_get and not is_put:
-        raise click.ClickException(
-            "Both source and target path are local. This operation is not supported."
-        )
-
-    if is_get:
-        return get(
-            recursive=recursive,
-            source_path=source_path,
-            destination_path=destination_path,
-            parallel=parallel,
-        )
-    return _put(
-        recursive=recursive,
-        source_path=source_path,
-        destination_path=destination_path,
-        parallel=parallel,
-        overwrite=overwrite,
-        auto_compress=auto_compress,
-    )
-
-
-@app.command("create", requires_connection=True)
-def stage_create(stage_name: FQN = StageNameArgument, **options) -> CommandResult:
-    """
-    Creates a named stage if it does not already exist.
-    """
-    cursor = StageManager().create(fqn=stage_name)
-    return SingleQueryResult(cursor)
-
-
-@app.command("remove", requires_connection=True)
-def stage_remove(
-    stage_name: FQN = StageNameArgument,
-    file_name: str = typer.Argument(
-        ...,
-        help="Name of the file to remove.",
-        show_default=False,
-    ),
-    **options,
-) -> CommandResult:
-    """
-    Removes a file from a stage.
-    """
-
-    cursor = StageManager().remove(stage_name=stage_name.identifier, path=file_name)
-    return SingleQueryResult(cursor)
-
-
-@app.command("diff", hidden=True, requires_connection=True)
-def stage_diff(
-    stage_name: str = typer.Argument(
-        help="Fully qualified name of a stage",
-        show_default=False,
-    ),
-    folder_name: str = typer.Argument(
-        help="Path to local folder",
-        show_default=False,
-    ),
-    **options,
-) -> Optional[CommandResult]:
-    """
-    Diffs a stage with a local folder.
-    """
-    diff: DiffResult = compute_stage_diff(
-        local_root=Path(folder_name), stage_fqn=stage_name
-    )
-    if get_cli_context().output_format == OutputFormat.JSON:
-        return ObjectResult(diff.to_dict())
-    else:
-        print_diff_to_console(diff)
-        return None  # don't print any output
-
-
-@app.command("execute", requires_connection=True)
-def execute(
-    stage_path: str = typer.Argument(
-        ...,
-        help="Stage path with files to be execute. For example `@stage/dev/*`.",
-        show_default=False,
-    ),
-    on_error: OnErrorType = OnErrorOption,
-    variables: Optional[List[str]] = ExecuteVariablesOption,
-    **options,
-):
-    """
-    Execute immediate all files from the stage path. Files can be filtered with glob like pattern,
-    e.g. `@stage/*.sql`, `@stage/dev/*`. Only files with `.sql` extension will be executed.
-    """
-    results = StageManager().execute(
-        stage_path=stage_path, on_error=on_error, variables=variables
-    )
-    return CollectionResult(results)
-
-
-def get(recursive: bool, source_path: str, destination_path: str, parallel: int):
-    target = Path(destination_path).resolve()
-    if not recursive:
-        cli_console.warning(
-            "Use `--recursive` flag, which copy files recursively with directory structure. This will be the default behavior in the future."
-        )
-        cursor = StageManager().get(
-            stage_path=source_path, dest_path=target, parallel=parallel
-        )
-        return QueryResult(cursor)
-
-    cursors = StageManager().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)
-
-
-def _put(
-    recursive: bool,
-    source_path: str,
-    destination_path: str,
-    parallel: int,
-    overwrite: bool,
-    auto_compress: bool,
-):
-    if recursive:
-        raise click.ClickException("Recursive flag for upload is not supported.")
-
-    source = Path(source_path).resolve()
-    local_path = str(source) + "/*" if source.is_dir() else str(source)
-
-    cursor = StageManager().put(
-        local_path=local_path,
-        stage_path=destination_path,
-        overwrite=overwrite,
-        parallel=parallel,
-        auto_compress=auto_compress,
-    )
-    return QueryResult(cursor)

snowflake/cli/_plugins/stage/diff.py

@@ -1,280 +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 logging
-from dataclasses import dataclass, field
-from pathlib import Path, PurePosixPath
-from typing import Collection, Dict, List, Optional, Tuple
-
-from snowflake.cli._plugins.nativeapp.artifacts import BundleMap
-from snowflake.cli.api.exceptions import (
-    SnowflakeSQLExecutionError,
-)
-from snowflake.connector.cursor import DictCursor
-
-from .manager import StageManager
-from .md5 import UnknownMD5FormatError, file_matches_md5sum
-
-log = logging.getLogger(__name__)
-
-StagePath = PurePosixPath  # alias PurePosixPath as StagePath for clarity
-
-
-@dataclass
-class DiffResult:
-    """
-    Each collection is a list of stage paths ('/'-separated, regardless of the platform), relative to the stage root.
-    """
-
-    identical: List[StagePath] = field(default_factory=list)
-    "Files with matching md5sums"
-
-    different: List[StagePath] = field(default_factory=list)
-    "Files that may be different between the stage and the local directory"
-
-    only_local: List[StagePath] = field(default_factory=list)
-    "Files that only exist in the local directory"
-
-    only_on_stage: List[StagePath] = field(default_factory=list)
-    "Files that only exist on the stage"
-
-    def has_changes(self) -> bool:
-        return (
-            len(self.different) > 0
-            or len(self.only_local) > 0
-            or len(self.only_on_stage) > 0
-        )
-
-    def to_dict(self) -> dict:
-        return {
-            "modified": [str(p) for p in sorted(self.different)],
-            "added": [str(p) for p in sorted(self.only_local)],
-            "deleted": [str(p) for p in sorted(self.only_on_stage)],
-        }
-
-
-def enumerate_files(path: Path) -> List[Path]:
-    """
-    Get a list of all files in a directory (recursively).
-    """
-    if not path.is_dir():
-        raise ValueError("Path must point to a directory")
-
-    paths: List[Path] = []
-    for child in sorted(path.iterdir()):
-        if child.is_dir():
-            paths += enumerate_files(child)
-        else:
-            paths.append(child)
-
-    return paths
-
-
-def strip_stage_name(path: str) -> StagePath:
-    """Returns the given stage path without the stage name as the first part."""
-    return StagePath(*path.split("/")[1:])
-
-
-def build_md5_map(list_stage_cursor: DictCursor) -> Dict[StagePath, Optional[str]]:
-    """
-    Returns a mapping of relative stage paths to their md5sums.
-    """
-    return {
-        strip_stage_name(file["name"]): file["md5"]
-        for file in list_stage_cursor.fetchall()
-    }
-
-
-def preserve_from_diff(
-    diff: DiffResult, stage_paths_to_sync: Collection[StagePath]
-) -> DiffResult:
-    """
-    Returns a filtered version of the provided diff, keeping only the provided stage paths.
-    """
-    paths_to_preserve = set(stage_paths_to_sync)
-    preserved_diff: DiffResult = DiffResult()
-    preserved_diff.identical = [i for i in diff.identical if i in paths_to_preserve]
-    preserved_diff.different = [i for i in diff.different if i in paths_to_preserve]
-    preserved_diff.only_local = [i for i in diff.only_local if i in paths_to_preserve]
-    preserved_diff.only_on_stage = [
-        i for i in diff.only_on_stage if i in paths_to_preserve
-    ]
-    return preserved_diff
-
-
-def compute_stage_diff(
-    local_root: Path,
-    stage_fqn: str,
-) -> DiffResult:
-    """
-    Diffs the files in a stage with a local folder.
-    """
-    stage_manager = StageManager()
-    local_files = enumerate_files(local_root)
-    remote_md5 = build_md5_map(stage_manager.list_files(stage_fqn))
-
-    result: DiffResult = DiffResult()
-
-    for local_file in local_files:
-        relpath = local_file.relative_to(local_root)
-        stage_path = to_stage_path(relpath)
-        if stage_path not in remote_md5:
-            # doesn't exist on the stage
-            result.only_local.append(stage_path)
-        else:
-            # N.B. file size on stage is not always accurate, so cannot fail fast
-            try:
-                if file_matches_md5sum(local_file, remote_md5[stage_path]):
-                    # We are assuming that we will not get accidental collisions here due to the
-                    # large space of the md5sum (32 * 4 = 128 bits means 1-in-9-trillion chance)
-                    # combined with the fact that the file name + path must also match elsewhere.
-                    result.identical.append(stage_path)
-                else:
-                    # either the file has changed, or we can't tell if it has
-                    result.different.append(stage_path)
-            except UnknownMD5FormatError:
-                log.warning(
-                    "Could not compare md5 for %s, assuming file has changed",
-                    local_file,
-                    exc_info=True,
-                )
-                result.different.append(stage_path)
-
-            # mark this file as seen
-            del remote_md5[stage_path]
-
-    # every entry here is a file we never saw locally
-    for stage_path in remote_md5.keys():
-        result.only_on_stage.append(stage_path)
-
-    return result
-
-
-def get_stage_subpath(stage_path: StagePath) -> str:
-    """
-    Returns the parent portion of a stage path, as a string, for inclusion in the fully qualified stage path. Note that
-    '.' treated specially here, and so the return value of this call is not a `StagePath` instance.
-    """
-    parent = str(stage_path.parent)
-    return "" if parent == "." else parent
-
-
-def to_stage_path(filename: Path) -> StagePath:
-    """
-    Returns the stage file name, with the path separator suitably transformed if needed.
-    """
-    return StagePath(*filename.parts)
-
-
-def to_local_path(stage_path: StagePath) -> Path:
-    return Path(*stage_path.parts)
-
-
-def delete_only_on_stage_files(
-    stage_manager: StageManager,
-    stage_fqn: str,
-    only_on_stage: List[StagePath],
-    role: Optional[str] = None,
-):
-    """
-    Deletes all files from a Snowflake stage according to the input list of filenames, using a custom role.
-    """
-    for _stage_path in only_on_stage:
-        stage_manager.remove(stage_name=stage_fqn, path=str(_stage_path), role=role)
-
-
-def put_files_on_stage(
-    stage_manager: StageManager,
-    stage_fqn: str,
-    deploy_root_path: Path,
-    stage_paths: List[StagePath],
-    role: Optional[str] = None,
-    overwrite: bool = False,
-):
-    """
-    Uploads all files given input list of filenames on your local filesystem, to a Snowflake stage, using a custom role.
-    """
-    for _stage_path in stage_paths:
-        stage_sub_path = get_stage_subpath(_stage_path)
-        full_stage_path = (
-            f"{stage_fqn}/{stage_sub_path}" if stage_sub_path else stage_fqn
-        )
-        stage_manager.put(
-            local_path=deploy_root_path / to_local_path(_stage_path),
-            stage_path=full_stage_path,
-            role=role,
-            overwrite=overwrite,
-        )
-
-
-def sync_local_diff_with_stage(
-    role: str, deploy_root_path: Path, diff_result: DiffResult, stage_fqn: str
-):
-    """
-    Syncs a given local directory's contents with a Snowflake stage, including removing old files, and re-uploading modified and new files.
-    """
-    stage_manager = StageManager()
-    log.info(
-        "Uploading diff-ed files from your local %s directory to the Snowflake stage.",
-        deploy_root_path,
-    )
-
-    try:
-        delete_only_on_stage_files(
-            stage_manager, stage_fqn, diff_result.only_on_stage, role
-        )
-        put_files_on_stage(
-            stage_manager,
-            stage_fqn,
-            deploy_root_path,
-            diff_result.different,
-            role,
-            overwrite=True,
-        )
-        put_files_on_stage(
-            stage_manager, stage_fqn, deploy_root_path, diff_result.only_local, role
-        )
-    except Exception as err:
-        # Could be ProgrammingError or IntegrityError from SnowflakeCursor
-        log.error(err)
-        raise SnowflakeSQLExecutionError()
-
-
-def _to_src_dest_pair(
-    stage_path: StagePath, bundle_map: Optional[BundleMap]
-) -> Tuple[Optional[str], str]:
-    if not bundle_map:
-        return None, str(stage_path)
-
-    dest_path = to_local_path(stage_path)
-    src = bundle_map.to_project_path(dest_path)
-    if src:
-        return str(src), str(stage_path)
-
-    return "?", str(stage_path)
-
-
-def _to_diff_line(status: str, src: Optional[str], dest: str) -> str:
-    if src is None:
-        src_prefix = ""
-    else:
-        src_prefix = f"{src} -> "
-
-    longest_status = "modified"
-    padding = " " * (len(longest_status) - len(status))
-    status_prefix = f"[red]{status}[/red]: {padding}"
-
-    return f"{status_prefix}{src_prefix}{dest}"