snowflake-cli 3.7.2__py3-none-any.whl → 3.8.1__py3-none-any.whl

snowflake/cli/_plugins/sql/statement_reader.py

@@ -0,0 +1,296 @@
+import enum
+import io
+import re
+import urllib.error
+from dataclasses import dataclass
+from typing import Any, Callable, Generator, List, Literal, Sequence, Tuple
+from urllib.request import urlopen
+
+from jinja2 import UndefinedError
+from snowflake.cli._plugins.sql.snowsql_commands import (
+    SnowSQLCommand,
+    compile_snowsql_command,
+)
+from snowflake.cli.api.secure_path import UNLIMITED, SecurePath
+from snowflake.connector.util_text import split_statements
+
+COMMAND_PATTERN = re.compile(
+    r"^!(\w+)\s*[\"']?(.*?)[\"']?\s*(?:;|$)",
+    flags=re.IGNORECASE,
+)
+URL_PATTERN = re.compile(r"^(\w+?):\/(\/.*)", flags=re.IGNORECASE)
+
+ASYNC_SUFFIX = ";>"
+
+
+SplitedStatements = Generator[
+    tuple[str, bool | None] | tuple[str, Literal[False]],
+    Any,
+    None,
+]
+
+SqlTransformFunc = Callable[[str], str]
+OperatorFunctions = Sequence[SqlTransformFunc]
+
+
+class StatementType(enum.Enum):
+    FILE = "file"
+    QUERY = "query"
+    UNKNOWN = "unknown"
+    URL = "url"
+    SNOWSQL_COMMAND = "snowsql_command"
+
+
+class ParsedStatement:
+    """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__ = ("statement", "statement_type", "source_path", "error")
+    __match_args__ = ("statement_type", "error")
+
+    statement: io.StringIO
+    statement_type: StatementType | None
+    source_path: str | None
+    error: str | None
+
+    def __init__(
+        self,
+        statement: str,
+        source_type: StatementType,
+        source_path: str | None,
+        error: str | None = None,
+    ):
+        self.statement = io.StringIO(statement)
+        self.statement_type = source_type
+        self.source_path = source_path
+        self.error = error
+
+    def __bool__(self):
+        return not self.error
+
+    def __eq__(self, other):
+        result = (
+            self.statement_type == other.statement_type,
+            self.source_path == other.source_path,
+            self.error == other.error,
+            self.statement.read() == other.statement.read(),
+        )
+        self.statement.seek(0)
+        other.statement.seek(0)
+        return all(result)
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}(statement_type={self.statement_type}, source_path={self.source_path}, error={self.error})"
+
+    @classmethod
+    def from_url(cls, path_part: str, raw_source: str) -> "ParsedStatement":
+        """Constructor for loading from URL."""
+        try:
+            payload = urlopen(path_part, timeout=10.0).read().decode()
+            return cls(payload, StatementType.URL, path_part)
+
+        except urllib.error.HTTPError as err:
+            error = f"Could not fetch {path_part}: {err}"
+            return cls(path_part, StatementType.URL, raw_source, error)
+
+    @classmethod
+    def from_file(cls, path_part: str, raw_source: str) -> "ParsedStatement":
+        """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, StatementType.FILE, path.as_posix())
+
+        error_msg = f"Could not read: {path_part}"
+        return cls(path_part, StatementType.FILE, raw_source, error_msg)
+
+
+RecursiveStatementReader = Generator[ParsedStatement, Any, Any]
+
+
+def parse_statement(source: str, operators: OperatorFunctions) -> ParsedStatement:
+    """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 ParsedStatement(source, StatementType.UNKNOWN, source, error_msg)
+
+    split_result = COMMAND_PATTERN.split(statement, maxsplit=1)
+    split_result = [p.strip() for p in split_result]
+
+    if len(split_result) == 1:
+        return ParsedStatement(statement, StatementType.QUERY, None)
+
+    _, command, command_args, *_ = split_result
+    _path_match = URL_PATTERN.split(command_args.lower())
+
+    match command.lower(), _path_match:
+        # load content from an URL
+        case "source" | "load", ("", "http" | "https", *_):
+            return ParsedStatement.from_url(command_args, statement)
+
+        # load content from a local file
+        case "source" | "load", (str(),):
+            return ParsedStatement.from_file(command_args, statement)
+
+        case "source" | "load", _:
+            return ParsedStatement(
+                statement,
+                StatementType.UNKNOWN,
+                command_args,
+                f"Unknown source: {command_args}",
+            )
+
+        case "queries" | "result" | "abort", (str(),):
+            return ParsedStatement(statement, StatementType.SNOWSQL_COMMAND, None)
+
+        case _:
+            error_msg = f"Unknown command: {command}"
+
+    return ParsedStatement(statement, StatementType.UNKNOWN, None, error_msg)
+
+
+def recursive_statement_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_statement(stmt, operators)
+
+        match parsed_source:
+            case ParsedStatement(StatementType.FILE | StatementType.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_statement_reader(
+                    split_statements(parsed_source.statement, remove_comments),
+                    seen_files,
+                    operators,
+                    remove_comments,
+                )
+
+                seen_files.pop()
+
+            case ParsedStatement(StatementType.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_statement_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."""
+    # known issue of split_statements (doesn't work in SnowSQL either):
+    # when the line starts with a command:
+    # '!queries amount=3; select 3;'
+    # it is treated as a single statement
+    stmts = split_statements(io.StringIO(source), remove_comments)
+    yield from recursive_statement_reader(stmts, [], operators, remove_comments)
+
+
+@dataclass
+class CompiledStatement:
+    statement: str | None = None
+    execute_async: bool = False
+    command: SnowSQLCommand | None = None
+
+
+def _is_empty_statement(statement: str) -> bool:
+    # checks whether all lines from the statement are empty or start with comment
+    for line in statement.splitlines():
+        if line.strip() and not line.lstrip().startswith("--"):
+            # nonempty uncommented line
+            return False
+    return True
+
+
+def compile_statements(
+    source: RecursiveStatementReader,
+) -> Tuple[List[str], int, List[CompiledStatement]]:
+    """Tracks statements evaluation and collects errors."""
+    errors = []
+    expected_results_cnt = 0
+    compiled = []
+
+    for stmt in source:
+        if stmt.statement_type == StatementType.QUERY:
+            statement = stmt.statement.read()
+            if not stmt.error and not _is_empty_statement(statement):
+                is_async = statement.endswith(ASYNC_SUFFIX)
+                compiled.append(
+                    CompiledStatement(
+                        statement=statement.removesuffix(ASYNC_SUFFIX),
+                        execute_async=is_async,
+                    )
+                )
+                if not is_async:
+                    expected_results_cnt += 1
+
+        if stmt.statement_type == StatementType.SNOWSQL_COMMAND:
+            if not stmt.error:
+                cmd = (
+                    stmt.statement.read()
+                    .removesuffix(ASYNC_SUFFIX)
+                    .removesuffix(";")
+                    .split()
+                )
+                parsed_command = compile_snowsql_command(
+                    command=cmd[0], cmd_args=cmd[1:]
+                )
+                if parsed_command.error_message:
+                    errors.append(parsed_command.error_message)
+                else:
+                    compiled.append(CompiledStatement(command=parsed_command.command))
+
+        if stmt.error:
+            errors.append(stmt.error)
+
+    return errors, expected_results_cnt, compiled

snowflake/cli/_plugins/streamlit/commands.py

@@ -25,12 +25,8 @@ from snowflake.cli._plugins.object.command_aliases import (
     scope_option,
 )
 from snowflake.cli._plugins.streamlit.manager import StreamlitManager
-from snowflake.cli._plugins.streamlit.streamlit_entity_model import (
-    StreamlitEntityModel,
-)
-from snowflake.cli._plugins.streamlit.streamlit_project_paths import (
-    StreamlitProjectPaths,
-)
+from snowflake.cli._plugins.streamlit.streamlit_entity import StreamlitEntity
+from snowflake.cli._plugins.workspace.context import ActionContext, WorkspaceContext
 from snowflake.cli.api.cli_global_context import get_cli_context
 from snowflake.cli.api.commands.decorators import (
     with_experimental_behaviour,
@@ -45,7 +41,9 @@ from snowflake.cli.api.commands.flags import (
 )
 from snowflake.cli.api.commands.snow_typer import SnowTyperFactory
 from snowflake.cli.api.commands.utils import get_entity_for_operation
+from snowflake.cli.api.console.console import CliConsole
 from snowflake.cli.api.constants import ObjectType
+from snowflake.cli.api.entities.utils import EntityActions
 from snowflake.cli.api.exceptions import NoProjectDefinitionError
 from snowflake.cli.api.identifiers import FQN
 from snowflake.cli.api.output.types import (
@@ -158,18 +156,24 @@ def streamlit_deploy(
             )
         pd = convert_project_definition_to_v2(cli_context.project_root, pd)
 
-    streamlit: StreamlitEntityModel = get_entity_for_operation(
-        cli_context=cli_context,
-        entity_id=entity_id,
-        project_definition=pd,
-        entity_type="streamlit",
+    streamlit: StreamlitEntity = StreamlitEntity(
+        entity_model=get_entity_for_operation(
+            cli_context=cli_context,
+            entity_id=entity_id,
+            project_definition=pd,
+            entity_type=ObjectType.STREAMLIT.value.cli_name,
+        ),
+        workspace_ctx=_get_current_workspace_context(),
     )
 
-    streamlit_project_paths = StreamlitProjectPaths(cli_context.project_root)
-    url = StreamlitManager().deploy(
-        streamlit=streamlit,
-        streamlit_project_paths=streamlit_project_paths,
+    url = streamlit.perform(
+        EntityActions.DEPLOY,
+        ActionContext(
+            get_entity=lambda *args: None,
+        ),
+        _open=open_,
         replace=replace,
+        experimental=options.get("experimental"),
         prune=prune,
     )
 
@@ -190,3 +194,14 @@ def get_url(
     if open_:
         typer.launch(url)
     return MessageResult(url)
+
+
+def _get_current_workspace_context():
+    ctx = get_cli_context()
+
+    return WorkspaceContext(
+        console=CliConsole(),
+        project_root=ctx.project_root,
+        get_default_role=lambda: ctx.connection.role,
+        get_default_warehouse=lambda: ctx.connection.warehouse,
+    )

snowflake/cli/_plugins/streamlit/manager.py

@@ -15,33 +15,18 @@
 from __future__ import annotations
 
 import logging
-from typing import List, Optional
 
-from click import ClickException
 from snowflake.cli._plugins.connection.util import (
     MissingConnectionAccountError,
     MissingConnectionRegionError,
     make_snowsight_url,
 )
-from snowflake.cli._plugins.object.manager import ObjectManager
-from snowflake.cli._plugins.stage.manager import StageManager
 from snowflake.cli._plugins.streamlit.streamlit_entity_model import (
     StreamlitEntityModel,
 )
-from snowflake.cli._plugins.streamlit.streamlit_project_paths import (
-    StreamlitProjectPaths,
-)
-from snowflake.cli.api.artifacts.upload import sync_artifacts_with_stage
-from snowflake.cli.api.commands.experimental_behaviour import (
-    experimental_behaviour_enabled,
-)
-from snowflake.cli.api.console import cli_console
-from snowflake.cli.api.feature_flags import FeatureFlag
 from snowflake.cli.api.identifiers import FQN
-from snowflake.cli.api.project.schemas.entities.common import PathMapping
 from snowflake.cli.api.sql_execution import SqlExecutionMixin
 from snowflake.connector.cursor import SnowflakeCursor
-from snowflake.connector.errors import ProgrammingError
 
 log = logging.getLogger(__name__)
 
@@ -56,174 +41,6 @@ class StreamlitManager(SqlExecutionMixin):
             f"grant usage on streamlit {streamlit_name.sql_identifier} to role {to_role}"
         )
 
-    def _upload_artifacts(
-        self,
-        streamlit_project_paths: StreamlitProjectPaths,
-        stage_root: str,
-        prune: bool,
-        artifacts: Optional[List[PathMapping]] = None,
-    ):
-        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,
-        streamlit: StreamlitEntityModel,
-        replace: Optional[bool] = None,
-        experimental: Optional[bool] = None,
-        from_stage_name: Optional[str] = None,
-    ):
-        streamlit_id = streamlit.fqn.using_connection(self._conn)
-        cli_console.step(f"Creating {streamlit_id} Streamlit")
-        query = []
-        if replace:
-            query.append(f"CREATE OR REPLACE STREAMLIT {streamlit_id.sql_identifier}")
-        elif experimental:
-            # For experimental behaviour, we need to use CREATE STREAMLIT IF NOT EXISTS
-            # for a streamlit app with an embedded stage
-            # because this is analogous to the behavior for non-experimental
-            # deploy which does CREATE STAGE IF NOT EXISTS
-            query.append(
-                f"CREATE STREAMLIT IF NOT EXISTS {streamlit_id.sql_identifier}"
-            )
-        else:
-            query.append(f"CREATE STREAMLIT {streamlit_id.sql_identifier}")
-
-        if from_stage_name:
-            query.append(f"ROOT_LOCATION = '{from_stage_name}'")
-
-        query.append(f"MAIN_FILE = '{streamlit.main_file}'")
-        if streamlit.imports:
-            query.append(streamlit.get_imports_sql())
-
-        if not streamlit.query_warehouse:
-            cli_console.warning(
-                "[Deprecation] In next major version we will remove default query_warehouse='streamlit'."
-            )
-            query.append(f"QUERY_WAREHOUSE = 'streamlit'")
-        else:
-            query.append(f"QUERY_WAREHOUSE = {streamlit.query_warehouse}")
-
-        if streamlit.title:
-            query.append(f"TITLE = '{streamlit.title}'")
-
-        if streamlit.comment:
-            query.append(f"COMMENT = '{streamlit.comment}'")
-
-        if streamlit.external_access_integrations:
-            query.append(streamlit.get_external_access_integrations_sql())
-
-        if streamlit.secrets:
-            query.append(streamlit.get_secrets_sql())
-
-        self.execute_query("\n".join(query))
-
-    def deploy(
-        self,
-        streamlit: StreamlitEntityModel,
-        streamlit_project_paths: StreamlitProjectPaths,
-        replace: bool = False,
-        prune: bool = False,
-    ):
-        streamlit_id = streamlit.fqn.using_connection(self._conn)
-        if (
-            ObjectManager().object_exists(object_type="streamlit", fqn=streamlit_id)
-            and not replace
-        ):
-            raise ClickException(
-                f"Streamlit {streamlit.fqn} already exist. If you want to replace it use --replace flag."
-            )
-
-        # for backwards compatibility - quoted stage path might be case-sensitive
-        # https://docs.snowflake.com/en/sql-reference/identifiers-syntax#double-quoted-identifiers
-        streamlit_name_for_root_location = streamlit_id.name
-        use_versioned_stage = FeatureFlag.ENABLE_STREAMLIT_VERSIONED_STAGE.is_enabled()
-        if (
-            experimental_behaviour_enabled()
-            or FeatureFlag.ENABLE_STREAMLIT_EMBEDDED_STAGE.is_enabled()
-            or use_versioned_stage
-        ):
-            """
-            1. Create streamlit object
-            2. Upload files to embedded stage
-            """
-            # TODO: Support from_stage
-            # from_stage_stmt = f"FROM_STAGE = '{stage_name}'" if stage_name else ""
-            self._create_streamlit(
-                streamlit=streamlit,
-                replace=replace,
-                experimental=True,
-            )
-            try:
-                if use_versioned_stage:
-                    self.execute_query(
-                        f"ALTER STREAMLIT {streamlit_id.identifier} ADD LIVE VERSION FROM LAST"
-                    )
-                elif not FeatureFlag.ENABLE_STREAMLIT_NO_CHECKOUTS.is_enabled():
-                    self.execute_query(
-                        f"ALTER streamlit {streamlit_id.identifier} CHECKOUT"
-                    )
-            except ProgrammingError as e:
-                # If an error is raised because a CHECKOUT has already occurred or a LIVE VERSION already exists, simply skip it and continue
-                if "Checkout already exists" in str(
-                    e
-                ) or "There is already a live version" in str(e):
-                    log.info("Checkout already exists, continuing")
-                else:
-                    raise
-
-            stage_path = streamlit_id.identifier
-            embedded_stage_name = f"snow://streamlit/{stage_path}"
-            if use_versioned_stage:
-                # "LIVE" is the only supported version for now, but this may change later.
-                stage_root = f"{embedded_stage_name}/versions/live"
-            else:
-                stage_root = f"{embedded_stage_name}/default_checkout"
-
-            self._upload_artifacts(
-                streamlit_project_paths,
-                stage_root,
-                prune=prune,
-                artifacts=streamlit.artifacts,
-            )
-        else:
-            """
-            1. Create stage
-            2. Upload files to created stage
-            3. Create streamlit from stage
-            """
-            stage_manager = StageManager()
-
-            stage_name = streamlit.stage or "streamlit"
-            stage_name = FQN.from_string(stage_name).using_connection(self._conn)
-
-            stage_root = stage_manager.get_standard_stage_prefix(
-                f"{stage_name}/{streamlit_name_for_root_location}"
-            )
-
-            self._upload_artifacts(
-                streamlit_project_paths,
-                stage_root,
-                prune=prune,
-                artifacts=streamlit.artifacts,
-            )
-
-            self._create_streamlit(
-                streamlit=streamlit,
-                replace=replace,
-                from_stage_name=stage_root,
-                experimental=False,
-            )
-
-            self.grant_privileges(streamlit)
-
-        return self.get_url(streamlit_name=streamlit_id)
-
     def grant_privileges(self, entity_model: StreamlitEntityModel):
         if not entity_model.grants:
             return