From 1324a450b20112b792950f51aa58921bf025ae12 Mon Sep 17 00:00:00 2001 From: Supriya Vasudevan Date: Mon, 11 May 2026 19:38:01 +0000 Subject: [PATCH 1/2] Add stage encryption type changes Description Testing --- .../cli/_plugins/spcs/services/commands.py | 37 +++- tests/__snapshots__/test_help_messages.ambr | 5 + tests/spcs/test_services.py | 190 ++++++++++++++++++ 3 files changed, 229 insertions(+), 3 deletions(-) diff --git a/src/snowflake/cli/_plugins/spcs/services/commands.py b/src/snowflake/cli/_plugins/spcs/services/commands.py index 0f0c2c7b3a..83d101fa81 100644 --- a/src/snowflake/cli/_plugins/spcs/services/commands.py +++ b/src/snowflake/cli/_plugins/spcs/services/commands.py @@ -35,7 +35,10 @@ from snowflake.cli._plugins.spcs.services.service_project_paths import ( ServiceProjectPaths, ) -from snowflake.cli._plugins.stage.manager import StageManager +from snowflake.cli._plugins.stage.manager import ( + InternalStageEncryptionType, + StageManager, +) 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 ( @@ -707,6 +710,15 @@ def build_image( help="Stage to store build context files. Format: [db.][schema.]stage_name. If not provided, a temporary stage will be created and dropped automatically. If provided, only the uploaded build context files will be removed after the build completes.", show_default=False, ), + stage_encryption: Optional[str] = typer.Option( + None, + "--stage-encryption", + help=( + "SNOWFLAKE_SSE or SNOWFLAKE_FULL for an auto-created temporary stage; " + "ignored when --stage is set. Omit for legacy CREATE STAGE (no ENCRYPTION clause)." + ), + show_default=False, + ), job_name: str = typer.Option( None, "--job-name", @@ -741,6 +753,9 @@ def build_image( If --stage is not provided, a stage will be automatically created using the current session's database and schema context, and dropped after the build completes. If your session doesn't have a database/schema set, you should provide --stage explicitly. + + Optional ``--stage-encryption`` applies only when the CLI creates that temporary stage + (for example use ``SNOWFLAKE_SSE`` when your deployment requires it for stage-mounted builds). """ # Verify Dockerfile exists in build context directory dockerfile_path = build_context_dir / "Dockerfile" @@ -776,15 +791,31 @@ def build_image( f"Invalid job name '{job_name}'. Must be a valid unquoted identifier." ) - stage_manager = StageManager() use_temporary_stage = stage is None + temp_encryption: InternalStageEncryptionType | None = None + if stage_encryption is not None: + key = stage_encryption.strip().upper() + try: + temp_encryption = InternalStageEncryptionType(key) + except ValueError: + allowed = ", ".join(sorted(e.value for e in InternalStageEncryptionType)) + raise CliArgumentError( + f"Invalid --stage-encryption {stage_encryption!r}. " + f"Expected one of: {allowed}." + ) + + stage_manager = StageManager() + if use_temporary_stage: # Create a stage stage = f"{job_name}_stage" cli_console.step(f"Creating temporary stage: {stage}") stage_fqn = FQN.from_string(stage).using_context() - stage_manager.create(fqn=stage_fqn) + if temp_encryption is not None: + stage_manager.create(fqn=stage_fqn, encryption=temp_encryption) + else: + stage_manager.create(fqn=stage_fqn) else: # Use the provided stage (ensure it exists) stage_fqn = FQN.from_string(stage) diff --git a/tests/__snapshots__/test_help_messages.ambr b/tests/__snapshots__/test_help_messages.ambr index dd499e52f8..db2f8d5239 100644 --- a/tests/__snapshots__/test_help_messages.ambr +++ b/tests/__snapshots__/test_help_messages.ambr @@ -18436,6 +18436,11 @@ | the uploaded build context files | | will be removed after the build | | completes. | + | --stage-encryption TEXT SNOWFLAKE_SSE or SNOWFLAKE_FULL for an | + | auto-created temporary stage; | + | ignored when --stage is set. Omit for| + | legacy CREATE STAGE (no ENCRYPTION | + | clause). | | --job-name TEXT Name for the build job service. | | If not provided, a name will be | | auto-generated. | diff --git a/tests/spcs/test_services.py b/tests/spcs/test_services.py index 6111ca498c..7afeebd483 100644 --- a/tests/spcs/test_services.py +++ b/tests/spcs/test_services.py @@ -27,6 +27,7 @@ from snowflake.cli._plugins.spcs.common import NoPropertiesProvidedError from snowflake.cli._plugins.spcs.services.commands import _service_name_callback from snowflake.cli._plugins.spcs.services.manager import ServiceManager +from snowflake.cli._plugins.stage.manager import InternalStageEncryptionType from snowflake.cli.api.constants import ObjectType from snowflake.cli.api.identifiers import FQN from snowflake.cli.api.project.util import to_string_literal @@ -2191,6 +2192,195 @@ def test_build_image_cli_parameter_validation(runner, temporary_directory): assert result.exit_code != 0 assert "Invalid job name" in result.output + # Test 4: Invalid stage encryption + result = runner.invoke( + [ + "spcs", + "service", + "build-image", + "--compute-pool", + "test_pool", + "--image-repository", + "db.schema.repo", + "--image-name", + "my_image", + "--image-tag", + "v1.0", + "--build-context-dir", + str(build_context), + "--stage-encryption", + "NOT_A_REAL_TYPE", + ], + catch_exceptions=False, + ) + assert result.exit_code != 0 + assert "Invalid --stage-encryption" in result.output + + +@pytest.mark.parametrize( + "stage_encryption_cli, expect_encryption_kw", + [ + ("snowflake_sse", InternalStageEncryptionType.SNOWFLAKE_SSE), + ("SNOWFLAKE_FULL", InternalStageEncryptionType.SNOWFLAKE_FULL), + (None, None), + ], +) +@patch("time.sleep") +@patch( + "snowflake.cli._plugins.spcs.services.commands.ObjectManager", +) +@patch( + "snowflake.cli._plugins.spcs.services.commands.ServiceManager", +) +@patch( + "snowflake.cli._plugins.stage.manager.StageManager.put", +) +@patch( + "snowflake.cli._plugins.stage.manager.StageManager.execute_query", +) +@patch( + "snowflake.cli._plugins.stage.manager.StageManager.create", +) +def test_build_image_cli_temp_stage_encryption_option( + mock_stage_create, + mock_stage_execute_query, + mock_stage_put, + mock_service_manager_class, + mock_object_manager_class, + mock_sleep, + stage_encryption_cli, + expect_encryption_kw, + runner, + temporary_directory, +): + """--stage-encryption is forwarded when the CLI creates a temporary stage; omit for legacy CREATE STAGE.""" + build_context = Path(temporary_directory) / "build_context" + build_context.mkdir() + (build_context / "Dockerfile").write_text("FROM alpine\n") + + mock_stage_put.return_value = Mock(fetchall=lambda: []) + + mock_service_manager = Mock() + mock_service_manager_class.return_value = mock_service_manager + mock_build_cursor = Mock(spec=SnowflakeCursor) + mock_build_cursor.__iter__ = Mock(return_value=iter([])) + mock_build_cursor.fetchone.return_value = {"status": "DONE"} + mock_build_cursor.description = [] + mock_build_cursor.query = "" + mock_service_manager.build_image.return_value = mock_build_cursor + mock_service_manager.stream_logs.return_value = iter( + [("__TERMINAL_STATUS__", "DONE")] + ) + + mock_object_manager = Mock() + mock_object_manager_class.return_value = mock_object_manager + mock_describe_cursor = Mock() + mock_describe_cursor.fetchone.return_value = {"status": "RUNNING"} + mock_object_manager.describe.return_value = mock_describe_cursor + + cmd = [ + "spcs", + "service", + "build-image", + "--compute-pool", + "test_pool", + "--image-repository", + "db.schema.repo", + "--image-name", + "my_image", + "--image-tag", + "v1.0", + "--build-context-dir", + str(build_context), + ] + if stage_encryption_cli is not None: + cmd.extend(["--stage-encryption", stage_encryption_cli]) + + result = runner.invoke(cmd, catch_exceptions=False) + assert result.exit_code == 0, f"Command failed with output: {result.output}" + mock_stage_create.assert_called_once() + if expect_encryption_kw is not None: + assert mock_stage_create.call_args.kwargs["encryption"] == expect_encryption_kw + else: + assert "encryption" not in mock_stage_create.call_args.kwargs + + +@patch("time.sleep") +@patch( + "snowflake.cli._plugins.spcs.services.commands.ObjectManager", +) +@patch( + "snowflake.cli._plugins.spcs.services.commands.ServiceManager", +) +@patch( + "snowflake.cli._plugins.stage.manager.StageManager.put", +) +@patch( + "snowflake.cli._plugins.stage.manager.StageManager.execute_query", +) +@patch( + "snowflake.cli._plugins.stage.manager.StageManager.create", +) +def test_build_image_cli_explicit_stage_does_not_call_create( + mock_stage_create, + mock_stage_execute_query, + mock_stage_put, + mock_service_manager_class, + mock_object_manager_class, + mock_sleep, + runner, + temporary_directory, +): + """With --stage, StageManager.create is not used (--stage-encryption ignored).""" + build_context = Path(temporary_directory) / "build_context" + build_context.mkdir() + (build_context / "Dockerfile").write_text("FROM alpine\n") + + mock_stage_put.return_value = Mock(fetchall=lambda: []) + + mock_service_manager = Mock() + mock_service_manager_class.return_value = mock_service_manager + mock_build_cursor = Mock(spec=SnowflakeCursor) + mock_build_cursor.__iter__ = Mock(return_value=iter([])) + mock_build_cursor.fetchone.return_value = {"status": "DONE"} + mock_build_cursor.description = [] + mock_build_cursor.query = "" + mock_service_manager.build_image.return_value = mock_build_cursor + mock_service_manager.stream_logs.return_value = iter( + [("__TERMINAL_STATUS__", "DONE")] + ) + + mock_object_manager = Mock() + mock_object_manager_class.return_value = mock_object_manager + mock_describe_cursor = Mock() + mock_describe_cursor.fetchone.return_value = {"status": "RUNNING"} + mock_object_manager.describe.return_value = mock_describe_cursor + + result = runner.invoke( + [ + "spcs", + "service", + "build-image", + "--compute-pool", + "test_pool", + "--image-repository", + "db.schema.repo", + "--image-name", + "my_image", + "--image-tag", + "v1.0", + "--build-context-dir", + str(build_context), + "--stage", + "test_stage", + "--stage-encryption", + "SNOWFLAKE_SSE", + ], + catch_exceptions=False, + ) + assert result.exit_code == 0, f"Command failed with output: {result.output}" + mock_stage_create.assert_not_called() + # Tests for check_terminal_status parameter in stream_logs @patch("snowflake.cli._plugins.spcs.services.manager.ObjectManager") From 4997b2d3e0adb0814818675c07c7fde77f9e09b7 Mon Sep 17 00:00:00 2001 From: Supriya Vasudevan Date: Tue, 12 May 2026 04:38:16 +0000 Subject: [PATCH 2/2] Fix tests Description Testing --- tests/spcs/test_services.py | 63 +++++++++++++++++-------------------- 1 file changed, 29 insertions(+), 34 deletions(-) diff --git a/tests/spcs/test_services.py b/tests/spcs/test_services.py index 7afeebd483..a8ced315e5 100644 --- a/tests/spcs/test_services.py +++ b/tests/spcs/test_services.py @@ -2225,6 +2225,7 @@ def test_build_image_cli_parameter_validation(runner, temporary_directory): (None, None), ], ) +@patch("snowflake.cli.api.cli_global_context.get_cli_context") @patch("time.sleep") @patch( "snowflake.cli._plugins.spcs.services.commands.ObjectManager", @@ -2233,32 +2234,32 @@ def test_build_image_cli_parameter_validation(runner, temporary_directory): "snowflake.cli._plugins.spcs.services.commands.ServiceManager", ) @patch( - "snowflake.cli._plugins.stage.manager.StageManager.put", -) -@patch( - "snowflake.cli._plugins.stage.manager.StageManager.execute_query", -) -@patch( - "snowflake.cli._plugins.stage.manager.StageManager.create", + "snowflake.cli._plugins.spcs.services.commands.StageManager", ) def test_build_image_cli_temp_stage_encryption_option( - mock_stage_create, - mock_stage_execute_query, - mock_stage_put, + mock_stage_manager_class, mock_service_manager_class, mock_object_manager_class, mock_sleep, + mock_get_cli_context, stage_encryption_cli, expect_encryption_kw, runner, temporary_directory, ): """--stage-encryption is forwarded when the CLI creates a temporary stage; omit for legacy CREATE STAGE.""" + mock_ctx = Mock() + mock_ctx.connection = Mock(database="TESTDB", schema="PUBLIC") + mock_get_cli_context.return_value = mock_ctx + build_context = Path(temporary_directory) / "build_context" build_context.mkdir() (build_context / "Dockerfile").write_text("FROM alpine\n") - mock_stage_put.return_value = Mock(fetchall=lambda: []) + mock_stage = Mock() + mock_stage_manager_class.return_value = mock_stage + mock_stage.put_recursive.return_value = iter([]) + mock_stage.create.return_value = Mock() mock_service_manager = Mock() mock_service_manager_class.return_value = mock_service_manager @@ -2298,11 +2299,11 @@ def test_build_image_cli_temp_stage_encryption_option( result = runner.invoke(cmd, catch_exceptions=False) assert result.exit_code == 0, f"Command failed with output: {result.output}" - mock_stage_create.assert_called_once() + mock_stage.create.assert_called_once() if expect_encryption_kw is not None: - assert mock_stage_create.call_args.kwargs["encryption"] == expect_encryption_kw + assert mock_stage.create.call_args.kwargs["encryption"] == expect_encryption_kw else: - assert "encryption" not in mock_stage_create.call_args.kwargs + assert "encryption" not in mock_stage.create.call_args.kwargs @patch("time.sleep") @@ -2313,18 +2314,10 @@ def test_build_image_cli_temp_stage_encryption_option( "snowflake.cli._plugins.spcs.services.commands.ServiceManager", ) @patch( - "snowflake.cli._plugins.stage.manager.StageManager.put", -) -@patch( - "snowflake.cli._plugins.stage.manager.StageManager.execute_query", -) -@patch( - "snowflake.cli._plugins.stage.manager.StageManager.create", + "snowflake.cli._plugins.spcs.services.commands.StageManager", ) def test_build_image_cli_explicit_stage_does_not_call_create( - mock_stage_create, - mock_stage_execute_query, - mock_stage_put, + mock_stage_manager_class, mock_service_manager_class, mock_object_manager_class, mock_sleep, @@ -2336,7 +2329,10 @@ def test_build_image_cli_explicit_stage_does_not_call_create( build_context.mkdir() (build_context / "Dockerfile").write_text("FROM alpine\n") - mock_stage_put.return_value = Mock(fetchall=lambda: []) + mock_stage = Mock() + mock_stage_manager_class.return_value = mock_stage + mock_stage.put_recursive.return_value = iter([]) + mock_stage.create.return_value = Mock() mock_service_manager = Mock() mock_service_manager_class.return_value = mock_service_manager @@ -2379,7 +2375,7 @@ def test_build_image_cli_explicit_stage_does_not_call_create( catch_exceptions=False, ) assert result.exit_code == 0, f"Command failed with output: {result.output}" - mock_stage_create.assert_not_called() + mock_stage.create.assert_not_called() # Tests for check_terminal_status parameter in stream_logs @@ -2529,14 +2525,10 @@ def test_stream_logs_without_terminal_status_check(mock_sleep, mock_logs): "snowflake.cli._plugins.spcs.services.commands.ServiceManager", ) @patch( - "snowflake.cli._plugins.stage.manager.StageManager.put", -) -@patch( - "snowflake.cli._plugins.stage.manager.StageManager.execute_query", + "snowflake.cli._plugins.spcs.services.commands.StageManager", ) def test_build_image_cli_recursive_upload_with_nested_dirs( - mock_stage_execute_query, - mock_stage_put, + mock_stage_manager_class, mock_service_manager_class, mock_object_manager_class, mock_sleep, @@ -2557,7 +2549,10 @@ def test_build_image_cli_recursive_upload_with_nested_dirs( partials_dir.mkdir() (partials_dir / "header.html").write_text("
Header
") - mock_stage_put.return_value = Mock(fetchall=lambda: []) + mock_stage = Mock() + mock_stage_manager_class.return_value = mock_stage + mock_stage.put_recursive.return_value = iter([]) + mock_stage.create.return_value = Mock() mock_service_manager = Mock() mock_service_manager_class.return_value = mock_service_manager @@ -2600,7 +2595,7 @@ def test_build_image_cli_recursive_upload_with_nested_dirs( assert result.exit_code == 0, f"Command failed with output: {result.output}" stage_paths = set() - for c in mock_stage_put.call_args_list: + for c in mock_stage.put_recursive.call_args_list: sp = c.kwargs.get("stage_path", None) if sp is not None: stage_paths.add(str(sp))