snowflake-cli 3.4.1__py3-none-any.whl → 3.6.0__py3-none-any.whl

snowflake/cli/_plugins/sql/manager.py

@@ -14,23 +14,29 @@
 
 from __future__ import annotations
 
+import logging
 import sys
-from io import StringIO
-from itertools import chain
+from functools import partial
 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._plugins.sql.source_reader import (
+    compile_statements,
+    files_reader,
+    query_reader,
+)
+from snowflake.cli.api.console import cli_console
 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.secure_path import 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
 
+logger = logging.getLogger(__name__)
+
 
 class SqlManager(SqlExecutionMixin):
     def execute(
@@ -41,57 +47,42 @@ class SqlManager(SqlExecutionMixin):
         data: Dict | None = None,
         retain_comments: bool = False,
     ) -> 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."
-            )
+        """Reads, transforms and execute statements from input.
 
-        if std_in:
-            query = sys.stdin.read()
-        if query:
-            return self._execute_single_query(
-                query=query, data=data, retain_comments=retain_comments
-            )
+        Only one input can be consumed at a time.
+        When no compilation errors are detected, the sequence on queries
+        in executed and returned as tuple.
 
-        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, retain_comments=retain_comments
-                )
-                results.append(result)
+        Throws an exception ff multiple inputs are provided.
+        """
+        query = sys.stdin.read() if std_in else query
 
-            # 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)
+        stmt_operators = (
+            transpile_snowsql_templates,
+            partial(snowflake_sql_jinja_render, data=data),
+        )
+        remove_comments = not retain_comments
 
-        # At that point, no stdin, query or files were provided
-        raise UsageError("Use either query, filename or input option.")
+        if query:
+            stmt_reader = query_reader(query, stmt_operators, remove_comments)
+        elif files:
+            secured_files = [SecurePath(f) for f in files]
+            stmt_reader = files_reader(secured_files, stmt_operators, remove_comments)
+        else:
+            raise UsageError("Use either query, filename or input option.")
 
-    def _execute_single_query(
-        self, query: str, data: Dict | None = None, retain_comments: bool = False
-    ) -> 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}")
+        errors, stmt_count, compiled_statements = compile_statements(stmt_reader)
+        if not any((errors, stmt_count, compiled_statements)):
+            raise UsageError("Use either query, filename or input option.")
 
-        statements = tuple(
-            statement
-            for statement, _ in split_statements(
-                StringIO(query), remove_comments=not retain_comments
-            )
-        )
-        single_statement = len(statements) == 1
+        if errors:
+            for error in errors:
+                logger.info("Statement compilation error: %s", error)
+                cli_console.warning(error)
+            raise ClickException("SQL rendering error")
 
-        return single_statement, self._execute_string(
-            "\n".join(statements), cursor_class=VerboseCursor
+        is_single_statement = not (stmt_count > 1)
+        return is_single_statement, self.execute_string(
+            "\n".join(compiled_statements),
+            cursor_class=VerboseCursor,
         )

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

@@ -439,9 +439,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
@@ -561,7 +565,7 @@ class StageManager(SqlExecutionMixin):
         )
 
         parsed_variables = parse_key_value_variables(variables)
-        sql_variables = self._parse_execute_variables(parsed_variables)
+        sql_variables = self.parse_execute_variables(parsed_variables)
         python_variables = self._parse_python_variables(parsed_variables)
         results = []
 
@@ -663,7 +667,7 @@ class StageManager(SqlExecutionMixin):
         return [f for f in files if Path(f).suffix in EXECUTE_SUPPORTED_FILES_FORMATS]
 
     @staticmethod
-    def _parse_execute_variables(variables: List[Variable]) -> Optional[str]:
+    def parse_execute_variables(variables: List[Variable]) -> Optional[str]:
         if not variables:
             return None
         query_parameters = [f"{v.key}=>{v.value}" for v in variables]
@@ -703,6 +707,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 +821,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

@@ -16,11 +16,10 @@ from __future__ import annotations
 
 import logging
 from pathlib import Path
-from typing import Dict
 
 import click
 import typer
-from click import ClickException, UsageError
+from click import ClickException
 from snowflake.cli._plugins.object.command_aliases import (
     add_object_command_aliases,
     scope_option,
@@ -44,6 +43,7 @@ from snowflake.cli.api.commands.flags import (
     like_option,
 )
 from snowflake.cli.api.commands.snow_typer import SnowTyperFactory
+from snowflake.cli.api.commands.utils import get_entity_for_operation
 from snowflake.cli.api.constants import ObjectType
 from snowflake.cli.api.exceptions import NoProjectDefinitionError
 from snowflake.cli.api.identifiers import FQN
@@ -133,7 +133,8 @@ def _default_file_callback(param_name: str):
 @with_experimental_behaviour()
 def streamlit_deploy(
     replace: bool = ReplaceOption(
-        help="Replace the Streamlit app if it already exists."
+        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."
     ),
     entity_id: str = entity_argument("streamlit"),
     open_: bool = OpenOption,
@@ -155,30 +156,14 @@ def streamlit_deploy(
             )
         pd = convert_project_definition_to_v2(cli_context.project_root, pd)
 
-    streamlits: Dict[str, StreamlitEntityModel] = pd.get_entities_by_type(
-        entity_type="streamlit"
+    streamlit: StreamlitEntityModel = get_entity_for_operation(
+        cli_context=cli_context,
+        entity_id=entity_id,
+        project_definition=pd,
+        entity_type="streamlit",
     )
 
     streamlit_project_paths = StreamlitProjectPaths(cli_context.project_root)
-
-    if not streamlits:
-        raise NoProjectDefinitionError(
-            project_type="streamlit", project_root=cli_context.project_root
-        )
-
-    if entity_id and entity_id not in streamlits:
-        raise UsageError(f"No '{entity_id}' entity in project definition file.")
-
-    if len(streamlits.keys()) == 1:
-        entity_id = list(streamlits.keys())[0]
-
-    if entity_id is None:
-        raise UsageError(
-            "Multiple Streamlit apps found. Please provide entity id for the operation."
-        )
-
-    # Get first streamlit
-    streamlit: StreamlitEntityModel = streamlits[entity_id]
     url = StreamlitManager().deploy(
         streamlit=streamlit,
         streamlit_project_paths=streamlit_project_paths,

snowflake/cli/_plugins/streamlit/manager.py

@@ -15,7 +15,6 @@
 from __future__ import annotations
 
 import logging
-from pathlib import PurePosixPath
 from typing import List, Optional
 
 from click import ClickException
@@ -32,8 +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.bundle_map import BundleMap
-from snowflake.cli.api.artifacts.utils import symlink_or_copy
+from snowflake.cli.api.artifacts.upload import put_files
 from snowflake.cli.api.commands.experimental_behaviour import (
     experimental_behaviour_enabled,
 )
@@ -65,40 +63,11 @@ class StreamlitManager(SqlExecutionMixin):
         artifacts: Optional[List[PathMapping]] = None,
     ):
         cli_console.step(f"Deploying files to {stage_root}")
-        if not artifacts:
-            return
-        stage_manager = StageManager()
-        # We treat the bundle root as deploy root
-        bundle_map = BundleMap(
-            project_root=streamlit_project_paths.project_root,
-            deploy_root=streamlit_project_paths.bundle_root,
+        put_files(
+            project_paths=streamlit_project_paths,
+            stage_root=stage_root,
+            artifacts=artifacts,
         )
-        for artifact in artifacts:
-            bundle_map.add(PathMapping(src=str(artifact.src), dest=artifact.dest))
-
-        # Clean up bundle root
-        streamlit_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=streamlit_project_paths.bundle_root,
-                )
-                # Temporary solution, will be replaced with diff
-                stage_path = (
-                    PurePosixPath(absolute_dest)
-                    .relative_to(streamlit_project_paths.bundle_root)
-                    .parent
-                )
-                full_stage_path = f"{stage_root}/{stage_path}".rstrip("/")
-                stage_manager.put(
-                    local_path=absolute_dest, stage_path=full_stage_path, overwrite=True
-                )
 
     def _create_streamlit(
         self,

snowflake/cli/api/artifacts/upload.py

@@ -0,0 +1,51 @@
+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.console import cli_console
+from snowflake.cli.api.project.project_paths import ProjectPaths
+from snowflake.cli.api.project.schemas.entities.common import PathMapping
+
+
+def put_files(
+    project_paths: ProjectPaths,
+    stage_root: str,
+    artifacts: Optional[List[PathMapping]] = None,
+):
+    if not artifacts:
+        return
+    stage_manager = StageManager()
+    # We treat the bundle root as deploy root
+    bundle_map = BundleMap(
+        project_root=project_paths.project_root,
+        deploy_root=project_paths.bundle_root,
+    )
+    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
+            )