snowflake-cli 3.5.0__py3-none-any.whl → 3.7.0__py3-none-any.whl

snowflake/cli/_plugins/sql/source_reader.py

@@ -0,0 +1,230 @@
+import enum
+import io
+import re
+import urllib.error
+from typing import Any, Callable, Generator, Literal, Sequence
+from urllib.request import urlopen
+
+from jinja2 import UndefinedError
+from snowflake.cli.api.secure_path import UNLIMITED, SecurePath
+from snowflake.connector.util_text import split_statements
+
+SOURCE_PATTERN = re.compile(
+    r"^!(source|load)\s+[\"']?(.*?)[\"']?\s*(?:;|$)",
+    flags=re.IGNORECASE,
+)
+
+URL_PATTERN = re.compile(r"^(\w+?):\/(\/.*)", flags=re.IGNORECASE)
+
+SplitedStatements = Generator[
+    tuple[str, bool | None] | tuple[str, Literal[False]],
+    Any,
+    None,
+]
+
+SqlTransformFunc = Callable[[str], str]
+OperatorFunctions = Sequence[SqlTransformFunc]
+
+
+class SourceType(enum.Enum):
+    FILE = "file"
+    QUERY = "query"
+    UNKNOWN = "unknown"
+    URL = "url"
+
+
+class ParsedSource:
+    """Container for parsed statement.
+
+    Holds:
+      - source: statement on command content
+      - source_type: type of source
+      - source_path: in case of URL or FILE path of the origin
+      - error: optional message
+    """
+
+    __slots__ = ("source", "source_type", "source_path", "error")
+    __match_args__ = ("source_type", "error")
+
+    source: io.StringIO
+    source_type: SourceType | None
+    source_path: str | None
+    error: str | None
+
+    def __init__(
+        self,
+        source: str,
+        source_type: SourceType,
+        source_path: str | None,
+        error: str | None = None,
+    ):
+        self.source = io.StringIO(source)
+        self.source_type = source_type
+        self.source_path = source_path
+        self.error = error
+
+    def __bool__(self):
+        return not self.error
+
+    def __eq__(self, other):
+        result = (
+            self.source_type == other.source_type,
+            self.source_path == other.source_path,
+            self.error == other.error,
+            self.source.read() == other.source.read(),
+        )
+        self.source.seek(0)
+        other.source.seek(0)
+        return all(result)
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}(source_type={self.source_type}, source_path={self.source_path}, error={self.error})"
+
+    @classmethod
+    def from_url(cls, path_part: str, raw_source: str) -> "ParsedSource":
+        """Constructor for loading from URL."""
+        try:
+            payload = urlopen(path_part, timeout=10.0).read().decode()
+            return cls(payload, SourceType.URL, path_part)
+
+        except urllib.error.HTTPError as err:
+            error = f"Could not fetch {path_part}: {err}"
+            return cls(path_part, SourceType.URL, raw_source, error)
+
+    @classmethod
+    def from_file(cls, path_part: str, raw_source: str) -> "ParsedSource":
+        """Constructor for loading from file."""
+        path = SecurePath(path_part)
+
+        if path.is_file():
+            payload = path.read_text(file_size_limit_mb=UNLIMITED)
+            return cls(payload, SourceType.FILE, path.as_posix())
+
+        error_msg = f"Could not read: {path_part}"
+        return cls(path_part, SourceType.FILE, raw_source, error_msg)
+
+
+RecursiveStatementReader = Generator[ParsedSource, Any, Any]
+
+
+def parse_source(source: str, operators: OperatorFunctions) -> ParsedSource:
+    """Evaluates templating and source commands.
+
+    Returns parsed source according to origin."""
+    try:
+        statement = source
+        for operator in operators:
+            statement = operator(statement)
+    except UndefinedError as e:
+        error_msg = f"SQL template rendering error: {e}"
+        return ParsedSource(source, SourceType.UNKNOWN, source, error_msg)
+
+    split_result = SOURCE_PATTERN.split(statement, maxsplit=1)
+    split_result = [p.strip() for p in split_result]
+
+    if len(split_result) == 1:
+        return ParsedSource(statement, SourceType.QUERY, None)
+
+    _, command, source_path, *_ = split_result
+    _path_match = URL_PATTERN.split(source_path.lower())
+
+    match command.lower(), _path_match:
+        # load content from an URL
+        case "source" | "load", ("", "http" | "https", *_):
+            return ParsedSource.from_url(source_path, statement)
+
+        # load content from a local file
+        case "source" | "load", (str(),):
+            return ParsedSource.from_file(source_path, statement)
+
+        case _:
+            error_msg = f"Unknown source: {source_path}"
+
+    return ParsedSource(source_path, SourceType.UNKNOWN, source, error_msg)
+
+
+def recursive_source_reader(
+    source: SplitedStatements,
+    seen_files: list,
+    operators: OperatorFunctions,
+    remove_comments: bool,
+) -> RecursiveStatementReader:
+    """Based on detected source command reads content of the source and tracks for recursion cycles."""
+    for stmt, _ in source:
+        if not stmt:
+            continue
+        parsed_source = parse_source(stmt, operators)
+
+        match parsed_source:
+            case ParsedSource(SourceType.FILE | SourceType.URL, None):
+                if parsed_source.source_path in seen_files:
+                    error = f"Recursion detected: {' -> '.join(seen_files)}"
+                    parsed_source.error = error
+                    yield parsed_source
+                    continue
+
+                seen_files.append(parsed_source.source_path)
+
+                yield from recursive_source_reader(
+                    split_statements(parsed_source.source, remove_comments),
+                    seen_files,
+                    operators,
+                    remove_comments,
+                )
+
+                seen_files.pop()
+
+            case ParsedSource(SourceType.URL, error) if error:
+                yield parsed_source
+
+            case _:
+                yield parsed_source
+    return
+
+
+def files_reader(
+    paths: Sequence[SecurePath],
+    operators: OperatorFunctions,
+    remove_comments: bool = False,
+) -> RecursiveStatementReader:
+    """Entry point for reading statements from files.
+
+    Returns a generator with statements."""
+    for path in paths:
+        with path.open(read_file_limit_mb=UNLIMITED) as f:
+            stmts = split_statements(io.StringIO(f.read()), remove_comments)
+            yield from recursive_source_reader(
+                stmts,
+                [path.as_posix()],
+                operators,
+                remove_comments,
+            )
+
+
+def query_reader(
+    source: str,
+    operators: OperatorFunctions,
+    remove_comments: bool = False,
+) -> RecursiveStatementReader:
+    """Entry point for reading statements from query.
+
+    Returns a generator with statements."""
+    stmts = split_statements(io.StringIO(source), remove_comments)
+    yield from recursive_source_reader(stmts, [], operators, remove_comments)
+
+
+def compile_statements(source: RecursiveStatementReader):
+    """Tracks statements evaluation and collects errors."""
+    errors = []
+    cnt = 0
+    compiled = []
+
+    for stmt in source:
+        if stmt.source_type == SourceType.QUERY:
+            cnt += 1
+            if not stmt.error:
+                compiled.append(stmt.source.read())
+        if stmt.error:
+            errors.append(stmt.error)
+
+    return errors, cnt, compiled

snowflake/cli/_plugins/stage/manager.py

@@ -65,7 +65,7 @@ EXECUTE_SUPPORTED_FILES_FORMATS = (
 
 # Replace magic numbers with constants
 OMIT_FIRST = slice(1, None)
-STAGE_PATH_REGEX = rf"(?P<prefix>@)?(?:(?P<first_qualifier>{VALID_IDENTIFIER_REGEX})\.)?(?:(?P<second_qualifier>{VALID_IDENTIFIER_REGEX})\.)?(?P<name>{VALID_IDENTIFIER_REGEX})/?(?P<directory>([^/]*/?)*)?"
+STAGE_PATH_REGEX = rf"(?P<prefix>(@|{re.escape('snow://')}))?(?:(?P<first_qualifier>{VALID_IDENTIFIER_REGEX})\.)?(?:(?P<second_qualifier>{VALID_IDENTIFIER_REGEX})\.)?(?P<name>{VALID_IDENTIFIER_REGEX})/?(?P<directory>([^/]*/?)*)?"
 
 
 @dataclass
@@ -119,6 +119,14 @@ class StagePathParts:
         raise NotImplementedError
 
 
+def _strip_standard_stage_prefix(path: str) -> str:
+    """Removes '@' or 'snow://' prefix from given string"""
+    for prefix in ["@", "snow://"]:
+        if path.startswith(prefix):
+            path = path.removeprefix(prefix)
+    return path
+
+
 @dataclass
 class DefaultStagePathParts(StagePathParts):
     """
@@ -126,8 +134,8 @@ class DefaultStagePathParts(StagePathParts):
         directory = dir
         stage = @db.schema.stage
         stage_name = stage
-    For `@stage/dir` to
-        stage -> @stage
+    For `snow://stage/dir` to
+        stage -> snow://stage
         stage_name -> stage
         directory -> dir
     """
@@ -138,12 +146,12 @@ class DefaultStagePathParts(StagePathParts):
             raise ClickException("Invalid stage path")
         self.directory = match.group("directory")
         self._schema = match.group("second_qualifier") or match.group("first_qualifier")
+        self._prefix = match.group("prefix") or "@"
         self.stage = stage_path.removesuffix(self.directory).rstrip("/")
 
         stage_name = FQN.from_stage(self.stage).name
-        stage_name = (
-            stage_name[OMIT_FIRST] if stage_name.startswith("@") else stage_name
-        )
+        if stage_name.startswith(self._prefix):
+            stage_name = stage_name.removeprefix(self._prefix)
         self.stage_name = stage_name
         self.is_directory = True if stage_path.endswith("/") else False
 
@@ -167,13 +175,12 @@ class DefaultStagePathParts(StagePathParts):
         return self._schema
 
     def replace_stage_prefix(self, file_path: str) -> str:
-        stage = Path(self.stage).parts[0]
+        file_path = _strip_standard_stage_prefix(file_path)
         file_path_without_prefix = Path(file_path).parts[OMIT_FIRST]
-        return f"{stage}/{'/'.join(file_path_without_prefix)}"
+        return f"{self.stage}/{'/'.join(file_path_without_prefix)}"
 
     def strip_stage_prefix(self, file_path: str) -> str:
-        if file_path.startswith("@"):
-            file_path = file_path[OMIT_FIRST]
+        file_path = _strip_standard_stage_prefix(file_path)
         if file_path.startswith(self.stage_name):
             return file_path[len(self.stage_name) :]
         return file_path
@@ -439,9 +446,13 @@ class StageManager(SqlExecutionMixin):
                 # We end if we reach the root directory
                 if directory == temp_dir_with_copy:
                     break
-
                 # Add parent directory to the list if it's not already there
-                if directory.parent not in deepest_dirs_list:
+                if directory.parent not in deepest_dirs_list and not any(
+                    (
+                        existing_dir.is_relative_to(directory.parent)
+                        for existing_dir in deepest_dirs_list
+                    )
+                ):
                     deepest_dirs_list.append(directory.parent)
 
                 # Remove the directory so the parent directory will contain only files
@@ -703,6 +714,7 @@ class StageManager(SqlExecutionMixin):
         original_file: str,
     ) -> Dict:
         try:
+            log.info("Executing SQL file: %s", file_stage_path)
             query = f"execute immediate from {self.quote_stage_name(file_stage_path)}"
             if variables:
                 query += variables
@@ -816,6 +828,7 @@ class StageManager(SqlExecutionMixin):
         from snowflake.snowpark.exceptions import SnowparkSQLException
 
         try:
+            log.info("Executing Python file: %s", file_stage_path)
             self._python_exe_procedure(self.get_standard_stage_prefix(file_stage_path), variables, session=self.snowpark_session)  # type: ignore
             return StageManager._success_result(file=original_file)
         except SnowparkSQLException as e:

snowflake/cli/_plugins/streamlit/commands.py

@@ -37,6 +37,7 @@ from snowflake.cli.api.commands.decorators import (
     with_project_definition,
 )
 from snowflake.cli.api.commands.flags import (
+    PruneOption,
     ReplaceOption,
     entity_argument,
     identifier_argument,
@@ -136,6 +137,7 @@ def streamlit_deploy(
         help="Replaces the Streamlit app if it already exists. It only uploads new and overwrites existing files, "
         "but does not remove any files already on the stage."
     ),
+    prune: bool = PruneOption(),
     entity_id: str = entity_argument("streamlit"),
     open_: bool = OpenOption,
     **options,
@@ -168,6 +170,7 @@ def streamlit_deploy(
         streamlit=streamlit,
         streamlit_project_paths=streamlit_project_paths,
         replace=replace,
+        prune=prune,
     )
 
     if open_:

snowflake/cli/_plugins/streamlit/manager.py

@@ -31,7 +31,7 @@ from snowflake.cli._plugins.streamlit.streamlit_entity_model import (
 from snowflake.cli._plugins.streamlit.streamlit_project_paths import (
     StreamlitProjectPaths,
 )
-from snowflake.cli.api.artifacts.upload import put_files
+from snowflake.cli.api.artifacts.upload import sync_artifacts_with_stage
 from snowflake.cli.api.commands.experimental_behaviour import (
     experimental_behaviour_enabled,
 )
@@ -56,18 +56,20 @@ class StreamlitManager(SqlExecutionMixin):
             f"grant usage on streamlit {streamlit_name.sql_identifier} to role {to_role}"
         )
 
-    def _put_streamlit_files(
+    def _upload_artifacts(
         self,
         streamlit_project_paths: StreamlitProjectPaths,
         stage_root: str,
+        prune: bool,
         artifacts: Optional[List[PathMapping]] = None,
     ):
-        cli_console.step(f"Deploying files to {stage_root}")
-        put_files(
-            project_paths=streamlit_project_paths,
-            stage_root=stage_root,
-            artifacts=artifacts,
-        )
+        with cli_console.phase(f"Deploying files to {stage_root}"):
+            sync_artifacts_with_stage(
+                project_paths=streamlit_project_paths,
+                stage_root=stage_root,
+                prune=prune,
+                artifacts=artifacts,
+            )
 
     def _create_streamlit(
         self,
@@ -126,6 +128,7 @@ class StreamlitManager(SqlExecutionMixin):
         streamlit: StreamlitEntityModel,
         streamlit_project_paths: StreamlitProjectPaths,
         replace: bool = False,
+        prune: bool = False,
     ):
         streamlit_id = streamlit.fqn.using_connection(self._conn)
         if (
@@ -182,10 +185,11 @@ class StreamlitManager(SqlExecutionMixin):
             else:
                 stage_root = f"{embedded_stage_name}/default_checkout"
 
-            self._put_streamlit_files(
+            self._upload_artifacts(
                 streamlit_project_paths,
                 stage_root,
-                streamlit.artifacts,
+                prune=prune,
+                artifacts=streamlit.artifacts,
             )
         else:
             """
@@ -198,15 +202,15 @@ class StreamlitManager(SqlExecutionMixin):
             stage_name = streamlit.stage or "streamlit"
             stage_name = FQN.from_string(stage_name).using_connection(self._conn)
 
-            cli_console.step(f"Creating {stage_name} stage")
-            stage_manager.create(fqn=stage_name)
-
             stage_root = stage_manager.get_standard_stage_prefix(
                 f"{stage_name}/{streamlit_name_for_root_location}"
             )
 
-            self._put_streamlit_files(
-                streamlit_project_paths, stage_root, streamlit.artifacts
+            self._upload_artifacts(
+                streamlit_project_paths,
+                stage_root,
+                prune=prune,
+                artifacts=streamlit.artifacts,
             )
 
             self._create_streamlit(

snowflake/cli/api/artifacts/upload.py

@@ -1,51 +1,47 @@
-from pathlib import PurePosixPath
 from typing import List, Optional
 
 from snowflake.cli._plugins.stage.manager import StageManager
-from snowflake.cli.api.artifacts.bundle_map import BundleMap
-from snowflake.cli.api.artifacts.utils import symlink_or_copy
+from snowflake.cli.api.artifacts.utils import bundle_artifacts
 from snowflake.cli.api.console import cli_console
+from snowflake.cli.api.entities.utils import sync_deploy_root_with_stage
 from snowflake.cli.api.project.project_paths import ProjectPaths
 from snowflake.cli.api.project.schemas.entities.common import PathMapping
 
 
-def put_files(
+def sync_artifacts_with_stage(
     project_paths: ProjectPaths,
     stage_root: str,
+    prune: bool = False,
     artifacts: Optional[List[PathMapping]] = None,
 ):
-    if not artifacts:
-        return
-    stage_manager = StageManager()
+    if artifacts is None:
+        artifacts = []
+
+    bundle_map = bundle_artifacts(project_paths, artifacts)
+    stage_path_parts = StageManager().stage_path_parts_from_str(stage_root)
     # We treat the bundle root as deploy root
-    bundle_map = BundleMap(
-        project_root=project_paths.project_root,
+    sync_deploy_root_with_stage(
+        console=cli_console,
         deploy_root=project_paths.bundle_root,
+        bundle_map=bundle_map,
+        prune=prune,
+        recursive=True,
+        stage_path=stage_path_parts,
+        print_diff=True,
     )
-    for artifact in artifacts:
-        bundle_map.add(PathMapping(src=str(artifact.src), dest=artifact.dest))
 
-    # Clean up bundle root
-    project_paths.remove_up_bundle_root()
 
-    for (absolute_src, absolute_dest) in bundle_map.all_mappings(
-        absolute=True, expand_directories=True
-    ):
-        if absolute_src.is_file():
-            # We treat the bundle/streamlit root as deploy root
-            symlink_or_copy(
-                absolute_src,
-                absolute_dest,
-                deploy_root=project_paths.bundle_root,
-            )
-            # Temporary solution, will be replaced with diff
-            stage_path = (
-                PurePosixPath(absolute_dest)
-                .relative_to(project_paths.bundle_root)
-                .parent
-            )
-            full_stage_path = f"{stage_root}/{stage_path}".rstrip("/")
-            cli_console.step(f"Uploading {absolute_dest} to {full_stage_path}")
-            stage_manager.put(
-                local_path=absolute_dest, stage_path=full_stage_path, overwrite=True
-            )
+def put_files(
+    project_paths: ProjectPaths,
+    stage_root: str,
+    artifacts: Optional[List[PathMapping]] = None,
+):
+    if not artifacts:
+        return
+
+    sync_artifacts_with_stage(
+        project_paths=project_paths,
+        stage_root=stage_root,
+        prune=False,
+        artifacts=artifacts,
+    )

snowflake/cli/api/artifacts/utils.py

@@ -69,14 +69,16 @@ def bundle_artifacts(project_paths: ProjectPaths, artifacts: Artifacts) -> Bundl
         bundle_map.add(artifact)
 
     project_paths.remove_up_bundle_root()
+    SecurePath(project_paths.bundle_root).mkdir(parents=True, exist_ok=True)
     for absolute_src, absolute_dest in bundle_map.all_mappings(
         absolute=True, expand_directories=True
     ):
-        # We treat the bundle root as deploy root
-        symlink_or_copy(
-            absolute_src,
-            absolute_dest,
-            deploy_root=project_paths.bundle_root,
-        )
+        if absolute_src.is_file():
+            # We treat the bundle root as deploy root
+            symlink_or_copy(
+                absolute_src,
+                absolute_dest,
+                deploy_root=project_paths.bundle_root,
+            )
 
     return bundle_map

snowflake/cli/api/cli_global_context.py

@@ -22,7 +22,7 @@ from pathlib import Path
 from typing import TYPE_CHECKING, Iterator, Optional
 
 from snowflake.cli.api.connections import ConnectionContext, OpenConnectionCache
-from snowflake.cli.api.exceptions import MissingConfiguration
+from snowflake.cli.api.exceptions import MissingConfigurationError
 from snowflake.cli.api.metrics import CLIMetrics
 from snowflake.cli.api.output.formats import OutputFormat
 from snowflake.cli.api.rendering.jinja import CONTEXT_KEY
@@ -61,6 +61,7 @@ class _CliGlobalContextManager:
     override_project_definition: ProjectDefinition | None = None
 
     _definition_manager: DefinitionManager | None = None
+    enhanced_exit_codes: bool = False
 
     # which properties invalidate our current DefinitionManager?
     DEFINITION_MANAGER_DEPENDENCIES = [
@@ -126,7 +127,7 @@ class _CliGlobalContextManager:
                 {CONTEXT_KEY: {"env": self.project_env_overrides_args}},
             )
             if not dm.has_definition_file and not self.project_is_optional:
-                raise MissingConfiguration(
+                raise MissingConfigurationError(
                     "Cannot find project definition (snowflake.yml). Please provide a path to the project or run this command in a valid project directory."
                 )
             self._definition_manager = dm
@@ -209,6 +210,10 @@ class _CliGlobalContextAccess:
         else:
             return None
 
+    @property
+    def enhanced_exit_codes(self) -> bool:
+        return self._manager.enhanced_exit_codes
+
 
 _CLI_CONTEXT_MANAGER: ContextVar[_CliGlobalContextManager | None] = ContextVar(
     "cli_context", default=None

snowflake/cli/api/commands/decorators.py

@@ -29,6 +29,7 @@ from snowflake.cli.api.commands.flags import (
     DiagAllowlistPathOption,
     DiagLogPathOption,
     EnableDiagOption,
+    EnhancedExitCodesOption,
     HostOption,
     MasterTokenOption,
     MfaPasscodeOption,
@@ -78,7 +79,6 @@ def global_options_with_connection(func: Callable):
 
 def with_project_definition(is_optional: bool = False):
     def _decorator(func: Callable):
-
         return _options_decorator_factory(
             func,
             additional_options=[
@@ -159,7 +159,10 @@ def _options_decorator_factory(
             execute_before_command_using_new_options(**options)
         return func(**options)
 
-    wrapper.__signature__ = _extend_signature_with_additional_options(func, additional_options)  # type: ignore
+    wrapper.__signature__ = _extend_signature_with_additional_options(  # type: ignore
+        func, additional_options
+    )
+
     return wrapper
 
 
@@ -353,6 +356,12 @@ GLOBAL_OPTIONS = [
         annotation=Optional[bool],
         default=SilentOption,
     ),
+    inspect.Parameter(
+        "enhanced_exit_codes",
+        inspect.Parameter.KEYWORD_ONLY,
+        annotation=Optional[bool],
+        default=EnhancedExitCodesOption,
+    ),
 ]