SNOW-2162936 Add --refresh option to snow stage copy (#2448)

* SNOW-2162936 Add `--refresh` option to `snow stage copy`

* update `test_help_messages.ambr`

* Add `--enable-directory`

* Add log about refreshing stage

* update RELEASE-NOTES.md

* explicit assertion

Author: sfc-gh-daniszewski

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

@@ -117,6 +117,10 @@ def copy(
         default=False,
         help="Specifies whether Snowflake uses gzip to compress files during upload. Ignored when downloading.",
     ),
+    refresh: bool = typer.Option(
+        default=False,
+        help="Specifies whether ALTER STAGE {name} REFRESH should be executed after uploading.",
+    ),
     **options,
 ) -> CommandResult:
     """
@@ -150,6 +154,7 @@ def copy(
         parallel=parallel,
         overwrite=overwrite,
         auto_compress=auto_compress,
+        refresh=refresh,
     )
 
 
@@ -161,12 +166,19 @@ def stage_create(
         "--encryption",
         help="Type of encryption supported for all files stored on the stage.",
     ),
+    enable_directory: bool = typer.Option(
+        False,
+        "--enable-directory",
+        help="Specifies whether directory support is enabled for the stage.",
+    ),
     **options,
 ) -> CommandResult:
     """
     Creates a named stage if it does not already exist.
     """
-    cursor = StageManager().create(fqn=stage_name, encryption=encryption)
+    cursor = StageManager().create(
+        fqn=stage_name, encryption=encryption, enable_directory=enable_directory
+    )
     return SingleQueryResult(cursor)
 
 
@@ -265,6 +277,7 @@ def _put(
     parallel: int,
     overwrite: bool,
     auto_compress: bool,
+    refresh: bool,
 ):
     if recursive and not source_path.is_file():
         cursor_generator = StageManager().put_recursive(
@@ -274,7 +287,7 @@ def _put(
             parallel=parallel,
             auto_compress=auto_compress,
         )
-        return CollectionResult(cursor_generator)
+        output = CollectionResult(cursor_generator)
     else:
         cursor = StageManager().put(
             local_path=source_path.resolve(),
@@ -283,4 +296,10 @@ def _put(
             parallel=parallel,
             auto_compress=auto_compress,
         )
-        return QueryResult(cursor)
+        output = QueryResult(cursor)
+
+    if refresh:
+        StageManager().refresh(
+            StageManager.stage_path_parts_from_str(destination_path).stage_name
+        )
+    return output

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

@@ -548,11 +548,14 @@ def create(
         comment: Optional[str] = None,
         temporary: bool = False,
         encryption: InternalStageEncryptionType | None = None,
+        enable_directory: bool = False,
     ) -> SnowflakeCursor:
         temporary_str = "temporary " if temporary else ""
         query = f"create {temporary_str}stage if not exists {fqn.sql_identifier}"
         if encryption:
             query += f" encryption = (type = '{encryption.value}')"
+        if enable_directory:
+            query += f" directory = (enable = true)"
         if comment:
             query += f" comment='{comment}'"
         return self.execute_query(query)
@@ -768,6 +771,11 @@ def stage_path_parts_from_str(stage_path: str) -> StagePathParts:
             return VStagePathParts(stage_path)
         return DefaultStagePathParts(stage_path)
 
+    def refresh(self, stage_name):
+        sql = f"ALTER STAGE {stage_name} REFRESH"
+        log.info("Refreshing stage %s", stage_name)
+        return self.execute_query(sql)
+
     def _check_for_requirements_file(self, stage_path: StagePath) -> List[str]:
         """Looks for requirements.txt file on stage."""
         current_dir = stage_path.parent if stage_path.is_file() else stage_path

tests/__snapshots__/test_help_messages.ambr

@@ -17958,6 +17958,11 @@
   |                                                     downloading.             |
   |                                                     [default:                |
   |                                                     no-auto-compress]        |
+  | --refresh            --no-refresh                   Specifies whether ALTER  |
+  |                                                     STAGE {name} REFRESH     |
+  |                                                     should be executed after |
+  |                                                     uploading.               |
+  |                                                     [default: no-refresh]    |
   | --help           -h                                 Show this message and    |
   |                                                     exit.                    |
   +------------------------------------------------------------------------------+
@@ -18086,12 +18091,16 @@
   |                            [required]                                        |
   +------------------------------------------------------------------------------+
   +- Options --------------------------------------------------------------------+
-  | --encryption          [SNOWFLAKE_FULL|SNOWFLAKE_  Type of encryption         |
-  |                       SSE]                        supported for all files    |
-  |                                                   stored on the stage.       |
-  |                                                   [default: SNOWFLAKE_FULL]  |
-  | --help        -h                                  Show this message and      |
-  |                                                   exit.                      |
+  | --encryption                [SNOWFLAKE_FULL|SNOWFL  Type of encryption       |
+  |                             AKE_SSE]                supported for all files  |
+  |                                                     stored on the stage.     |
+  |                                                     [default:                |
+  |                                                     SNOWFLAKE_FULL]          |
+  | --enable-directory                                  Specifies whether        |
+  |                                                     directory support is     |
+  |                                                     enabled for the stage.   |
+  | --help              -h                              Show this message and    |
+  |                                                     exit.                    |
   +------------------------------------------------------------------------------+
   +- Connection configuration ---------------------------------------------------+
   | --connection,--environment    -c      TEXT     Name of the connection, as    |

tests/stage/test_stage.py

@@ -1427,3 +1427,34 @@ def test_recursive_upload_with_provided_temp_directory():
         tester.prepare(structure=NESTED_STRUCTURE)
         StageManager().copy_to_tmp_dir(Path(source_directory), temp_directory_path)
         tester.execute(local_path=source_directory)
+
+
+@mock.patch(f"{STAGE_MANAGER}.execute_query")
+def test_stage_create_enable_directory(mock_execute, runner, mock_cursor):
+    mock_execute.return_value = mock_cursor(["row"], [])
+    result = runner.invoke(["stage", "create", "stageName", "--enable-directory"])
+    assert result.exit_code == 0, result.output
+    mock_execute.assert_called_once_with(
+        "create stage if not exists IDENTIFIER('stageName') encryption = (type = 'SNOWFLAKE_FULL') directory = (enable = true)"
+    )
+
+
+@mock.patch(f"{STAGE_MANAGER}.execute_query")
+def test_stage_create_with_encryption_and_directory_options(
+    mock_execute, runner, mock_cursor
+):
+    mock_execute.return_value = mock_cursor(["row"], [])
+    result = runner.invoke(
+        [
+            "stage",
+            "create",
+            "stageName",
+            "--encryption",
+            "SNOWFLAKE_SSE",
+            "--enable-directory",
+        ]
+    )
+    assert result.exit_code == 0, result.output
+    mock_execute.assert_called_once_with(
+        "create stage if not exists IDENTIFIER('stageName') encryption = (type = 'SNOWFLAKE_SSE') directory = (enable = true)"
+    )

tests_integration/test_stage.py

@@ -18,6 +18,7 @@
 import tempfile
 import time
 from pathlib import Path
+
 import pytest
 
 from tests.stage.test_stage import RecursiveUploadTester, NESTED_STRUCTURE
@@ -788,6 +789,54 @@ def test_recursive_upload_glob_file_pattern(temporary_directory, runner, test_da
     ]
 
 
+@pytest.mark.integration
+def test_stage_copy_refreshes_stream(
+    runner, snowflake_session, test_database, tmp_path
+):
+    stage_name = "test_stage_stream"
+    stage_name_with_at = "@" + stage_name
+    stream_name = "test_stream"
+
+    # Create stage and stream
+    runner.invoke_with_connection_json(
+        ["stage", "create", stage_name_with_at, "--enable-directory"]
+    )
+    snowflake_session.execute_string(
+        f"CREATE OR REPLACE STREAM {stream_name} ON STAGE {stage_name}"
+    )
+
+    # Create test file
+    test_file = tmp_path / "test.txt"
+    test_file.write_text("test content")
+
+    # Initial copy without --refresh
+    result = runner.invoke_with_connection_json(
+        ["stage", "copy", str(test_file), stage_name_with_at]
+    )
+    assert result.exit_code == 0
+
+    # Check stream has no changes
+    stream_changes_cursor = snowflake_session.execute_string(
+        f"SELECT * FROM {stream_name}"
+    )
+    stream_changes = row_from_snowflake_session(stream_changes_cursor)
+    assert len(stream_changes) == 0
+
+    # Copy again with --refresh flag
+    result = runner.invoke_with_connection_json(
+        ["stage", "copy", str(test_file), stage_name_with_at, "--refresh"]
+    )
+    assert result.exit_code == 0
+
+    # Check stream has changes after --refresh was specified
+    stream_changes_cursor = snowflake_session.execute_string(
+        f"SELECT * FROM {stream_name}"
+    )
+    stream_changes = row_from_snowflake_session(stream_changes_cursor)
+    assert len(stream_changes) == 1
+    assert stream_changes[0]["RELATIVE_PATH"] == "test.txt"
+
+
 @pytest.mark.integration
 def test_recursive_upload_no_recursive_glob_pattern(
     temporary_directory, runner, test_database
@@ -866,3 +915,24 @@ def test_create_encryption(runner, test_database, encryption):
     )
     assert result.exit_code == 0, result.output
     assert result.json == {"status": f"Stage area A_STAGE successfully created."}
+
+
+@pytest.mark.integration
+def test_create_directory_option(runner, test_database):
+    stage_name = "stage_with_directory"
+    result = runner.invoke_with_connection_json(
+        ["stage", "create", stage_name, "--enable-directory"]
+    )
+    assert result.exit_code == 0, result.output
+    assert result.json == {
+        "status": f"Stage area {stage_name.upper()} successfully created."
+    }
+
+    # Verify directory is enabled
+    result = runner.invoke_with_connection_json(["stage", "describe", stage_name])
+    assert result.exit_code == 0, result.output
+    assert any(
+        row.get("parent_property") == "DIRECTORY"
+        and row.get("property_value") == "true"
+        for row in result.json
+    )