Skip to content

Add stage encryption type changes #3011

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
sfc-gh-svasudevan wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Add stage encryption type changes #3011

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
1324a45
Add stage encryption type changes
sfc-gh-svasudevan May 11, 2026
4997b2d
Fix tests
sfc-gh-svasudevan May 12, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 34 additions & 3 deletions src/snowflake/cli/_plugins/spcs/services/commands.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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"
Expand Down Expand Up @@ -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)
Expand Down
5 changes: 5 additions & 0 deletions tests/__snapshots__/test_help_messages.ambr
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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). |
Comment on lines +18439 to +18443

@sfc-gh-mdrach sfc-gh-mdrach May 12, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

looks like the | are crooked

| --job-name TEXT Name for the build job service. |
| If not provided, a name will be |
| auto-generated. |
Expand Down
201 changes: 193 additions & 8 deletions tests/spcs/test_services.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -2191,6 +2192,191 @@ 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("snowflake.cli.api.cli_global_context.get_cli_context")
@patch("time.sleep")
@patch(
"snowflake.cli._plugins.spcs.services.commands.ObjectManager",
)
@patch(
"snowflake.cli._plugins.spcs.services.commands.ServiceManager",
)
@patch(
"snowflake.cli._plugins.spcs.services.commands.StageManager",
)
def test_build_image_cli_temp_stage_encryption_option(
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 = 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
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.spcs.services.commands.StageManager",
)
def test_build_image_cli_explicit_stage_does_not_call_create(
mock_stage_manager_class,
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 = 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
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")
Expand Down Expand Up @@ -2339,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,
Expand All @@ -2367,7 +2549,10 @@ def test_build_image_cli_recursive_upload_with_nested_dirs(
partials_dir.mkdir()
(partials_dir / "header.html").write_text("<header>Header</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
Expand Down Expand Up @@ -2410,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))
Expand Down
Loading