diff --git a/src/snowflake/cli/_plugins/spcs/services/commands.py b/src/snowflake/cli/_plugins/spcs/services/commands.py index 0f0c2c7b3a..ab23a46c99 100644 --- a/src/snowflake/cli/_plugins/spcs/services/commands.py +++ b/src/snowflake/cli/_plugins/spcs/services/commands.py @@ -35,11 +35,13 @@ 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 ( - IfExistsOption, IfNotExistsOption, OverrideableOption, entity_argument, @@ -178,32 +180,9 @@ def _service_name_callback(name: FQN) -> FQN: help_example='`list --like "my%"` lists all services that begin with “my”.' ), scope_option=scope_option(help_example="`list --in compute-pool my_pool`"), - ommit_commands=["drop"], ) -@app.command(requires_connection=True) -def drop( - name: FQN = ServiceNameArgument, - if_exists: bool = IfExistsOption(), - force: bool = typer.Option( - False, - "--force", - help="Drops the service along with any block storage volumes it contains. Without this flag, dropping a service that contains block storage volumes fails.", - is_flag=True, - ), - **options, -) -> CommandResult: - """ - Drops a service. - """ - return SingleQueryResult( - ServiceManager().drop( - service_name=name.sql_identifier, if_exists=if_exists, force=force - ) - ) - - @app.command(requires_connection=True) def create( name: FQN = ServiceNameArgument, @@ -707,6 +686,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 +729,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 to that auto-created stage; use SNOWFLAKE_SSE + when stage-mounted builds require server-side encryption on your deployment. """ # Verify Dockerfile exists in build context directory dockerfile_path = build_context_dir / "Dockerfile" @@ -776,15 +767,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/src/snowflake/cli/_plugins/spcs/services/manager.py b/src/snowflake/cli/_plugins/spcs/services/manager.py index 4a20077bcd..fbc5dea174 100644 --- a/src/snowflake/cli/_plugins/spcs/services/manager.py +++ b/src/snowflake/cli/_plugins/spcs/services/manager.py @@ -304,50 +304,26 @@ def execute_job( replicas=replicas, ) - def build_image( + def build_image_registry_url( self, - job_service_name: str, - compute_pool: str, image_repository: str, - image_name: str, - image_tag: str, - stage: str, - build_context_path: str, - external_access_integrations: Optional[List[str]], - async_mode: bool = True, - ) -> SnowflakeCursor: - """ - Builds a container image using EXECUTE JOB SERVICE with the Snowflake image builder. + ) -> str: + """Build the canonical image-registry URL for ``image_repository``. - Args: - job_service_name: Name for the build job service - compute_pool: Compute pool to run the build job - image_repository: Repository path in format [db.][schema.]repo_name - image_name: Name for the built image - image_tag: Tag for the built image - stage: Stage where the build context is uploaded - build_context_path: Path on stage where build context is located - external_access_integrations: List of EAI names - async_mode: Whether to run the build asynchronously + Returns ``-.registry-local.snowflakecomputing.com///``, + the same shape ``build_image`` injects as the ``IMAGE_REGISTRY_URL`` env var. Exposed + separately so the post-push validation hook can build the URL it passes to + ``SYSTEM$START_IMAGE_VALIDATION_SERVICE`` without re-implementing FQN parsing. + + Raises ``ValueError`` if ``image_repository`` is missing the database or schema parts. """ - # Get organization name - # Using execute_string (same as get_account) for consistency *_, org_cursor = self._conn.execute_string( "SELECT CURRENT_ORGANIZATION_NAME()", cursor_class=DictCursor ) org_name = org_cursor.fetchone()["CURRENT_ORGANIZATION_NAME()"].lower() - - # Get account name using existing utility function account_name = get_account(self._conn) - # Parse the image repository using FQN utility - # This handles [db.][schema.]repo_name format automatically repo_fqn = FQN.from_string(image_repository).using_connection(self._conn) - - # Validate that we have all required parts - # Note: FQN parsing gives us: "db.schema.repo" → both present, - # "schema.repo" → only schema, "repo" → neither - # It's impossible to have database without schema if not repo_fqn.database or not repo_fqn.schema: missing_parts = [] if not repo_fqn.database: @@ -362,26 +338,53 @@ def build_image( f"Provided: '{image_repository}'" ) - # Build the image registry URL - # Format: -.registry-local.snowflakecomputing.com/// - image_registry_url = ( + return ( f"{org_name}-{account_name}.registry-local.snowflakecomputing.com/" f"{repo_fqn.database.lower()}/{repo_fqn.schema.lower()}/{repo_fqn.name.lower()}" ) - # Create the specification for the image build job + def build_image( + self, + job_service_name: str, + compute_pool: str, + image_repository: str, + image_name: str, + image_tag: str, + stage: str, + build_context_path: str, + external_access_integrations: Optional[List[str]], + async_mode: bool = True, + ) -> SnowflakeCursor: + """ + Builds a container image using EXECUTE JOB SERVICE with the Snowflake image builder. + + Args: + job_service_name: Name for the build job service + compute_pool: Compute pool to run the build job + image_repository: Repository path in format [db.][schema.]repo_name + image_name: Name for the built image + image_tag: Tag for the built image + stage: Stage where the build context is uploaded + build_context_path: Path on stage where build context is located + external_access_integrations: List of EAI names + async_mode: Whether to run the build asynchronously + """ + image_registry_url = self.build_image_registry_url(image_repository) + + env_vars = { + "IMAGE_REGISTRY_URL": image_registry_url, + "IMAGE_NAME": image_name, + "IMAGE_TAG": image_tag, + "BUILD_CONTEXT": "/app", + } + spec = { "spec": { "containers": [ { "name": "main", "image": "/snowflake/images/snowflake_images/sf-image-build:0.0.1", - "env": { - "IMAGE_REGISTRY_URL": image_registry_url, - "IMAGE_NAME": image_name, - "IMAGE_TAG": image_tag, - "BUILD_CONTEXT": "/app", - }, + "env": env_vars, "volumeMounts": [ { "name": "code-volume", @@ -711,15 +714,6 @@ def suspend(self, service_name: str): def resume(self, service_name: str): return self.execute_query(f"alter service {service_name} resume") - def drop( - self, service_name: str, if_exists: bool = False, force: bool = False - ) -> SnowflakeCursor: - if_exists_clause = " if exists" if if_exists else "" - force_clause = " force" if force else "" - return self.execute_query( - f"drop service{if_exists_clause} {service_name}{force_clause}" - ) - def set_property( self, service_name: str, 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..77533a001a 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 @@ -1550,71 +1551,6 @@ def test_resume_cli(mock_resume, mock_cursor, runner): assert "Statement executed successfully" in result.output -@pytest.mark.parametrize( - "if_exists, force, expected_query", - [ - (False, False, "drop service test_service"), - (True, False, "drop service if exists test_service"), - (False, True, "drop service test_service force"), - (True, True, "drop service if exists test_service force"), - ], -) -@patch(EXECUTE_QUERY) -def test_drop(mock_execute_query, if_exists, force, expected_query): - service_name = "test_service" - cursor = Mock(spec=SnowflakeCursor) - mock_execute_query.return_value = cursor - result = ServiceManager().drop( - service_name=service_name, if_exists=if_exists, force=force - ) - mock_execute_query.assert_called_once_with(expected_query) - assert result == cursor - - -@patch("snowflake.cli._plugins.spcs.services.manager.ServiceManager.drop") -def test_drop_cli_default(mock_drop, mock_cursor, runner): - service_name = "test_service" - cursor = mock_cursor( - rows=[["Statement executed successfully."]], columns=["status"] - ) - mock_drop.return_value = cursor - result = runner.invoke(["spcs", "service", "drop", service_name]) - assert result.exit_code == 0, result.output - mock_drop.assert_called_once_with( - service_name=f"IDENTIFIER('{service_name}')", if_exists=False, force=False - ) - - -@patch("snowflake.cli._plugins.spcs.services.manager.ServiceManager.drop") -def test_drop_cli_force(mock_drop, mock_cursor, runner): - service_name = "test_service" - cursor = mock_cursor( - rows=[["Statement executed successfully."]], columns=["status"] - ) - mock_drop.return_value = cursor - result = runner.invoke(["spcs", "service", "drop", service_name, "--force"]) - assert result.exit_code == 0, result.output - mock_drop.assert_called_once_with( - service_name=f"IDENTIFIER('{service_name}')", if_exists=False, force=True - ) - - -@patch("snowflake.cli._plugins.spcs.services.manager.ServiceManager.drop") -def test_drop_cli_if_exists_and_force(mock_drop, mock_cursor, runner): - service_name = "test_service" - cursor = mock_cursor( - rows=[["Statement executed successfully."]], columns=["status"] - ) - mock_drop.return_value = cursor - result = runner.invoke( - ["spcs", "service", "drop", service_name, "--if-exists", "--force"] - ) - assert result.exit_code == 0, result.output - mock_drop.assert_called_once_with( - service_name=f"IDENTIFIER('{service_name}')", if_exists=True, force=True - ) - - @patch(EXECUTE_QUERY) def test_set_property(mock_execute_query): service_name = "test_service" @@ -1999,6 +1935,8 @@ def test_build_image(mock_execute_query, mock_get_account): assert container["env"]["IMAGE_NAME"] == image_name assert container["env"]["IMAGE_TAG"] == image_tag assert container["env"]["BUILD_CONTEXT"] == "/app" + assert "WORKLOAD_TYPE" not in container["env"] + assert "SKIP_VALIDATION" not in container["env"] assert result == cursor @@ -2191,6 +2129,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") @@ -2431,3 +2558,39 @@ def test_build_image_hidden_by_default(runner): result = runner.invoke(["spcs", "service", "--help"]) assert result.exit_code == 0 assert "build-image" not in result.output + + +@patch("snowflake.cli._plugins.connection.util.get_account") +def test_build_image_registry_url(mock_get_account): + """build_image_registry_url builds the canonical -/// URL.""" + mock_org_cursor = Mock() + mock_org_cursor.fetchone.return_value = {"CURRENT_ORGANIZATION_NAME()": "TEST_ORG"} + mock_conn = Mock() + mock_conn.execute_string.return_value = (None, mock_org_cursor) + mock_conn.database = None + mock_conn.schema = None + mock_conn.account = "test_account" + mock_get_account.return_value = "test_account" + + manager = ServiceManager(connection=mock_conn) + url = manager.build_image_registry_url("MyDB.MySchema.MyRepo") + assert ( + url + == "test_org-test_account.registry-local.snowflakecomputing.com/mydb/myschema/myrepo" + ) + + +@patch("snowflake.cli._plugins.connection.util.get_account") +def test_build_image_registry_url_missing_parts(mock_get_account): + mock_org_cursor = Mock() + mock_org_cursor.fetchone.return_value = {"CURRENT_ORGANIZATION_NAME()": "TEST_ORG"} + mock_conn = Mock() + mock_conn.execute_string.return_value = (None, mock_org_cursor) + mock_conn.database = None + mock_conn.schema = None + mock_conn.account = "test_account" + mock_get_account.return_value = "test_account" + + manager = ServiceManager(connection=mock_conn) + with pytest.raises(ValueError, match="Image repository requires database and schema"): + manager.build_image_registry_url("repo_only")