Jw/snow 2306289 support local output path (#2591)

* feat: [SNOW-2306289] support copying output files to local path

* refactor: [SNOW-2306289] extract logic related to generating identifiers for temp stages to FQN.related_to_resource

* feat: [SNOW-2306289] add more tests

* refactor: [SNOW-2306289] apply review suggestions

Author: sfc-gh-jwilkowski

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

@@ -27,7 +27,6 @@
 from snowflake.cli.api.constants import DEFAULT_SIZE_LIMIT_MB, ObjectType
 from snowflake.cli.api.exceptions import CliError
 from snowflake.cli.api.identifiers import FQN
-from snowflake.cli.api.project.util import unquote_identifier
 from snowflake.cli.api.secure_path import SecurePath
 from snowflake.cli.api.sql_execution import SqlExecutionMixin
 from snowflake.connector.cursor import SnowflakeCursor
@@ -104,10 +103,9 @@ def deploy(
 
         with cli_console.phase("Creating temporary stage"):
             stage_manager = StageManager()
-            unquoted_name = unquote_identifier(fqn.name)
-            stage_fqn = FQN.from_string(f"DBT_{unquoted_name}_STAGE").using_context()
-            stage_name = stage_manager.get_standard_stage_prefix(stage_fqn)
+            stage_fqn = FQN.from_resource(ObjectType.DBT_PROJECT, fqn, "STAGE")
             stage_manager.create(stage_fqn, temporary=True)
+            stage_name = stage_manager.get_standard_stage_prefix(stage_fqn)
 
         with cli_console.phase("Copying project files to stage"):
             with TemporaryDirectory() as tmp:

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

@@ -142,7 +142,7 @@ def plan(
     variables: Optional[List[str]] = variables_flag,
     configuration: Optional[str] = configuration_flag,
     output_path: Optional[str] = output_path_option(
-        help="Stage path where the deployment plan output will be stored."
+        help="Path where the deployment plan output will be stored. Can be a stage path (starting with '@') or a local directory path."
     ),
     **options,
 ):

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

@@ -11,30 +11,72 @@
 # 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.
-import time
+from contextlib import contextmanager, nullcontext
 from pathlib import Path
-from typing import List
+from typing import Generator, List
 
 import yaml
 from snowflake.cli._plugins.stage.manager import StageManager
 from snowflake.cli.api.artifacts.upload import sync_artifacts_with_stage
 from snowflake.cli.api.commands.utils import parse_key_value_variables
 from snowflake.cli.api.console.console import cli_console
-from snowflake.cli.api.constants import DEFAULT_SIZE_LIMIT_MB, PatternMatchingType
+from snowflake.cli.api.constants import (
+    DEFAULT_SIZE_LIMIT_MB,
+    ObjectType,
+    PatternMatchingType,
+)
 from snowflake.cli.api.exceptions import CliError
 from snowflake.cli.api.identifiers import FQN
 from snowflake.cli.api.project.project_paths import ProjectPaths
 from snowflake.cli.api.project.schemas.entities.common import PathMapping
-from snowflake.cli.api.project.util import unquote_identifier
 from snowflake.cli.api.secure_path import SecurePath
 from snowflake.cli.api.sql_execution import SqlExecutionMixin
 from snowflake.cli.api.stage_path import StagePath
+from snowflake.cli.api.utils.path_utils import is_stage_path
 
 MANIFEST_FILE_NAME = "manifest.yml"
 DCM_PROJECT_TYPE = "dcm_project"
 
 
 class DCMProjectManager(SqlExecutionMixin):
+    @contextmanager
+    def _collect_output(
+        self, project_identifier: FQN, output_path: str
+    ) -> Generator[str, None, None]:
+        """
+        Context manager for handling output path - creates temporary stage for local paths,
+        downloads files after execution, and ensures proper cleanup.
+
+        Args:
+            project_identifier: The DCM project identifier
+            output_path: Either a stage path (@stage/path) or local directory path
+
+        Yields:
+            str: The effective output path to use in the DCM command
+        """
+        temp_stage_for_local_output = None
+        stage_manager = StageManager()
+
+        if should_download_files := not is_stage_path(output_path):
+            temp_stage_fqn = FQN.from_resource(
+                ObjectType.DCM_PROJECT, project_identifier, "OUTPUT_TMP_STAGE"
+            )
+            stage_manager.create(temp_stage_fqn, temporary=True)
+            effective_output_path = StagePath.from_stage_str(temp_stage_fqn.identifier)
+            temp_stage_for_local_output = (temp_stage_fqn.identifier, Path(output_path))
+        else:
+            effective_output_path = StagePath.from_stage_str(output_path)
+
+        yield effective_output_path.absolute_path()
+
+        if should_download_files:
+            assert temp_stage_for_local_output is not None
+            stage_path, local_path = temp_stage_for_local_output
+            stage_manager.get_recursive(stage_path=stage_path, dest_path=local_path)
+            cli_console.step(f"Plan output saved to: {local_path.resolve()}")
+        else:
+            cli_console.step(f"Plan output saved to: {output_path}")
+
     def execute(
         self,
         project_identifier: FQN,
@@ -45,28 +87,31 @@ def execute(
         alias: str | None = None,
         output_path: str | None = None,
     ):
+        with self._collect_output(project_identifier, output_path) if (
+            output_path and dry_run
+        ) else nullcontext() as output_stage:
+            query = f"EXECUTE DCM PROJECT {project_identifier.sql_identifier}"
+            if dry_run:
+                query += " PLAN"
+            else:
+                query += " DEPLOY"
+                if alias:
+                    query += f' AS "{alias}"'
+            if configuration or variables:
+                query += f" USING"
+            if configuration:
+                query += f" CONFIGURATION {configuration}"
+            if variables:
+                query += StageManager.parse_execute_variables(
+                    parse_key_value_variables(variables)
+                ).removeprefix(" using")
+            stage_path = StagePath.from_stage_str(from_stage)
+            query += f" FROM {stage_path.absolute_path()}"
+            if output_stage is not None:
+                query += f" OUTPUT_PATH {output_stage}"
+            result = self.execute_query(query=query)
 
-        query = f"EXECUTE DCM PROJECT {project_identifier.sql_identifier}"
-        if dry_run:
-            query += " PLAN"
-        else:
-            query += " DEPLOY"
-            if alias:
-                query += f' AS "{alias}"'
-        if configuration or variables:
-            query += f" USING"
-        if configuration:
-            query += f" CONFIGURATION {configuration}"
-        if variables:
-            query += StageManager.parse_execute_variables(
-                parse_key_value_variables(variables)
-            ).removeprefix(" using")
-        stage_path = StagePath.from_stage_str(from_stage)
-        query += f" FROM {stage_path.absolute_path()}"
-        if output_path:
-            output_stage_path = StagePath.from_stage_str(output_path)
-            query += f" OUTPUT_PATH {output_stage_path.absolute_path()}"
-        return self.execute_query(query=query)
+        return result
 
     def create(self, project_identifier: FQN) -> None:
         query = f"CREATE DCM PROJECT {project_identifier.sql_identifier}"
@@ -114,10 +159,9 @@ def sync_local_files(project_identifier: FQN) -> str:
                 definitions.append(MANIFEST_FILE_NAME)
 
         with cli_console.phase(f"Uploading definition files"):
-            unquoted_name = unquote_identifier(project_identifier.name)
-            stage_fqn = FQN.from_string(
-                f"DCM_{unquoted_name}_{int(time.time())}_TMP_STAGE"
-            ).using_context()
+            stage_fqn = FQN.from_resource(
+                ObjectType.DCM_PROJECT, project_identifier, "_TMP_STAGE"
+            )
             sync_artifacts_with_stage(
                 project_paths=ProjectPaths(project_root=Path.cwd()),
                 stage_root=stage_fqn.identifier,

src/snowflake/cli/api/identifiers.py

@@ -15,14 +15,20 @@
 from __future__ import annotations
 
 import re
+import time
 from pathlib import Path
 
 from click import ClickException
+from snowflake.cli.api.constants import ObjectType
 from snowflake.cli.api.exceptions import FQNInconsistencyError, FQNNameError
 from snowflake.cli.api.project.schemas.v1.identifier_model import (
     ObjectIdentifierBaseModel,
 )
-from snowflake.cli.api.project.util import VALID_IDENTIFIER_REGEX, identifier_for_url
+from snowflake.cli.api.project.util import (
+    VALID_IDENTIFIER_REGEX,
+    identifier_for_url,
+    unquote_identifier,
+)
 
 
 class FQN:
@@ -167,6 +173,17 @@ def from_identifier_model_v2(cls, model) -> "FQN":
 
         return fqn.set_database(model.database).set_schema(model.schema_)
 
+    @classmethod
+    def from_resource(
+        cls, resource_type: ObjectType, resource_fqn: FQN, purpose: str
+    ) -> "FQN":
+        """Create an instance related to another Snowflake resource."""
+        unquoted_name = unquote_identifier(resource_fqn.name)
+        safe_cli_name = resource_type.value.cli_name.upper().replace("-", "_")
+        return cls.from_string(
+            f"{safe_cli_name}_{unquoted_name}_{int(time.time())}_{purpose.upper()}"
+        ).using_context()
+
     def set_database(self, database: str | None) -> "FQN":
         if database:
             self._database = database

tests/__snapshots__/test_help_messages.ambr

@@ -8210,8 +8210,9 @@
   |                                example: -D "<key>=<value>".                  |
   | --configuration          TEXT  Configuration of the DCM Project to use. If   |
   |                                not specified default configuration is used.  |
-  | --output-path            TEXT  Stage path where the deployment plan output   |
-  |                                will be stored.                               |
+  | --output-path            TEXT  Path where the deployment plan output will be |
+  |                                stored. Can be a stage path (starting with    |
+  |                                '@') or a local directory path.               |
   | --help           -h            Show this message and exit.                   |
   +------------------------------------------------------------------------------+
   +- Connection configuration ---------------------------------------------------+

tests/api/test_fqn.py

@@ -16,6 +16,7 @@
 from unittest.mock import MagicMock
 
 import pytest
+from snowflake.cli.api.constants import ObjectType
 from snowflake.cli.api.exceptions import FQNInconsistencyError, FQNNameError
 from snowflake.cli.api.identifiers import FQN
 from snowflake.cli.api.project.schemas.v1.streamlit.streamlit import Streamlit
@@ -197,3 +198,54 @@ def test_using_context(mock_ctx):
 def test_git_fqn():
     fqn = FQN.from_stage_path("@git_repo/branches/main/devops/")
     assert fqn.name == "git_repo"
+
+
+class TestRelatedToResource:
+    @pytest.fixture
+    def mock_time(self):
+        with mock.patch(
+            "snowflake.cli.api.identifiers.time.time", return_value=1234567890
+        ) as _fixture:
+            yield _fixture
+
+    @pytest.fixture
+    def mock_ctx(self):
+        with mock.patch(
+            "snowflake.cli.api.cli_global_context.get_cli_context"
+        ) as _fixture:
+            _fixture().connection = MagicMock(database="test_db", schema="test_schema")
+            yield _fixture
+
+    def test_basic_functionality(self, mock_ctx, mock_time):
+        resource_fqn = FQN(name="my_pipeline", database=None, schema=None)
+
+        result = FQN.from_resource(ObjectType.DBT_PROJECT, resource_fqn, "STAGE")
+
+        assert (
+            result.identifier
+            == "test_db.test_schema.DBT_PROJECT_MY_PIPELINE_1234567890_STAGE"
+        )
+
+    def test_with_quoted_identifier(self, mock_ctx, mock_time):
+        resource_fqn = FQN(name='"caseSenSITIVEnAME"', database=None, schema=None)
+
+        result = FQN.from_resource(ObjectType.DCM_PROJECT, resource_fqn, "TEMP_STAGE")
+
+        assert (
+            result.identifier
+            == "test_db.test_schema.DCM_caseSenSITIVEnAME_1234567890_TEMP_STAGE"
+        )
+
+    def test_with_fqn_resource(self, mock_ctx, mock_time):
+        mock_ctx().connection = MagicMock(
+            database="context_db", schema="context_schema"
+        )
+        resource_fqn = FQN(
+            name="resource", database="resource_db", schema="resource_schema"
+        )
+
+        result = FQN.from_resource(ObjectType.STAGE, resource_fqn, "TEST")
+
+        assert result.database == "context_db"
+        assert result.schema == "context_schema"
+        assert result.name == "STAGE_RESOURCE_1234567890_TEST"