snowflake-cli-labs 2.8.0rc1__py3-none-any.whl → 3.0.0rc1__py3-none-any.whl

snowflake/cli/{plugins → _plugins}/stage/manager.py

@@ -27,11 +27,12 @@ from textwrap import dedent
 from typing import Dict, List, Optional, Union
 
 from click import ClickException
-from snowflake.cli.api.commands.flags import (
+from snowflake.cli._plugins.snowpark.package_utils import parse_requirements
+from snowflake.cli.api.commands.common import (
     OnErrorType,
     Variable,
-    parse_key_value_variables,
 )
+from snowflake.cli.api.commands.utils import parse_key_value_variables
 from snowflake.cli.api.console import cli_console
 from snowflake.cli.api.constants import PYTHON_3_12
 from snowflake.cli.api.identifiers import FQN
@@ -39,7 +40,6 @@ from snowflake.cli.api.project.util import to_string_literal
 from snowflake.cli.api.secure_path import SecurePath
 from snowflake.cli.api.sql_execution import SqlExecutionMixin
 from snowflake.cli.api.utils.path_utils import path_resolver
-from snowflake.cli.plugins.snowpark.package_utils import parse_requirements
 from snowflake.connector import DictCursor, ProgrammingError
 from snowflake.connector.cursor import SnowflakeCursor
 
@@ -65,14 +65,21 @@ class StagePathParts:
     stage_name: str
     is_directory: bool
 
-    @staticmethod
-    def get_directory(stage_path: str) -> str:
+    @classmethod
+    def get_directory(cls, stage_path: str) -> str:
         return "/".join(Path(stage_path).parts[1:])
 
     @property
     def path(self) -> str:
         raise NotImplementedError
 
+    @property
+    def full_path(self) -> str:
+        raise NotImplementedError
+
+    def replace_stage_prefix(self, file_path: str) -> str:
+        raise NotImplementedError
+
     def add_stage_prefix(self, file_path: str) -> str:
         raise NotImplementedError
 
@@ -112,24 +119,27 @@ class DefaultStagePathParts(StagePathParts):
         self.directory = self.get_directory(stage_path)
         self.stage = StageManager.get_stage_from_path(stage_path)
         stage_name = self.stage.split(".")[-1]
-        if stage_name.startswith("@"):
-            stage_name = stage_name[1:]
+        stage_name = stage_name[1:] if stage_name.startswith("@") else stage_name
         self.stage_name = stage_name
         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}"
-        )
+        return f"{self.stage_name.rstrip('/')}/{self.directory}"
 
-    def add_stage_prefix(self, file_path: str) -> str:
+    @property
+    def full_path(self) -> str:
+        return f"{self.stage.rstrip('/')}/{self.directory}"
+
+    def replace_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 add_stage_prefix(self, file_path: str) -> str:
+        stage = self.stage.rstrip("/")
+        return f"{stage}/{file_path.lstrip('/')}"
+
     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[1 + stage_path_length : -1])
@@ -146,14 +156,29 @@ class UserStagePathParts(StagePathParts):
 
     def __init__(self, stage_path: str):
         self.directory = self.get_directory(stage_path)
-        self.stage = "@~"
-        self.stage_name = "@~"
+        self.stage = USER_STAGE_PREFIX
+        self.stage_name = USER_STAGE_PREFIX
         self.is_directory = True if stage_path.endswith("/") else False
 
+    @classmethod
+    def get_directory(cls, stage_path: str) -> str:
+        if Path(stage_path).parts[0] == USER_STAGE_PREFIX:
+            return super().get_directory(stage_path)
+        return stage_path
+
     @property
     def path(self) -> str:
         return f"{self.directory}"
 
+    @property
+    def full_path(self) -> str:
+        return f"{self.stage}/{self.directory}"
+
+    def replace_stage_prefix(self, file_path: str) -> str:
+        if Path(file_path).parts[0] == self.stage_name:
+            return file_path
+        return f"{self.stage}/{file_path}"
+
     def add_stage_prefix(self, file_path: str) -> str:
         return f"{self.stage}/{file_path}"
 
@@ -168,7 +193,9 @@ class StageManager(SqlExecutionMixin):
         self._python_exe_procedure = None
 
     @staticmethod
-    def get_standard_stage_prefix(name: str) -> str:
+    def get_standard_stage_prefix(name: str | FQN) -> str:
+        if isinstance(name, FQN):
+            name = name.identifier
         # Handle embedded stages
         if name.startswith("snow://") or name.startswith("@"):
             return name
@@ -239,7 +266,7 @@ class StageManager(SqlExecutionMixin):
             self._assure_is_existing_directory(dest_directory)
 
             result = self._execute_query(
-                f"get {self.quote_stage_name(stage_path_parts.add_stage_prefix(file_path))} {self._to_uri(f'{dest_directory}/')} parallel={parallel}"
+                f"get {self.quote_stage_name(stage_path_parts.replace_stage_prefix(file_path))} {self._to_uri(f'{dest_directory}/')} parallel={parallel}"
             )
             results.append(result)
 
@@ -300,8 +327,8 @@ class StageManager(SqlExecutionMixin):
             quoted_stage_name = self.quote_stage_name(f"{stage_name}{path}")
             return self._execute_query(f"remove {quoted_stage_name}")
 
-    def create(self, stage_name: str, comment: Optional[str] = None) -> SnowflakeCursor:
-        query = f"create stage if not exists {stage_name}"
+    def create(self, fqn: FQN, comment: Optional[str] = None) -> SnowflakeCursor:
+        query = f"create stage if not exists {fqn.sql_identifier}"
         if comment:
             query += f" comment='{comment}'"
         return self._execute_query(query)
@@ -319,8 +346,14 @@ class StageManager(SqlExecutionMixin):
         stage_path_parts = self._stage_path_part_factory(stage_path)
         all_files_list = self._get_files_list_from_stage(stage_path_parts)
 
+        all_files_with_stage_name_prefix = [
+            stage_path_parts.get_directory(file) for file in all_files_list
+        ]
+
         # filter files from stage if match stage_path pattern
-        filtered_file_list = self._filter_files_list(stage_path_parts, all_files_list)
+        filtered_file_list = self._filter_files_list(
+            stage_path_parts, all_files_with_stage_name_prefix
+        )
 
         if not filtered_file_list:
             raise ClickException(f"No files matched pattern '{stage_path}'")
@@ -376,7 +409,7 @@ class StageManager(SqlExecutionMixin):
         if not stage_path_parts.directory:
             return self._filter_supported_files(files_on_stage)
 
-        stage_path = stage_path_parts.path.lower()
+        stage_path = stage_path_parts.directory
 
         # Exact file path was provided if stage_path in file list
         if stage_path in files_on_stage:

snowflake/cli/{plugins → _plugins}/stage/plugin_spec.py

@@ -12,13 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from snowflake.cli._plugins.stage import commands
 from snowflake.cli.api.plugins.command import (
     SNOWCLI_ROOT_COMMAND_PATH,
     CommandSpec,
     CommandType,
     plugin_hook_impl,
 )
-from snowflake.cli.plugins.stage import commands
 
 
 @plugin_hook_impl

snowflake/cli/_plugins/stage/utils.py

@@ -0,0 +1,54 @@
+from typing import Optional
+
+from snowflake.cli._plugins.nativeapp.artifacts import BundleMap
+from snowflake.cli._plugins.stage.diff import (
+    DiffResult,
+    _to_diff_line,
+    _to_src_dest_pair,
+)
+from snowflake.cli.api.console import cli_console as cc
+
+
+def print_diff_to_console(
+    diff: DiffResult,
+    bundle_map: Optional[BundleMap] = None,
+):
+    if not diff.has_changes():
+        cc.message("Your stage is up-to-date with your local deploy root.")
+        return
+
+    blank_line_needed = False
+    if diff.only_local or diff.different:
+        cc.message("Local changes to be deployed:")
+        messages_to_output = []
+        for p in diff.different:
+            src_dest_pair = _to_src_dest_pair(p, bundle_map)
+            messages_to_output.append(
+                (
+                    src_dest_pair,
+                    _to_diff_line("modified", src_dest_pair[0], src_dest_pair[1]),
+                )
+            )
+        for p in diff.only_local:
+            src_dest_pair = _to_src_dest_pair(p, bundle_map)
+            messages_to_output.append(
+                (
+                    src_dest_pair,
+                    _to_diff_line("added", src_dest_pair[0], src_dest_pair[1]),
+                )
+            )
+
+        with cc.indented():
+            for key, message in sorted(messages_to_output, key=lambda pair: pair[0]):
+                cc.message(message)
+
+        blank_line_needed = True
+
+    if diff.only_on_stage:
+        if blank_line_needed:
+            cc.message("")
+        cc.message(f"Deleted paths to be removed from your stage:")
+        with cc.indented():
+            for p in sorted(diff.only_on_stage):
+                diff_line = _to_diff_line("deleted", src=None, dest=str(p))
+                cc.message(diff_line)

snowflake/cli/{plugins → _plugins}/streamlit/commands.py

@@ -16,32 +16,42 @@ from __future__ import annotations
 
 import logging
 from pathlib import Path
+from typing import Dict
 
 import click
 import typer
-from click import ClickException
-from snowflake.cli.api.cli_global_context import cli_context
+from click import ClickException, UsageError
+from snowflake.cli._plugins.object.command_aliases import (
+    add_object_command_aliases,
+    scope_option,
+)
+from snowflake.cli._plugins.streamlit.manager import StreamlitManager
+from snowflake.cli.api.cli_global_context import get_cli_context
 from snowflake.cli.api.commands.decorators import (
     with_experimental_behaviour,
     with_project_definition,
 )
-from snowflake.cli.api.commands.flags import ReplaceOption, like_option
-from snowflake.cli.api.commands.project_initialisation import add_init_command
+from snowflake.cli.api.commands.flags import (
+    ReplaceOption,
+    entity_argument,
+    identifier_argument,
+    like_option,
+)
 from snowflake.cli.api.commands.snow_typer import SnowTyperFactory
 from snowflake.cli.api.constants import ObjectType
+from snowflake.cli.api.exceptions import NoProjectDefinitionError
 from snowflake.cli.api.identifiers import FQN
 from snowflake.cli.api.output.types import (
     CommandResult,
     MessageResult,
     SingleQueryResult,
 )
-from snowflake.cli.api.project.project_verification import assert_project_type
-from snowflake.cli.api.project.schemas.streamlit.streamlit import Streamlit
-from snowflake.cli.plugins.object.command_aliases import (
-    add_object_command_aliases,
-    scope_option,
+from snowflake.cli.api.project.definition_conversion import (
+    convert_project_definition_to_v2,
+)
+from snowflake.cli.api.project.schemas.entities.streamlit_entity_model import (
+    StreamlitEntityModel,
 )
-from snowflake.cli.plugins.streamlit.manager import StreamlitManager
 
 app = SnowTyperFactory(
     name="streamlit",
@@ -49,19 +59,8 @@ app = SnowTyperFactory(
 )
 log = logging.getLogger(__name__)
 
-
-class IdentifierType(click.ParamType):
-    name = "TEXT"
-
-    def convert(self, value, param, ctx):
-        return FQN.from_string(value)
-
-
-StreamlitNameArgument = typer.Argument(
-    ...,
-    help="Name of the Streamlit app.",
-    show_default=False,
-    click_type=IdentifierType(),
+StreamlitNameArgument = identifier_argument(
+    sf_object="Streamlit app", example="my_streamlit"
 )
 OpenOption = typer.Option(
     False,
@@ -70,12 +69,6 @@ OpenOption = typer.Option(
     is_flag=True,
 )
 
-add_init_command(
-    app,
-    project_type="Streamlit",
-    template="default_streamlit",
-    help_message="Name of the Streamlit app project directory you want to create. Defaults to `example_streamlit`.",
-)
 
 add_object_command_aliases(
     app=app,
@@ -92,7 +85,9 @@ add_object_command_aliases(
 def streamlit_share(
     name: FQN = StreamlitNameArgument,
     to_role: str = typer.Argument(
-        ..., help="Role with which to share the Streamlit app."
+        ...,
+        help="Role with which to share the Streamlit app.",
+        show_default=False,
     ),
     **options,
 ) -> CommandResult:
@@ -125,48 +120,50 @@ def streamlit_deploy(
     replace: bool = ReplaceOption(
         help="Replace the Streamlit app if it already exists."
     ),
+    entity_id: str = entity_argument("streamlit"),
     open_: bool = OpenOption,
     **options,
 ) -> CommandResult:
     """
     Deploys a Streamlit app defined in the project definition file (snowflake.yml). By default, the command uploads
     environment.yml and any other pages or folders, if present. If you don’t specify a stage name, the `streamlit`
-    stage is used. If the specified stage does not exist, the command creates it.
+    stage is used. If the specified stage does not exist, the command creates it. If multiple Streamlits are defined
+    in snowflake.yml and no entity_id is provided then command will raise an error.
     """
 
-    assert_project_type("streamlit")
-
-    streamlit: Streamlit = cli_context.project_definition.streamlit
-    if not streamlit:
-        return MessageResult("No streamlit were specified in project definition.")
-
-    environment_file = streamlit.env_file
-    if environment_file and not Path(environment_file).exists():
-        raise ClickException(f"Provided file {environment_file} does not exist")
-    elif environment_file is None:
-        environment_file = "environment.yml"
-
-    pages_dir = streamlit.pages_dir
-    if pages_dir and not Path(pages_dir).exists():
-        raise ClickException(f"Provided file {pages_dir} does not exist")
-    elif pages_dir is None:
-        pages_dir = "pages"
-
-    streamlit_id = FQN.from_identifier_model(streamlit).using_context()
-
-    url = StreamlitManager().deploy(
-        streamlit_id=streamlit_id,
-        environment_file=Path(environment_file),
-        pages_dir=Path(pages_dir),
-        stage_name=streamlit.stage,
-        main_file=Path(streamlit.main_file),
-        replace=replace,
-        query_warehouse=streamlit.query_warehouse,
-        additional_source_files=streamlit.additional_source_files,
-        title=streamlit.title,
-        **options,
+    cli_context = get_cli_context()
+    pd = cli_context.project_definition
+    if not pd.meets_version_requirement("2"):
+        if not pd.streamlit:
+            raise NoProjectDefinitionError(
+                project_type="streamlit", project_root=cli_context.project_root
+            )
+        pd = convert_project_definition_to_v2(pd)
+
+    streamlits: Dict[str, StreamlitEntityModel] = pd.get_entities_by_type(
+        entity_type="streamlit"
     )
 
+    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, replace=replace)
+
     if open_:
         typer.launch(url)
 

snowflake/cli/{plugins → _plugins}/streamlit/manager.py

@@ -15,22 +15,25 @@
 from __future__ import annotations
 
 import logging
-import os
 from pathlib import Path
 from typing import List, Optional
 
+from snowflake.cli._plugins.connection.util import (
+    MissingConnectionAccountError,
+    MissingConnectionRegionError,
+    make_snowsight_url,
+)
+from snowflake.cli._plugins.stage.manager import StageManager
 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.sql_execution import SqlExecutionMixin
-from snowflake.cli.plugins.connection.util import (
-    MissingConnectionAccountError,
-    MissingConnectionRegionError,
-    make_snowsight_url,
+from snowflake.cli.api.project.schemas.entities.streamlit_entity_model import (
+    StreamlitEntityModel,
 )
-from snowflake.cli.plugins.stage.manager import StageManager
+from snowflake.cli.api.sql_execution import SqlExecutionMixin
 from snowflake.connector.cursor import SnowflakeCursor
 from snowflake.connector.errors import ProgrammingError
 
@@ -46,39 +49,35 @@ class StreamlitManager(SqlExecutionMixin):
     def _put_streamlit_files(
         self,
         root_location: str,
-        main_file: Path,
-        environment_file: Optional[Path],
-        pages_dir: Optional[Path],
-        additional_source_files: Optional[List[Path]],
+        artifacts: Optional[List[Path]] = None,
     ):
+        cli_console.step(f"Deploying files to {root_location}")
+        if not artifacts:
+            return
         stage_manager = StageManager()
-
-        stage_manager.put(main_file, root_location, 4, True)
-
-        if environment_file and environment_file.exists():
-            stage_manager.put(environment_file, root_location, 4, True)
-
-        if pages_dir and pages_dir.exists():
-            stage_manager.put(pages_dir / "*.py", f"{root_location}/pages", 4, True)
-
-        if additional_source_files:
-            for file in additional_source_files:
-                if os.sep in str(file):
-                    destination = f"{root_location}/{str(file.parent)}"
-                else:
-                    destination = root_location
-                stage_manager.put(file, destination, 4, True)
+        for file in artifacts:
+            if file.is_dir():
+                if not any(file.iterdir()):
+                    cli_console.warning(f"Skipping empty directory: {file}")
+                    continue
+
+                stage_manager.put(
+                    f"{file.joinpath('*')}", f"{root_location}/{file}", 4, True
+                )
+            elif len(file.parts) > 1:
+                stage_manager.put(file, f"{root_location}/{file.parent}", 4, True)
+            else:
+                stage_manager.put(file, root_location, 4, True)
 
     def _create_streamlit(
         self,
-        streamlit_id: FQN,
-        main_file: Path,
+        streamlit: StreamlitEntityModel,
         replace: Optional[bool] = None,
         experimental: Optional[bool] = None,
-        query_warehouse: Optional[str] = None,
         from_stage_name: Optional[str] = None,
-        title: 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}")
@@ -96,28 +95,24 @@ class StreamlitManager(SqlExecutionMixin):
         if from_stage_name:
             query.append(f"ROOT_LOCATION = '{from_stage_name}'")
 
-        query.append(f"MAIN_FILE = '{main_file.name}'")
+        query.append(f"MAIN_FILE = '{streamlit.main_file}'")
 
-        if query_warehouse:
-            query.append(f"QUERY_WAREHOUSE = {query_warehouse}")
-        if title:
-            query.append(f"TITLE = '{title}'")
+        if streamlit.query_warehouse:
+            query.append(f"QUERY_WAREHOUSE = {streamlit.query_warehouse}")
+        if streamlit.title:
+            query.append(f"TITLE = '{streamlit.title}'")
+
+        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_id: FQN,
-        main_file: Path,
-        environment_file: Optional[Path] = None,
-        pages_dir: Optional[Path] = None,
-        stage_name: Optional[str] = None,
-        query_warehouse: Optional[str] = None,
-        replace: Optional[bool] = False,
-        additional_source_files: Optional[List[Path]] = None,
-        title: Optional[str] = None,
-        **options,
-    ):
+    def deploy(self, streamlit: StreamlitEntityModel, replace: bool = False):
+        streamlit_id = streamlit.fqn.using_connection(self._conn)
+
         # 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
@@ -134,12 +129,9 @@ class StreamlitManager(SqlExecutionMixin):
             # TODO: Support from_stage
             # from_stage_stmt = f"FROM_STAGE = '{stage_name}'" if stage_name else ""
             self._create_streamlit(
-                streamlit_id,
-                main_file,
+                streamlit=streamlit,
                 replace=replace,
-                query_warehouse=query_warehouse,
                 experimental=True,
-                title=title,
             )
             try:
                 if use_versioned_stage:
@@ -169,10 +161,7 @@ class StreamlitManager(SqlExecutionMixin):
 
             self._put_streamlit_files(
                 root_location,
-                main_file,
-                environment_file,
-                pages_dir,
-                additional_source_files,
+                streamlit.artifacts,
             )
         else:
             """
@@ -182,31 +171,23 @@ class StreamlitManager(SqlExecutionMixin):
             """
             stage_manager = StageManager()
 
-            stage_name = stage_name or "streamlit"
+            stage_name = streamlit.stage or "streamlit"
             stage_name = FQN.from_string(stage_name).using_connection(self._conn)
 
-            stage_manager.create(stage_name=stage_name)
+            cli_console.step(f"Creating {stage_name} stage")
+            stage_manager.create(fqn=stage_name)
 
             root_location = stage_manager.get_standard_stage_prefix(
                 f"{stage_name}/{streamlit_name_for_root_location}"
             )
 
-            self._put_streamlit_files(
-                root_location,
-                main_file,
-                environment_file,
-                pages_dir,
-                additional_source_files,
-            )
+            self._put_streamlit_files(root_location, streamlit.artifacts)
 
             self._create_streamlit(
-                streamlit_id,
-                main_file,
+                streamlit=streamlit,
                 replace=replace,
-                query_warehouse=query_warehouse,
                 from_stage_name=root_location,
                 experimental=False,
-                title=title,
             )
 
         return self.get_url(streamlit_name=streamlit_id)

snowflake/cli/_plugins/streamlit/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._plugins.streamlit import commands
+from snowflake.cli.api.plugins.command import (
+    SNOWCLI_ROOT_COMMAND_PATH,
+    CommandSpec,
+    CommandType,
+    plugin_hook_impl,
+)
+
+
+@plugin_hook_impl
+def command_spec():
+    return CommandSpec(
+        parent_command_path=SNOWCLI_ROOT_COMMAND_PATH,
+        command_type=CommandType.COMMAND_GROUP,
+        typer_instance=commands.app.create_instance(),
+    )

snowflake/cli/_plugins/workspace/action_context.py

@@ -0,0 +1,17 @@
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Optional
+
+from snowflake.cli.api.console.abc import AbstractConsole
+
+
+@dataclass
+class ActionContext:
+    """
+    An object that is passed to each action when called by WorkspaceManager
+    """
+
+    console: AbstractConsole
+    project_root: Path
+    default_role: str
+    default_warehouse: Optional[str]