Snow 1922856 refactor and unify uploading artifact to stage (#2133)

* notebooks

* fix stage sync utils

* Fix notebooks

* refactor; add Streamlit

* refactor notebook

* Fix streamlit tests

* Add support for snow:// prefix in StagePathParts

* Fix notebook tests

* snowpark

* test streamlit

* snowpark test

* update release notes

* Review fixes

* fix dcm tests

* fix dcm tests

Author: sfc-gh-pczajka

RELEASE-NOTES.md

@@ -19,6 +19,8 @@
 ## Deprecations
 
 ## New additions
+* Added `--prune` flag to `deploy` commands, which removes files that exist in the stage,
+  but not in the local filesystem.
 
 ## Fixes and improvements
 

src/snowflake/cli/_plugins/notebook/commands.py

@@ -23,6 +23,7 @@
 from snowflake.cli.api.cli_global_context import get_cli_context
 from snowflake.cli.api.commands.decorators import with_project_definition
 from snowflake.cli.api.commands.flags import (
+    PruneOption,
     ReplaceOption,
     entity_argument,
     identifier_argument,
@@ -108,6 +109,7 @@ def deploy(
         help="Replace notebook object 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(),
     **options,
 ) -> CommandResult:
     """Uploads a notebook and required files to a stage and creates a Snowflake notebook."""
@@ -132,6 +134,7 @@ def deploy(
         notebook.entity_id,
         EntityActions.DEPLOY,
         replace=replace,
+        prune=prune,
     )
     return MessageResult(
         f"Notebook successfully deployed and available under {notebook_url}"

src/snowflake/cli/_plugins/notebook/notebook_entity.py

@@ -4,9 +4,8 @@
 from snowflake.cli._plugins.connection.util import make_snowsight_url
 from snowflake.cli._plugins.notebook.notebook_entity_model import NotebookEntityModel
 from snowflake.cli._plugins.notebook.notebook_project_paths import NotebookProjectPaths
-from snowflake.cli._plugins.stage.manager import StageManager
 from snowflake.cli._plugins.workspace.context import ActionContext
-from snowflake.cli.api.artifacts.utils import bundle_artifacts
+from snowflake.cli.api.artifacts.upload import sync_artifacts_with_stage
 from snowflake.cli.api.cli_global_context import get_cli_context
 from snowflake.cli.api.console.console import cli_console
 from snowflake.cli.api.entities.common import EntityBase
@@ -22,12 +21,15 @@ class NotebookEntity(EntityBase[NotebookEntityModel]):
     A notebook.
     """
 
+    @property
+    def _stage_path_from_model(self) -> str:
+        if self.model.stage_path is None:
+            return f"{_DEFAULT_NOTEBOOK_STAGE_NAME}/{self.fqn.name}"
+        return self.model.stage_path
+
     @functools.cached_property
     def _stage_path(self) -> StagePath:
-        stage_path = self.model.stage_path
-        if stage_path is None:
-            stage_path = f"{_DEFAULT_NOTEBOOK_STAGE_NAME}/{self.fqn.name}"
-        return StagePath.from_stage_str(stage_path)
+        return StagePath.from_stage_str(self._stage_path_from_model)
 
     @functools.cached_property
     def _project_paths(self):
@@ -41,26 +43,6 @@ def _object_exists(self) -> bool:
         except ProgrammingError:
             return False
 
-    def _upload_artifacts(self):
-        stage_fqn = self._stage_path.stage_fqn
-        stage_manager = StageManager()
-        cli_console.step(f"Creating stage {stage_fqn} if not exists")
-        stage_manager.create(fqn=stage_fqn)
-
-        cli_console.step("Uploading artifacts")
-
-        # creating bundle map to handle glob patterns logic
-        bundle_map = bundle_artifacts(self._project_paths, self.model.artifacts)
-        for absolute_src, absolute_dest in bundle_map.all_mappings(
-            absolute=True, expand_directories=True
-        ):
-            artifact_stage_path = self._stage_path / (
-                absolute_dest.relative_to(self._project_paths.bundle_root).parent
-            )
-            stage_manager.put(
-                local_path=absolute_src, stage_path=artifact_stage_path, overwrite=True
-            )
-
     def get_create_sql(self, replace: bool) -> str:
         main_file_stage_path = self._stage_path / (
             self.model.notebook_file.absolute().relative_to(
@@ -99,6 +81,7 @@ def action_deploy(
         self,
         action_ctx: ActionContext,
         replace: bool,
+        prune: bool,
         *args,
         **kwargs,
     ) -> str:
@@ -108,7 +91,13 @@ def action_deploy(
                     f"Notebook {self.fqn.name} already exists. Consider using --replace."
                 )
         with cli_console.phase(f"Uploading artifacts to {self._stage_path}"):
-            self._upload_artifacts()
+            sync_artifacts_with_stage(
+                project_paths=self._project_paths,
+                stage_root=self._stage_path_from_model,
+                prune=prune,
+                artifacts=self.model.artifacts,
+            )
+
         with cli_console.phase(f"Creating notebook {self.fqn}"):
             return self.action_create(replace=replace)
 

src/snowflake/cli/_plugins/project/commands.py

@@ -22,8 +22,7 @@
 from snowflake.cli._plugins.project.project_entity_model import (
     ProjectEntityModel,
 )
-from snowflake.cli._plugins.stage.manager import StageManager
-from snowflake.cli.api.artifacts.upload import put_files
+from snowflake.cli.api.artifacts.upload import sync_artifacts_with_stage
 from snowflake.cli.api.cli_global_context import get_cli_context
 from snowflake.cli.api.commands.decorators import with_project_definition
 from snowflake.cli.api.commands.flags import (
@@ -118,20 +117,16 @@ def create_version(
 
     # Sync state
     with cli_console.phase("Syncing project state"):
-        stage_name = FQN.from_stage(project.stage)
-        sm = StageManager()
-
-        cli_console.step(f"Creating stage {stage_name}")
-        sm.create(fqn=stage_name)
-
-        put_files(
+        sync_artifacts_with_stage(
             project_paths=ProjectPaths(project_root=cli_context.project_root),
             stage_root=project.stage,
             artifacts=project.artifacts,
         )
 
     # Create project and version
     with cli_console.phase("Creating project and version"):
+        stage_name = FQN.from_stage(project.stage)
+
         pm = ProjectManager()
         cli_console.step(f"Creating project {project.fqn}")
         pm.create(project_name=project.fqn)

src/snowflake/cli/_plugins/snowpark/commands.py

@@ -71,6 +71,7 @@
 )
 from snowflake.cli.api.commands.flags import (
     ForceReplaceOption,
+    PruneOption,
     ReplaceOption,
     execution_identifier_argument,
     identifier_argument,
@@ -133,6 +134,9 @@ def deploy(
         "overwrites existing files, but does not remove any files already on the stage."
     ),
     force_replace: bool = ForceReplaceOption(),
+    prune: bool = PruneOption(
+        help="Remove contents of the stage before uploading artifacts."
+    ),
     **options,
 ) -> CommandResult:
     """
@@ -174,7 +178,7 @@ def deploy(
             snowpark_entities=snowpark_entities,
         )
 
-        create_stages_and_upload_artifacts(stages_to_artifact_map)
+        create_stages_and_upload_artifacts(stages_to_artifact_map, prune=prune)
 
     # Create snowpark entities
     with cli_console.phase("Creating Snowpark entities"):
@@ -242,8 +246,17 @@ def build_artifacts_mappings(
     return entities_to_imports_map, stages_to_artifact_map
 
 
-def create_stages_and_upload_artifacts(stages_to_artifact_map: StageToArtifactMapping):
+def create_stages_and_upload_artifacts(
+    stages_to_artifact_map: StageToArtifactMapping, prune: bool
+):
     stage_manager = StageManager()
+    if prune:
+        # snowflake.cli._plugins.snowpark.snowpark_project_paths.Artifact class assumes that "stage"
+        # is a stage object, not path on stage - whole stage is managed by snowpark - it can be removed
+        for stage in stages_to_artifact_map.keys():
+            cli_console.step(f"Removing contents of stage {stage}")
+            stage_manager.remove(stage, path="")
+
     for stage, artifacts in stages_to_artifact_map.items():
         cli_console.step(f"Creating (if not exists) stage: {stage}")
         stage = FQN.from_stage(stage).using_context()

src/snowflake/cli/_plugins/stage/manager.py

@@ -65,7 +65,7 @@
 
 # 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,15 +119,23 @@ def strip_stage_prefix(self, path: str):
         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):
     """
     For path like @db.schema.stage/dir the values will be:
         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 @@ def __init__(self, stage_path: str):
             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 @@ def schema(self) -> str | None:
         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

src/snowflake/cli/_plugins/streamlit/commands.py

@@ -37,6 +37,7 @@
     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_:

src/snowflake/cli/_plugins/streamlit/manager.py

@@ -31,7 +31,7 @@
 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 @@ def share(self, streamlit_name: FQN, to_role: str) -> SnowflakeCursor:
             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 @@ def deploy(
         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 @@ def deploy(
             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 @@ def deploy(
             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(