From 174de37859af21b194ea172a9cac93bcba1f6610 Mon Sep 17 00:00:00 2001 From: Supriya Vasudevan Date: Mon, 11 May 2026 18:11:37 +0000 Subject: [PATCH 1/2] Add stage type property to build-image Description Testing --- .../cli/_plugins/spcs/services/commands.py | 330 ++++++- .../cli/_plugins/spcs/services/manager.py | 177 ++-- tests/__snapshots__/test_help_messages.ambr | 5 + tests/spcs/test_services.py | 825 ++++++++++++++++-- 4 files changed, 1193 insertions(+), 144 deletions(-) diff --git a/src/snowflake/cli/_plugins/spcs/services/commands.py b/src/snowflake/cli/_plugins/spcs/services/commands.py index 0f0c2c7b3a..9dd5abf4c5 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, @@ -78,6 +80,20 @@ short_help="Manages services.", ) +# Exit codes used by the post-push image-validation hook in build-image. +# Distinct from the default Click "command failed" exit (1) so callers / CI can +# tell apart "image push failed" (1), "image was pushed but workload-image +# smoke test reported the image is broken" (_EXIT_VALIDATION_FAILED), and +# "smoke test could not run to completion (transient infra)" (_EXIT_VALIDATION_TRANSIENT). +_EXIT_VALIDATION_FAILED = 2 +_EXIT_VALIDATION_TRANSIENT = 3 + +# Status strings the validator uses to label terminal outcomes for the user. +_VALIDATION_STATUS_PASS = "pass" +_VALIDATION_STATUS_FAIL = "fail" +_VALIDATION_STATUS_TRANSIENT = "transient" +_VALIDATION_STATUS_SKIPPED = "skipped" + # Define common options container_name_option = typer.Option( ..., @@ -178,32 +194,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 +700,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", @@ -718,6 +720,32 @@ def build_image( "--eai-name", help="Identifies external access integrations (EAI) that the build job can access. This option may be specified multiple times for multiple EAIs.", ), + workload_type: Optional[str] = typer.Option( + None, + "--workload-type", + help=( + "Workload identifier (e.g. 'notebook', 'ml_job') that selects " + "which validation YAML the sf-image-build runner applies and, " + "after a successful push, which server-side smoke test the CLI " + "asks for via SYSTEM$START_IMAGE_VALIDATION_SERVICE. Propagated " + "to the build job as the WORKLOAD_TYPE env var. If omitted, both " + "the runner-side advisory checks and the post-push smoke test " + "are skipped." + ), + show_default=False, + ), + skip_validation: bool = typer.Option( + False, + "--skip-validation", + help=( + "Skip both the runner-side advisory validation and the post-push " + "server-side smoke test. The CLI propagates SKIP_VALIDATION=true " + "into the build job's env (which the sf-image-build runner reads " + "to bypass its in-process checks) and short-circuits before the " + "smoke-test system function is called. Image is built and pushed " + "as usual." + ), + ), **options, ) -> CommandResult: """ @@ -741,6 +769,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" @@ -764,6 +795,18 @@ def build_image( f"Invalid image tag '{image_tag}'. Must contain only alphanumeric characters, hyphens, underscores, and dots." ) + # Defensive validation: workload_type becomes a JSON env value that is then + # used by the runner to construct /etc/sf-image-build/configs/.yaml. + # Restricting to alphanumerics + _ - . prevents path traversal and shell + # metacharacters from sneaking through. The runner validates further by + # rejecting any value that does not map to an existing baked-in YAML. + if workload_type is not None and not workload_type.replace("_", "").replace( + "-", "" + ).replace(".", "").isalnum(): + raise CliArgumentError( + f"Invalid workload type '{workload_type}'. Must contain only alphanumeric characters, hyphens, underscores, and dots." + ) + # Generate a unique identifier for this build operation build_uuid = uuid.uuid4().hex[:8] @@ -776,15 +819,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) @@ -819,6 +878,8 @@ def build_image( build_context_path=build_context_stage_path, external_access_integrations=external_access_integrations, async_mode=True, # Always async so we can stream logs + workload_type=workload_type, + skip_validation=skip_validation, ) cli_console.step(f"Waiting for job to start...") @@ -962,4 +1023,219 @@ def build_image( except ProgrammingError as e: cli_console.warning(f"Failed to clean up build context files: {e}") + # Post-push image validation: only if (a) the build job actually finished + # successfully (otherwise the image isn't even there to validate), (b) the + # caller opted in via --workload-type, and (c) they didn't explicitly + # disable validation via --skip-validation. + if ( + final_status == "DONE" + and workload_type is not None + and not skip_validation + ): + validation_result = _run_image_validation( + service_manager=service_manager, + image_repository=image_repository, + image_name=image_name, + image_tag=image_tag, + workload_type=workload_type, + ) + _report_validation_result(validation_result) + status = validation_result.get("status") + if status == _VALIDATION_STATUS_FAIL: + raise typer.Exit(_EXIT_VALIDATION_FAILED) + if status == _VALIDATION_STATUS_TRANSIENT: + raise typer.Exit(_EXIT_VALIDATION_TRANSIENT) + return SingleQueryResult(cursor) + + +def _run_image_validation( + service_manager: ServiceManager, + image_repository: str, + image_name: str, + image_tag: str, + workload_type: str, +) -> dict: + """ + Server-orchestrated post-push smoke test. + + Calls SYSTEM$START_IMAGE_VALIDATION_SERVICE to provision a per-workload smoke + service against the just-pushed image, polls DESCRIBE SERVICE until it leaves + PENDING, streams its logs to the terminal, classifies the terminal status, and + always calls SYSTEM$CLEANUP_IMAGE_VALIDATION_SERVICE in a finally so a Ctrl-C + or unexpected exception still triggers cleanup. + + Returns ``{"status": "pass"|"fail"|"transient"|"skipped", ...}`` so the caller + can translate to a CLI exit code. + """ + image_url = service_manager.build_image_registry_url( + image_repository=image_repository, + ) + image_url_with_tag = f"{image_url}/{image_name}:{image_tag}" + + cli_console.step( + f"Starting server-side image validation for workload type '{workload_type}'..." + ) + try: + start_payload = service_manager.start_image_validation_service( + image_url=image_url_with_tag, + workload_type=workload_type, + ) + except ProgrammingError as e: + # Treat any infra-level error from the start function as transient: the + # image was pushed, but we couldn't even kick off validation. The user + # can re-run validation manually. + cli_console.warning(f"Could not start image validation: {e}") + return { + "status": _VALIDATION_STATUS_TRANSIENT, + "reason": str(e), + "workload_type": workload_type, + } + + if start_payload.get("status") == "skipped": + return { + "status": _VALIDATION_STATUS_SKIPPED, + "reason": start_payload.get("reason") + or f"validation not enabled for workload {workload_type}", + "workload_type": workload_type, + } + + service_name = start_payload.get("service_name") + if not service_name: + return { + "status": _VALIDATION_STATUS_TRANSIENT, + "reason": "validation start response missing service_name", + "workload_type": workload_type, + "raw": start_payload, + } + + cli_console.message(f"Validation service: {service_name}") + + final_status: Optional[str] = None + try: + final_status = _poll_then_stream_validation_logs(service_manager, service_name) + except KeyboardInterrupt: + cli_console.warning( + f"\nImage validation '{service_name}' interrupted; attempting cleanup." + ) + return _classify_validation_status("CANCELLED", service_name, workload_type) + finally: + try: + cleanup_payload = service_manager.cleanup_image_validation_service( + service_name=service_name, + ) + if cleanup_payload.get("error"): + cli_console.warning( + f"Cleanup of validation service '{service_name}' reported: " + f"{cleanup_payload['error']}" + ) + except ProgrammingError as e: + cli_console.warning( + f"Cleanup of validation service '{service_name}' failed: {e}" + ) + + return _classify_validation_status(final_status, service_name, workload_type) + + +def _poll_then_stream_validation_logs( + service_manager: ServiceManager, service_name: str +) -> Optional[str]: + """Mirror the build-image polling/streaming loop for the validation service.""" + object_manager = ObjectManager() + fqn = FQN.from_string(service_name) + poll_interval = 5 + max_wait_time = 300 + elapsed = 0 + current_status: Optional[str] = None + while elapsed < max_wait_time: + try: + describe_result = object_manager.describe( + object_type="service", fqn=fqn, cursor_class=DictCursor + ) + row = describe_result.fetchone() + if row and "status" in row: + current_status = row["status"] + if current_status and current_status != "PENDING": + break + except Exception: + cli_console.message( + f"Waiting for validation service to be available... ({elapsed}s elapsed)" + ) + time.sleep(poll_interval) + elapsed += poll_interval + + if current_status is None or current_status == "PENDING": + cli_console.warning( + f"Validation service did not leave PENDING within {max_wait_time}s" + f" (current: {current_status or 'UNKNOWN'})" + ) + return None + + cli_console.step(f"Validation service status: {current_status}. Streaming logs...") + final_status: Optional[str] = None + log_stream = service_manager.stream_logs( + service_name=service_name, + instance_id="0", + container_name="main", + num_lines=1000, + since_timestamp="", + include_timestamps=False, + interval_seconds=2, + check_terminal_status=True, + ) + for log_entry in log_stream: + if isinstance(log_entry, tuple) and log_entry[0] == "__TERMINAL_STATUS__": + final_status = log_entry[1] + break + cli_console.message(log_entry) + + return final_status or current_status + + +def _classify_validation_status( + final_status: Optional[str], service_name: str, workload_type: str +) -> dict: + if final_status == "DONE": + return { + "status": _VALIDATION_STATUS_PASS, + "service_name": service_name, + "workload_type": workload_type, + } + if final_status == "FAILED": + return { + "status": _VALIDATION_STATUS_FAIL, + "service_name": service_name, + "workload_type": workload_type, + "reason": "validation service reached FAILED state", + } + # CANCELLED, INTERNAL_ERROR, timed-out, no-status -- transient. + return { + "status": _VALIDATION_STATUS_TRANSIENT, + "service_name": service_name, + "workload_type": workload_type, + "reason": ( + f"validation service did not reach DONE (final={final_status or 'UNKNOWN'})" + ), + } + + +def _report_validation_result(result: dict) -> None: + status = result.get("status") + workload_type = result.get("workload_type", "?") + if status == _VALIDATION_STATUS_PASS: + cli_console.message( + f"✓ Image validation passed (workload: {workload_type})." + ) + elif status == _VALIDATION_STATUS_FAIL: + cli_console.warning( + f"✗ Image validation FAILED (workload: {workload_type}): {result.get('reason', '')}" + ) + elif status == _VALIDATION_STATUS_TRANSIENT: + cli_console.warning( + f"! Image validation could not run to completion (workload:" + f" {workload_type}): {result.get('reason', '')}" + ) + elif status == _VALIDATION_STATUS_SKIPPED: + cli_console.message( + f"Image validation skipped (workload: {workload_type}): {result.get('reason', '')}" + ) diff --git a/src/snowflake/cli/_plugins/spcs/services/manager.py b/src/snowflake/cli/_plugins/spcs/services/manager.py index 4a20077bcd..543419310a 100644 --- a/src/snowflake/cli/_plugins/spcs/services/manager.py +++ b/src/snowflake/cli/_plugins/spcs/services/manager.py @@ -46,6 +46,7 @@ from snowflake.cli.api.constants import DEFAULT_SIZE_LIMIT_MB, ObjectType from snowflake.cli.api.identifiers import FQN from snowflake.cli.api.project.schemas.entities.common import Artifacts +from snowflake.cli.api.project.util import to_string_literal from snowflake.cli.api.secure_path import SecurePath from snowflake.cli.api.sql_execution import SqlExecutionMixin from snowflake.cli.api.stage_path import StagePath @@ -304,50 +305,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 +339,76 @@ 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, + workload_type: Optional[str] = None, + skip_validation: bool = False, + ) -> 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 + workload_type: Optional workload identifier (e.g. "notebook", "ml_job") + that selects which validation YAML the sf-image-build runner + applies. Propagated as the WORKLOAD_TYPE env var; the runner + resolves it to /etc/sf-image-build/configs/.yaml and + fails fast if no such file is baked into the runner image. + When None, the WORKLOAD_TYPE env var is omitted, which causes + the runner to skip validation entirely (with a logged warning). + skip_validation: When True, propagated as ``SKIP_VALIDATION=true`` on + the runner job's env. The runner reads this flag and bypasses the + advisory in-process validation phases entirely. Independent of + ``workload_type`` -- callers can still pass a workload type for + the post-push smoke test path while skipping the runner-side + advisory checks. + """ + 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", + } + if workload_type: + # Opts the runner into validation against + # /etc/sf-image-build/configs/.yaml. Setting an + # unknown workload type causes the runner to fail fast with a + # message listing every YAML baked into the runner image. + env_vars["WORKLOAD_TYPE"] = workload_type + if skip_validation: + env_vars["SKIP_VALIDATION"] = "true" + 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", @@ -410,6 +437,61 @@ def build_image( async_mode=async_mode, ) + def start_image_validation_service( + self, image_url: str, workload_type: str + ) -> dict: + """ + Start a server-side image validation smoke service. + + Wraps ``SELECT SYSTEM$START_IMAGE_VALIDATION_SERVICE(?, ?)`` and parses its JSON + return value. The system function provisions a per-service stage with the + workload-specific smoke driver, submits an SPCS service against ``image_url``, and + returns ``{"status": "started"|"skipped", "service_name": ..., "schema_fqn": ..., + "stage_fqn": ..., "workload_type": ...}``. Callers tail the service's logs via + ``stream_logs`` and call ``cleanup_image_validation_service`` once the service + reaches a terminal state. + """ + query = ( + "SELECT PARSE_JSON(SYSTEM$START_IMAGE_VALIDATION_SERVICE(" + f"{to_string_literal(image_url)}, {to_string_literal(workload_type)}" + ")) AS payload" + ) + cursor = self.execute_query(query) + row = cursor.fetchone() + if row is None: + raise ProgrammingError( + "SYSTEM$START_IMAGE_VALIDATION_SERVICE returned no rows; check the" + " session has CURRENT_DATABASE/CURRENT_SCHEMA set and the workload" + " type is supported." + ) + payload = row[0] + if isinstance(payload, str): + payload = json.loads(payload) + return payload + + def cleanup_image_validation_service(self, service_name: str) -> dict: + """ + Drop the validation service and its per-service payload stage. + + Wraps ``SELECT SYSTEM$CLEANUP_IMAGE_VALIDATION_SERVICE(?)`` and parses its JSON + return value. Returns ``{"dropped_service": bool, "dropped_stage": bool, "error": + str|None}``. Designed to be called from a ``finally`` block so a Ctrl-C or + unexpected exception during log streaming still triggers cleanup. + """ + query = ( + "SELECT PARSE_JSON(SYSTEM$CLEANUP_IMAGE_VALIDATION_SERVICE(" + f"{to_string_literal(service_name)}" + ")) AS payload" + ) + cursor = self.execute_query(query) + row = cursor.fetchone() + if row is None: + return {"dropped_service": False, "dropped_stage": False, "error": "no rows returned"} + payload = row[0] + if isinstance(payload, str): + payload = json.loads(payload) + return payload + def _read_yaml(self, path: Path) -> str: # TODO(aivanou): Add validation towards schema with SecurePath(path).open("r", read_file_limit_mb=DEFAULT_SIZE_LIMIT_MB) as fh: @@ -711,15 +793,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..56ac49a015 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,10 +1935,54 @@ 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" + # WORKLOAD_TYPE must NOT be set when the caller did not pass it: the + # runner treats absence as "skip validation entirely" and we never want + # to silently opt callers into a wrong ruleset. + assert "WORKLOAD_TYPE" not in container["env"] assert result == cursor +@patch("snowflake.cli._plugins.connection.util.get_account") +@patch(EXECUTE_QUERY) +def test_build_image_with_workload_type(mock_execute_query, mock_get_account): + """Passing workload_type propagates to the WORKLOAD_TYPE env var.""" + 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" + + cursor = Mock(spec=SnowflakeCursor) + mock_execute_query.return_value = cursor + + manager = ServiceManager(connection=mock_conn) + manager.build_image( + job_service_name="test_build_job", + compute_pool="test_pool", + image_repository="db.schema.repo", + image_name="my_image", + image_tag="v1.0", + stage="test_stage", + build_context_path="ctx", + external_access_integrations=None, + async_mode=True, + workload_type="notebook", + ) + + spec_arg = mock_execute_query.mock_calls[0].args[0] + import re + + match = re.search(r"\$\$\s*(\{.*?\})\s*\$\$", spec_arg, re.DOTALL) + assert match is not None + spec = json.loads(match.group(1)) + env = spec["spec"]["containers"][0]["env"] + assert env["WORKLOAD_TYPE"] == "notebook" + + @patch("snowflake.cli._plugins.connection.util.get_account") def test_build_image_repository_validation(mock_get_account): """Test build_image validates image repository format. @@ -2191,6 +2171,220 @@ 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 workload type (path-traversal-style value rejected before + # we ever shell anything out, in addition to the runner's fail-fast). + 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), + "--workload-type", + "../etc/passwd", + ], + catch_exceptions=False, + ) + assert result.exit_code != 0 + assert "Invalid workload type" in result.output + + # Test 5: 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 +2625,504 @@ 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 + + +# Tests for the post-push image-validation flow. +# Layered as: build_image_registry_url helper -> start/cleanup helpers -> the +# CLI integration that wires them together. The integration tests use the same +# build-image runner harness as test_build_image_cli_recursive_upload_with_nested_dirs. + + +@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") + + +@patch("snowflake.cli._plugins.connection.util.get_account") +@patch(EXECUTE_QUERY) +def test_build_image_skip_validation_propagates_env(mock_execute_query, mock_get_account): + """--skip-validation must inject SKIP_VALIDATION=true on the runner job env.""" + 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" + + cursor = Mock(spec=SnowflakeCursor) + mock_execute_query.return_value = cursor + + manager = ServiceManager(connection=mock_conn) + manager.build_image( + job_service_name="job", + compute_pool="pool", + image_repository="db.schema.repo", + image_name="img", + image_tag="v1", + stage="stg", + build_context_path="ctx", + external_access_integrations=None, + async_mode=True, + workload_type="ml_job", + skip_validation=True, + ) + + spec_arg = mock_execute_query.mock_calls[0].args[0] + match = re.search(r"\$\$\s*(\{.*?\})\s*\$\$", spec_arg, re.DOTALL) + assert match is not None + env = json.loads(match.group(1))["spec"]["containers"][0]["env"] + assert env["SKIP_VALIDATION"] == "true" + # Workload type still propagates so the smoke test (post-push) can be invoked + # if the user later re-runs without --skip-validation; the runner side just + # honours SKIP_VALIDATION=true and bypasses its own checks. + assert env["WORKLOAD_TYPE"] == "ml_job" + + +@patch(EXECUTE_QUERY) +def test_start_image_validation_service_parses_payload(mock_execute_query): + cursor = Mock(spec=SnowflakeCursor) + cursor.fetchone.return_value = ( + '{"status":"started","service_name":"IMG123","schema_fqn":"DB.SCH",' + '"stage_fqn":"@DB.SCH._SNOWFLAKE_IMAGE_VALIDATION_PAYLOAD_IMG123",' + '"workload_type":"ML_JOB"}', + ) + mock_execute_query.return_value = cursor + + manager = ServiceManager(connection=Mock()) + payload = manager.start_image_validation_service( + image_url="org-acct.registry-local.snowflakecomputing.com/db/sch/repo/img:v1", + workload_type="ML_JOB", + ) + + assert payload["status"] == "started" + assert payload["service_name"] == "IMG123" + + sql = mock_execute_query.mock_calls[0].args[0] + assert "SYSTEM$START_IMAGE_VALIDATION_SERVICE" in sql + assert "PARSE_JSON" in sql + # to_string_literal wraps in single quotes + assert "'org-acct.registry-local.snowflakecomputing.com/db/sch/repo/img:v1'" in sql + assert "'ML_JOB'" in sql + + +@patch(EXECUTE_QUERY) +def test_start_image_validation_service_skipped(mock_execute_query): + cursor = Mock(spec=SnowflakeCursor) + cursor.fetchone.return_value = ( + '{"status":"skipped","workload_type":"ML_JOB"}', + ) + mock_execute_query.return_value = cursor + + manager = ServiceManager(connection=Mock()) + payload = manager.start_image_validation_service( + image_url="org-acct.registry-local.snowflakecomputing.com/db/sch/repo/img:v1", + workload_type="ML_JOB", + ) + assert payload["status"] == "skipped" + + +@patch(EXECUTE_QUERY) +def test_cleanup_image_validation_service_parses_payload(mock_execute_query): + cursor = Mock(spec=SnowflakeCursor) + cursor.fetchone.return_value = ( + '{"dropped_service":true,"dropped_stage":true}', + ) + mock_execute_query.return_value = cursor + + manager = ServiceManager(connection=Mock()) + payload = manager.cleanup_image_validation_service(service_name="IMG123") + + assert payload == {"dropped_service": True, "dropped_stage": True} + sql = mock_execute_query.mock_calls[0].args[0] + assert "SYSTEM$CLEANUP_IMAGE_VALIDATION_SERVICE" in sql + assert "'IMG123'" in sql + + +@patch(EXECUTE_QUERY) +def test_cleanup_image_validation_service_no_rows(mock_execute_query): + cursor = Mock(spec=SnowflakeCursor) + cursor.fetchone.return_value = None + mock_execute_query.return_value = cursor + + manager = ServiceManager(connection=Mock()) + payload = manager.cleanup_image_validation_service(service_name="IMG_GONE") + assert payload["dropped_service"] is False + assert payload["dropped_stage"] is False + assert "no rows returned" in payload["error"] + + +def _make_validation_runner_mocks( + *, + start_payload: dict, + cleanup_payload: dict, + describe_status_seq: list, + stream_log_seq: list, +): + """Shared harness for the build-image CLI validation-flow tests. + + Builds the same Mock-graph as test_build_image_cli_recursive_upload_with_nested_dirs + but additionally lets each test choose how the validation start, describe, stream, + and cleanup calls behave. + """ + mock_service_manager = Mock() + 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.side_effect = [ + iter([("__TERMINAL_STATUS__", "DONE")]), # build-image stream + iter(stream_log_seq), # validation stream + ] + mock_service_manager.build_image_registry_url.return_value = ( + "org-acct.registry-local.snowflakecomputing.com/db/sch/repo" + ) + mock_service_manager.start_image_validation_service.return_value = start_payload + mock_service_manager.cleanup_image_validation_service.return_value = cleanup_payload + return mock_service_manager, describe_status_seq + + +@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") +def test_build_image_cli_validation_pass( + mock_stage_execute_query, + mock_stage_put, + mock_service_manager_class, + mock_object_manager_class, + mock_sleep, + runner, + temporary_directory, +): + """When the validation service hits DONE the CLI exits 0 and runs cleanup once.""" + build_context = Path(temporary_directory) / "build_context" + build_context.mkdir() + (build_context / "Dockerfile").write_text("FROM python:3.10-alpine") + mock_stage_put.return_value = Mock(fetchall=lambda: []) + + mock_service_manager, _ = _make_validation_runner_mocks( + start_payload={ + "status": "started", + "service_name": "IMG123", + "schema_fqn": "DB.SCH", + "stage_fqn": "@DB.SCH._SNOWFLAKE_IMAGE_VALIDATION_PAYLOAD_IMG123", + "workload_type": "ML_JOB", + }, + cleanup_payload={"dropped_service": True, "dropped_stage": True}, + describe_status_seq=[{"status": "RUNNING"}, {"status": "RUNNING"}], + stream_log_seq=[ + "validation log line 1", + ("__TERMINAL_STATUS__", "DONE"), + ], + ) + mock_service_manager_class.return_value = mock_service_manager + + mock_object_manager = Mock() + mock_object_manager_class.return_value = mock_object_manager + # Two describe paths: the build job (which the CLI polls first) and the + # validation service. Both leave PENDING immediately. + describe_cursor = Mock() + describe_cursor.fetchone.return_value = {"status": "RUNNING"} + mock_object_manager.describe.return_value = 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", + "--workload-type", + "ml_job", + ], + catch_exceptions=False, + ) + + assert result.exit_code == 0, f"unexpected exit; output: {result.output}" + mock_service_manager.start_image_validation_service.assert_called_once() + mock_service_manager.cleanup_image_validation_service.assert_called_once_with( + service_name="IMG123" + ) + assert "Image validation passed" in result.output + + +@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") +def test_build_image_cli_validation_failed( + mock_stage_execute_query, + mock_stage_put, + mock_service_manager_class, + mock_object_manager_class, + mock_sleep, + runner, + temporary_directory, +): + """When the validation service hits FAILED the CLI exits 2 (validation fail).""" + build_context = Path(temporary_directory) / "build_context" + build_context.mkdir() + (build_context / "Dockerfile").write_text("FROM python:3.10-alpine") + mock_stage_put.return_value = Mock(fetchall=lambda: []) + + mock_service_manager, _ = _make_validation_runner_mocks( + start_payload={ + "status": "started", + "service_name": "IMG456", + "schema_fqn": "DB.SCH", + "stage_fqn": "@DB.SCH._SNOWFLAKE_IMAGE_VALIDATION_PAYLOAD_IMG456", + "workload_type": "ML_JOB", + }, + cleanup_payload={"dropped_service": True, "dropped_stage": True}, + describe_status_seq=[{"status": "RUNNING"}], + stream_log_seq=[("__TERMINAL_STATUS__", "FAILED")], + ) + mock_service_manager_class.return_value = mock_service_manager + mock_object_manager = Mock() + mock_object_manager_class.return_value = mock_object_manager + describe_cursor = Mock() + describe_cursor.fetchone.return_value = {"status": "RUNNING"} + mock_object_manager.describe.return_value = 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", + "--workload-type", + "ml_job", + ], + catch_exceptions=False, + ) + + # _EXIT_VALIDATION_FAILED == 2 + assert result.exit_code == 2, f"unexpected exit; output: {result.output}" + mock_service_manager.cleanup_image_validation_service.assert_called_once_with( + service_name="IMG456" + ) + assert "Image validation FAILED" in result.output + + +@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") +def test_build_image_cli_skip_validation( + mock_stage_execute_query, + mock_stage_put, + mock_service_manager_class, + mock_object_manager_class, + mock_sleep, + runner, + temporary_directory, +): + """--skip-validation propagates SKIP_VALIDATION=true AND skips the post-push smoke test.""" + build_context = Path(temporary_directory) / "build_context" + build_context.mkdir() + (build_context / "Dockerfile").write_text("FROM python:3.10-alpine") + mock_stage_put.return_value = Mock(fetchall=lambda: []) + + mock_service_manager = Mock() + 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_service_manager_class.return_value = mock_service_manager + + mock_object_manager = Mock() + mock_object_manager_class.return_value = mock_object_manager + describe_cursor = Mock() + describe_cursor.fetchone.return_value = {"status": "RUNNING"} + mock_object_manager.describe.return_value = 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", + "--workload-type", + "ml_job", + "--skip-validation", + ], + catch_exceptions=False, + ) + + assert result.exit_code == 0, f"unexpected exit; output: {result.output}" + # The runner-side flag was passed to build_image: + build_image_kwargs = mock_service_manager.build_image.call_args.kwargs + assert build_image_kwargs["skip_validation"] is True + assert build_image_kwargs["workload_type"] == "ml_job" + # And the post-push smoke test was NOT invoked: + mock_service_manager.start_image_validation_service.assert_not_called() + mock_service_manager.cleanup_image_validation_service.assert_not_called() + + +@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") +def test_build_image_cli_validation_keyboard_interrupt( + mock_stage_execute_query, + mock_stage_put, + mock_service_manager_class, + mock_object_manager_class, + mock_sleep, + runner, + temporary_directory, +): + """Ctrl-C during log streaming still triggers cleanup and reports transient.""" + build_context = Path(temporary_directory) / "build_context" + build_context.mkdir() + (build_context / "Dockerfile").write_text("FROM python:3.10-alpine") + mock_stage_put.return_value = Mock(fetchall=lambda: []) + + mock_service_manager = Mock() + 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 + + def _stream_then_interrupt(): + # Build-image stream returns DONE quickly, then validation stream raises. + if not getattr(_stream_then_interrupt, "called_once", False): + _stream_then_interrupt.called_once = True + return iter([("__TERMINAL_STATUS__", "DONE")]) + + def gen(): + yield "validation log line" + raise KeyboardInterrupt() + + return gen() + + mock_service_manager.stream_logs.side_effect = lambda *a, **kw: _stream_then_interrupt() + mock_service_manager.build_image_registry_url.return_value = ( + "org-acct.registry-local.snowflakecomputing.com/db/sch/repo" + ) + mock_service_manager.start_image_validation_service.return_value = { + "status": "started", + "service_name": "IMG_INTR", + "schema_fqn": "DB.SCH", + "stage_fqn": "@DB.SCH._SNOWFLAKE_IMAGE_VALIDATION_PAYLOAD_IMG_INTR", + "workload_type": "ML_JOB", + } + mock_service_manager.cleanup_image_validation_service.return_value = { + "dropped_service": True, + "dropped_stage": True, + } + mock_service_manager_class.return_value = mock_service_manager + + mock_object_manager = Mock() + mock_object_manager_class.return_value = mock_object_manager + describe_cursor = Mock() + describe_cursor.fetchone.return_value = {"status": "RUNNING"} + mock_object_manager.describe.return_value = 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", + "--workload-type", + "ml_job", + ], + catch_exceptions=False, + ) + + # Ctrl-C => transient (exit 3); cleanup must still run exactly once. + assert result.exit_code == 3, f"unexpected exit; output: {result.output}" + mock_service_manager.cleanup_image_validation_service.assert_called_once_with( + service_name="IMG_INTR" + ) + assert "could not run to completion" in result.output From 4992e851042c58bc5c0aae7c64afa4313250f146 Mon Sep 17 00:00:00 2001 From: Supriya Vasudevan Date: Mon, 11 May 2026 19:23:24 +0000 Subject: [PATCH 2/2] Fix unintednded changes Description Testing --- .../cli/_plugins/spcs/services/commands.py | 269 --------- .../cli/_plugins/spcs/services/manager.py | 79 --- tests/spcs/test_services.py | 536 +----------------- 3 files changed, 2 insertions(+), 882 deletions(-) diff --git a/src/snowflake/cli/_plugins/spcs/services/commands.py b/src/snowflake/cli/_plugins/spcs/services/commands.py index 9dd5abf4c5..ab23a46c99 100644 --- a/src/snowflake/cli/_plugins/spcs/services/commands.py +++ b/src/snowflake/cli/_plugins/spcs/services/commands.py @@ -80,20 +80,6 @@ short_help="Manages services.", ) -# Exit codes used by the post-push image-validation hook in build-image. -# Distinct from the default Click "command failed" exit (1) so callers / CI can -# tell apart "image push failed" (1), "image was pushed but workload-image -# smoke test reported the image is broken" (_EXIT_VALIDATION_FAILED), and -# "smoke test could not run to completion (transient infra)" (_EXIT_VALIDATION_TRANSIENT). -_EXIT_VALIDATION_FAILED = 2 -_EXIT_VALIDATION_TRANSIENT = 3 - -# Status strings the validator uses to label terminal outcomes for the user. -_VALIDATION_STATUS_PASS = "pass" -_VALIDATION_STATUS_FAIL = "fail" -_VALIDATION_STATUS_TRANSIENT = "transient" -_VALIDATION_STATUS_SKIPPED = "skipped" - # Define common options container_name_option = typer.Option( ..., @@ -720,32 +706,6 @@ def build_image( "--eai-name", help="Identifies external access integrations (EAI) that the build job can access. This option may be specified multiple times for multiple EAIs.", ), - workload_type: Optional[str] = typer.Option( - None, - "--workload-type", - help=( - "Workload identifier (e.g. 'notebook', 'ml_job') that selects " - "which validation YAML the sf-image-build runner applies and, " - "after a successful push, which server-side smoke test the CLI " - "asks for via SYSTEM$START_IMAGE_VALIDATION_SERVICE. Propagated " - "to the build job as the WORKLOAD_TYPE env var. If omitted, both " - "the runner-side advisory checks and the post-push smoke test " - "are skipped." - ), - show_default=False, - ), - skip_validation: bool = typer.Option( - False, - "--skip-validation", - help=( - "Skip both the runner-side advisory validation and the post-push " - "server-side smoke test. The CLI propagates SKIP_VALIDATION=true " - "into the build job's env (which the sf-image-build runner reads " - "to bypass its in-process checks) and short-circuits before the " - "smoke-test system function is called. Image is built and pushed " - "as usual." - ), - ), **options, ) -> CommandResult: """ @@ -795,18 +755,6 @@ def build_image( f"Invalid image tag '{image_tag}'. Must contain only alphanumeric characters, hyphens, underscores, and dots." ) - # Defensive validation: workload_type becomes a JSON env value that is then - # used by the runner to construct /etc/sf-image-build/configs/.yaml. - # Restricting to alphanumerics + _ - . prevents path traversal and shell - # metacharacters from sneaking through. The runner validates further by - # rejecting any value that does not map to an existing baked-in YAML. - if workload_type is not None and not workload_type.replace("_", "").replace( - "-", "" - ).replace(".", "").isalnum(): - raise CliArgumentError( - f"Invalid workload type '{workload_type}'. Must contain only alphanumeric characters, hyphens, underscores, and dots." - ) - # Generate a unique identifier for this build operation build_uuid = uuid.uuid4().hex[:8] @@ -878,8 +826,6 @@ def build_image( build_context_path=build_context_stage_path, external_access_integrations=external_access_integrations, async_mode=True, # Always async so we can stream logs - workload_type=workload_type, - skip_validation=skip_validation, ) cli_console.step(f"Waiting for job to start...") @@ -1023,219 +969,4 @@ def build_image( except ProgrammingError as e: cli_console.warning(f"Failed to clean up build context files: {e}") - # Post-push image validation: only if (a) the build job actually finished - # successfully (otherwise the image isn't even there to validate), (b) the - # caller opted in via --workload-type, and (c) they didn't explicitly - # disable validation via --skip-validation. - if ( - final_status == "DONE" - and workload_type is not None - and not skip_validation - ): - validation_result = _run_image_validation( - service_manager=service_manager, - image_repository=image_repository, - image_name=image_name, - image_tag=image_tag, - workload_type=workload_type, - ) - _report_validation_result(validation_result) - status = validation_result.get("status") - if status == _VALIDATION_STATUS_FAIL: - raise typer.Exit(_EXIT_VALIDATION_FAILED) - if status == _VALIDATION_STATUS_TRANSIENT: - raise typer.Exit(_EXIT_VALIDATION_TRANSIENT) - return SingleQueryResult(cursor) - - -def _run_image_validation( - service_manager: ServiceManager, - image_repository: str, - image_name: str, - image_tag: str, - workload_type: str, -) -> dict: - """ - Server-orchestrated post-push smoke test. - - Calls SYSTEM$START_IMAGE_VALIDATION_SERVICE to provision a per-workload smoke - service against the just-pushed image, polls DESCRIBE SERVICE until it leaves - PENDING, streams its logs to the terminal, classifies the terminal status, and - always calls SYSTEM$CLEANUP_IMAGE_VALIDATION_SERVICE in a finally so a Ctrl-C - or unexpected exception still triggers cleanup. - - Returns ``{"status": "pass"|"fail"|"transient"|"skipped", ...}`` so the caller - can translate to a CLI exit code. - """ - image_url = service_manager.build_image_registry_url( - image_repository=image_repository, - ) - image_url_with_tag = f"{image_url}/{image_name}:{image_tag}" - - cli_console.step( - f"Starting server-side image validation for workload type '{workload_type}'..." - ) - try: - start_payload = service_manager.start_image_validation_service( - image_url=image_url_with_tag, - workload_type=workload_type, - ) - except ProgrammingError as e: - # Treat any infra-level error from the start function as transient: the - # image was pushed, but we couldn't even kick off validation. The user - # can re-run validation manually. - cli_console.warning(f"Could not start image validation: {e}") - return { - "status": _VALIDATION_STATUS_TRANSIENT, - "reason": str(e), - "workload_type": workload_type, - } - - if start_payload.get("status") == "skipped": - return { - "status": _VALIDATION_STATUS_SKIPPED, - "reason": start_payload.get("reason") - or f"validation not enabled for workload {workload_type}", - "workload_type": workload_type, - } - - service_name = start_payload.get("service_name") - if not service_name: - return { - "status": _VALIDATION_STATUS_TRANSIENT, - "reason": "validation start response missing service_name", - "workload_type": workload_type, - "raw": start_payload, - } - - cli_console.message(f"Validation service: {service_name}") - - final_status: Optional[str] = None - try: - final_status = _poll_then_stream_validation_logs(service_manager, service_name) - except KeyboardInterrupt: - cli_console.warning( - f"\nImage validation '{service_name}' interrupted; attempting cleanup." - ) - return _classify_validation_status("CANCELLED", service_name, workload_type) - finally: - try: - cleanup_payload = service_manager.cleanup_image_validation_service( - service_name=service_name, - ) - if cleanup_payload.get("error"): - cli_console.warning( - f"Cleanup of validation service '{service_name}' reported: " - f"{cleanup_payload['error']}" - ) - except ProgrammingError as e: - cli_console.warning( - f"Cleanup of validation service '{service_name}' failed: {e}" - ) - - return _classify_validation_status(final_status, service_name, workload_type) - - -def _poll_then_stream_validation_logs( - service_manager: ServiceManager, service_name: str -) -> Optional[str]: - """Mirror the build-image polling/streaming loop for the validation service.""" - object_manager = ObjectManager() - fqn = FQN.from_string(service_name) - poll_interval = 5 - max_wait_time = 300 - elapsed = 0 - current_status: Optional[str] = None - while elapsed < max_wait_time: - try: - describe_result = object_manager.describe( - object_type="service", fqn=fqn, cursor_class=DictCursor - ) - row = describe_result.fetchone() - if row and "status" in row: - current_status = row["status"] - if current_status and current_status != "PENDING": - break - except Exception: - cli_console.message( - f"Waiting for validation service to be available... ({elapsed}s elapsed)" - ) - time.sleep(poll_interval) - elapsed += poll_interval - - if current_status is None or current_status == "PENDING": - cli_console.warning( - f"Validation service did not leave PENDING within {max_wait_time}s" - f" (current: {current_status or 'UNKNOWN'})" - ) - return None - - cli_console.step(f"Validation service status: {current_status}. Streaming logs...") - final_status: Optional[str] = None - log_stream = service_manager.stream_logs( - service_name=service_name, - instance_id="0", - container_name="main", - num_lines=1000, - since_timestamp="", - include_timestamps=False, - interval_seconds=2, - check_terminal_status=True, - ) - for log_entry in log_stream: - if isinstance(log_entry, tuple) and log_entry[0] == "__TERMINAL_STATUS__": - final_status = log_entry[1] - break - cli_console.message(log_entry) - - return final_status or current_status - - -def _classify_validation_status( - final_status: Optional[str], service_name: str, workload_type: str -) -> dict: - if final_status == "DONE": - return { - "status": _VALIDATION_STATUS_PASS, - "service_name": service_name, - "workload_type": workload_type, - } - if final_status == "FAILED": - return { - "status": _VALIDATION_STATUS_FAIL, - "service_name": service_name, - "workload_type": workload_type, - "reason": "validation service reached FAILED state", - } - # CANCELLED, INTERNAL_ERROR, timed-out, no-status -- transient. - return { - "status": _VALIDATION_STATUS_TRANSIENT, - "service_name": service_name, - "workload_type": workload_type, - "reason": ( - f"validation service did not reach DONE (final={final_status or 'UNKNOWN'})" - ), - } - - -def _report_validation_result(result: dict) -> None: - status = result.get("status") - workload_type = result.get("workload_type", "?") - if status == _VALIDATION_STATUS_PASS: - cli_console.message( - f"✓ Image validation passed (workload: {workload_type})." - ) - elif status == _VALIDATION_STATUS_FAIL: - cli_console.warning( - f"✗ Image validation FAILED (workload: {workload_type}): {result.get('reason', '')}" - ) - elif status == _VALIDATION_STATUS_TRANSIENT: - cli_console.warning( - f"! Image validation could not run to completion (workload:" - f" {workload_type}): {result.get('reason', '')}" - ) - elif status == _VALIDATION_STATUS_SKIPPED: - cli_console.message( - f"Image validation skipped (workload: {workload_type}): {result.get('reason', '')}" - ) diff --git a/src/snowflake/cli/_plugins/spcs/services/manager.py b/src/snowflake/cli/_plugins/spcs/services/manager.py index 543419310a..fbc5dea174 100644 --- a/src/snowflake/cli/_plugins/spcs/services/manager.py +++ b/src/snowflake/cli/_plugins/spcs/services/manager.py @@ -46,7 +46,6 @@ from snowflake.cli.api.constants import DEFAULT_SIZE_LIMIT_MB, ObjectType from snowflake.cli.api.identifiers import FQN from snowflake.cli.api.project.schemas.entities.common import Artifacts -from snowflake.cli.api.project.util import to_string_literal from snowflake.cli.api.secure_path import SecurePath from snowflake.cli.api.sql_execution import SqlExecutionMixin from snowflake.cli.api.stage_path import StagePath @@ -355,8 +354,6 @@ def build_image( build_context_path: str, external_access_integrations: Optional[List[str]], async_mode: bool = True, - workload_type: Optional[str] = None, - skip_validation: bool = False, ) -> SnowflakeCursor: """ Builds a container image using EXECUTE JOB SERVICE with the Snowflake image builder. @@ -371,19 +368,6 @@ def build_image( 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 - workload_type: Optional workload identifier (e.g. "notebook", "ml_job") - that selects which validation YAML the sf-image-build runner - applies. Propagated as the WORKLOAD_TYPE env var; the runner - resolves it to /etc/sf-image-build/configs/.yaml and - fails fast if no such file is baked into the runner image. - When None, the WORKLOAD_TYPE env var is omitted, which causes - the runner to skip validation entirely (with a logged warning). - skip_validation: When True, propagated as ``SKIP_VALIDATION=true`` on - the runner job's env. The runner reads this flag and bypasses the - advisory in-process validation phases entirely. Independent of - ``workload_type`` -- callers can still pass a workload type for - the post-push smoke test path while skipping the runner-side - advisory checks. """ image_registry_url = self.build_image_registry_url(image_repository) @@ -393,14 +377,6 @@ def build_image( "IMAGE_TAG": image_tag, "BUILD_CONTEXT": "/app", } - if workload_type: - # Opts the runner into validation against - # /etc/sf-image-build/configs/.yaml. Setting an - # unknown workload type causes the runner to fail fast with a - # message listing every YAML baked into the runner image. - env_vars["WORKLOAD_TYPE"] = workload_type - if skip_validation: - env_vars["SKIP_VALIDATION"] = "true" spec = { "spec": { @@ -437,61 +413,6 @@ def build_image( async_mode=async_mode, ) - def start_image_validation_service( - self, image_url: str, workload_type: str - ) -> dict: - """ - Start a server-side image validation smoke service. - - Wraps ``SELECT SYSTEM$START_IMAGE_VALIDATION_SERVICE(?, ?)`` and parses its JSON - return value. The system function provisions a per-service stage with the - workload-specific smoke driver, submits an SPCS service against ``image_url``, and - returns ``{"status": "started"|"skipped", "service_name": ..., "schema_fqn": ..., - "stage_fqn": ..., "workload_type": ...}``. Callers tail the service's logs via - ``stream_logs`` and call ``cleanup_image_validation_service`` once the service - reaches a terminal state. - """ - query = ( - "SELECT PARSE_JSON(SYSTEM$START_IMAGE_VALIDATION_SERVICE(" - f"{to_string_literal(image_url)}, {to_string_literal(workload_type)}" - ")) AS payload" - ) - cursor = self.execute_query(query) - row = cursor.fetchone() - if row is None: - raise ProgrammingError( - "SYSTEM$START_IMAGE_VALIDATION_SERVICE returned no rows; check the" - " session has CURRENT_DATABASE/CURRENT_SCHEMA set and the workload" - " type is supported." - ) - payload = row[0] - if isinstance(payload, str): - payload = json.loads(payload) - return payload - - def cleanup_image_validation_service(self, service_name: str) -> dict: - """ - Drop the validation service and its per-service payload stage. - - Wraps ``SELECT SYSTEM$CLEANUP_IMAGE_VALIDATION_SERVICE(?)`` and parses its JSON - return value. Returns ``{"dropped_service": bool, "dropped_stage": bool, "error": - str|None}``. Designed to be called from a ``finally`` block so a Ctrl-C or - unexpected exception during log streaming still triggers cleanup. - """ - query = ( - "SELECT PARSE_JSON(SYSTEM$CLEANUP_IMAGE_VALIDATION_SERVICE(" - f"{to_string_literal(service_name)}" - ")) AS payload" - ) - cursor = self.execute_query(query) - row = cursor.fetchone() - if row is None: - return {"dropped_service": False, "dropped_stage": False, "error": "no rows returned"} - payload = row[0] - if isinstance(payload, str): - payload = json.loads(payload) - return payload - def _read_yaml(self, path: Path) -> str: # TODO(aivanou): Add validation towards schema with SecurePath(path).open("r", read_file_limit_mb=DEFAULT_SIZE_LIMIT_MB) as fh: diff --git a/tests/spcs/test_services.py b/tests/spcs/test_services.py index 56ac49a015..77533a001a 100644 --- a/tests/spcs/test_services.py +++ b/tests/spcs/test_services.py @@ -1935,54 +1935,12 @@ 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" - # WORKLOAD_TYPE must NOT be set when the caller did not pass it: the - # runner treats absence as "skip validation entirely" and we never want - # to silently opt callers into a wrong ruleset. assert "WORKLOAD_TYPE" not in container["env"] + assert "SKIP_VALIDATION" not in container["env"] assert result == cursor -@patch("snowflake.cli._plugins.connection.util.get_account") -@patch(EXECUTE_QUERY) -def test_build_image_with_workload_type(mock_execute_query, mock_get_account): - """Passing workload_type propagates to the WORKLOAD_TYPE env var.""" - 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" - - cursor = Mock(spec=SnowflakeCursor) - mock_execute_query.return_value = cursor - - manager = ServiceManager(connection=mock_conn) - manager.build_image( - job_service_name="test_build_job", - compute_pool="test_pool", - image_repository="db.schema.repo", - image_name="my_image", - image_tag="v1.0", - stage="test_stage", - build_context_path="ctx", - external_access_integrations=None, - async_mode=True, - workload_type="notebook", - ) - - spec_arg = mock_execute_query.mock_calls[0].args[0] - import re - - match = re.search(r"\$\$\s*(\{.*?\})\s*\$\$", spec_arg, re.DOTALL) - assert match is not None - spec = json.loads(match.group(1)) - env = spec["spec"]["containers"][0]["env"] - assert env["WORKLOAD_TYPE"] == "notebook" - - @patch("snowflake.cli._plugins.connection.util.get_account") def test_build_image_repository_validation(mock_get_account): """Test build_image validates image repository format. @@ -2171,32 +2129,7 @@ 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 workload type (path-traversal-style value rejected before - # we ever shell anything out, in addition to the runner's fail-fast). - 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), - "--workload-type", - "../etc/passwd", - ], - catch_exceptions=False, - ) - assert result.exit_code != 0 - assert "Invalid workload type" in result.output - - # Test 5: Invalid stage encryption + # Test 4: Invalid stage encryption result = runner.invoke( [ "spcs", @@ -2627,12 +2560,6 @@ def test_build_image_hidden_by_default(runner): assert "build-image" not in result.output -# Tests for the post-push image-validation flow. -# Layered as: build_image_registry_url helper -> start/cleanup helpers -> the -# CLI integration that wires them together. The integration tests use the same -# build-image runner harness as test_build_image_cli_recursive_upload_with_nested_dirs. - - @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.""" @@ -2667,462 +2594,3 @@ def test_build_image_registry_url_missing_parts(mock_get_account): manager = ServiceManager(connection=mock_conn) with pytest.raises(ValueError, match="Image repository requires database and schema"): manager.build_image_registry_url("repo_only") - - -@patch("snowflake.cli._plugins.connection.util.get_account") -@patch(EXECUTE_QUERY) -def test_build_image_skip_validation_propagates_env(mock_execute_query, mock_get_account): - """--skip-validation must inject SKIP_VALIDATION=true on the runner job env.""" - 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" - - cursor = Mock(spec=SnowflakeCursor) - mock_execute_query.return_value = cursor - - manager = ServiceManager(connection=mock_conn) - manager.build_image( - job_service_name="job", - compute_pool="pool", - image_repository="db.schema.repo", - image_name="img", - image_tag="v1", - stage="stg", - build_context_path="ctx", - external_access_integrations=None, - async_mode=True, - workload_type="ml_job", - skip_validation=True, - ) - - spec_arg = mock_execute_query.mock_calls[0].args[0] - match = re.search(r"\$\$\s*(\{.*?\})\s*\$\$", spec_arg, re.DOTALL) - assert match is not None - env = json.loads(match.group(1))["spec"]["containers"][0]["env"] - assert env["SKIP_VALIDATION"] == "true" - # Workload type still propagates so the smoke test (post-push) can be invoked - # if the user later re-runs without --skip-validation; the runner side just - # honours SKIP_VALIDATION=true and bypasses its own checks. - assert env["WORKLOAD_TYPE"] == "ml_job" - - -@patch(EXECUTE_QUERY) -def test_start_image_validation_service_parses_payload(mock_execute_query): - cursor = Mock(spec=SnowflakeCursor) - cursor.fetchone.return_value = ( - '{"status":"started","service_name":"IMG123","schema_fqn":"DB.SCH",' - '"stage_fqn":"@DB.SCH._SNOWFLAKE_IMAGE_VALIDATION_PAYLOAD_IMG123",' - '"workload_type":"ML_JOB"}', - ) - mock_execute_query.return_value = cursor - - manager = ServiceManager(connection=Mock()) - payload = manager.start_image_validation_service( - image_url="org-acct.registry-local.snowflakecomputing.com/db/sch/repo/img:v1", - workload_type="ML_JOB", - ) - - assert payload["status"] == "started" - assert payload["service_name"] == "IMG123" - - sql = mock_execute_query.mock_calls[0].args[0] - assert "SYSTEM$START_IMAGE_VALIDATION_SERVICE" in sql - assert "PARSE_JSON" in sql - # to_string_literal wraps in single quotes - assert "'org-acct.registry-local.snowflakecomputing.com/db/sch/repo/img:v1'" in sql - assert "'ML_JOB'" in sql - - -@patch(EXECUTE_QUERY) -def test_start_image_validation_service_skipped(mock_execute_query): - cursor = Mock(spec=SnowflakeCursor) - cursor.fetchone.return_value = ( - '{"status":"skipped","workload_type":"ML_JOB"}', - ) - mock_execute_query.return_value = cursor - - manager = ServiceManager(connection=Mock()) - payload = manager.start_image_validation_service( - image_url="org-acct.registry-local.snowflakecomputing.com/db/sch/repo/img:v1", - workload_type="ML_JOB", - ) - assert payload["status"] == "skipped" - - -@patch(EXECUTE_QUERY) -def test_cleanup_image_validation_service_parses_payload(mock_execute_query): - cursor = Mock(spec=SnowflakeCursor) - cursor.fetchone.return_value = ( - '{"dropped_service":true,"dropped_stage":true}', - ) - mock_execute_query.return_value = cursor - - manager = ServiceManager(connection=Mock()) - payload = manager.cleanup_image_validation_service(service_name="IMG123") - - assert payload == {"dropped_service": True, "dropped_stage": True} - sql = mock_execute_query.mock_calls[0].args[0] - assert "SYSTEM$CLEANUP_IMAGE_VALIDATION_SERVICE" in sql - assert "'IMG123'" in sql - - -@patch(EXECUTE_QUERY) -def test_cleanup_image_validation_service_no_rows(mock_execute_query): - cursor = Mock(spec=SnowflakeCursor) - cursor.fetchone.return_value = None - mock_execute_query.return_value = cursor - - manager = ServiceManager(connection=Mock()) - payload = manager.cleanup_image_validation_service(service_name="IMG_GONE") - assert payload["dropped_service"] is False - assert payload["dropped_stage"] is False - assert "no rows returned" in payload["error"] - - -def _make_validation_runner_mocks( - *, - start_payload: dict, - cleanup_payload: dict, - describe_status_seq: list, - stream_log_seq: list, -): - """Shared harness for the build-image CLI validation-flow tests. - - Builds the same Mock-graph as test_build_image_cli_recursive_upload_with_nested_dirs - but additionally lets each test choose how the validation start, describe, stream, - and cleanup calls behave. - """ - mock_service_manager = Mock() - 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.side_effect = [ - iter([("__TERMINAL_STATUS__", "DONE")]), # build-image stream - iter(stream_log_seq), # validation stream - ] - mock_service_manager.build_image_registry_url.return_value = ( - "org-acct.registry-local.snowflakecomputing.com/db/sch/repo" - ) - mock_service_manager.start_image_validation_service.return_value = start_payload - mock_service_manager.cleanup_image_validation_service.return_value = cleanup_payload - return mock_service_manager, describe_status_seq - - -@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") -def test_build_image_cli_validation_pass( - mock_stage_execute_query, - mock_stage_put, - mock_service_manager_class, - mock_object_manager_class, - mock_sleep, - runner, - temporary_directory, -): - """When the validation service hits DONE the CLI exits 0 and runs cleanup once.""" - build_context = Path(temporary_directory) / "build_context" - build_context.mkdir() - (build_context / "Dockerfile").write_text("FROM python:3.10-alpine") - mock_stage_put.return_value = Mock(fetchall=lambda: []) - - mock_service_manager, _ = _make_validation_runner_mocks( - start_payload={ - "status": "started", - "service_name": "IMG123", - "schema_fqn": "DB.SCH", - "stage_fqn": "@DB.SCH._SNOWFLAKE_IMAGE_VALIDATION_PAYLOAD_IMG123", - "workload_type": "ML_JOB", - }, - cleanup_payload={"dropped_service": True, "dropped_stage": True}, - describe_status_seq=[{"status": "RUNNING"}, {"status": "RUNNING"}], - stream_log_seq=[ - "validation log line 1", - ("__TERMINAL_STATUS__", "DONE"), - ], - ) - mock_service_manager_class.return_value = mock_service_manager - - mock_object_manager = Mock() - mock_object_manager_class.return_value = mock_object_manager - # Two describe paths: the build job (which the CLI polls first) and the - # validation service. Both leave PENDING immediately. - describe_cursor = Mock() - describe_cursor.fetchone.return_value = {"status": "RUNNING"} - mock_object_manager.describe.return_value = 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", - "--workload-type", - "ml_job", - ], - catch_exceptions=False, - ) - - assert result.exit_code == 0, f"unexpected exit; output: {result.output}" - mock_service_manager.start_image_validation_service.assert_called_once() - mock_service_manager.cleanup_image_validation_service.assert_called_once_with( - service_name="IMG123" - ) - assert "Image validation passed" in result.output - - -@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") -def test_build_image_cli_validation_failed( - mock_stage_execute_query, - mock_stage_put, - mock_service_manager_class, - mock_object_manager_class, - mock_sleep, - runner, - temporary_directory, -): - """When the validation service hits FAILED the CLI exits 2 (validation fail).""" - build_context = Path(temporary_directory) / "build_context" - build_context.mkdir() - (build_context / "Dockerfile").write_text("FROM python:3.10-alpine") - mock_stage_put.return_value = Mock(fetchall=lambda: []) - - mock_service_manager, _ = _make_validation_runner_mocks( - start_payload={ - "status": "started", - "service_name": "IMG456", - "schema_fqn": "DB.SCH", - "stage_fqn": "@DB.SCH._SNOWFLAKE_IMAGE_VALIDATION_PAYLOAD_IMG456", - "workload_type": "ML_JOB", - }, - cleanup_payload={"dropped_service": True, "dropped_stage": True}, - describe_status_seq=[{"status": "RUNNING"}], - stream_log_seq=[("__TERMINAL_STATUS__", "FAILED")], - ) - mock_service_manager_class.return_value = mock_service_manager - mock_object_manager = Mock() - mock_object_manager_class.return_value = mock_object_manager - describe_cursor = Mock() - describe_cursor.fetchone.return_value = {"status": "RUNNING"} - mock_object_manager.describe.return_value = 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", - "--workload-type", - "ml_job", - ], - catch_exceptions=False, - ) - - # _EXIT_VALIDATION_FAILED == 2 - assert result.exit_code == 2, f"unexpected exit; output: {result.output}" - mock_service_manager.cleanup_image_validation_service.assert_called_once_with( - service_name="IMG456" - ) - assert "Image validation FAILED" in result.output - - -@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") -def test_build_image_cli_skip_validation( - mock_stage_execute_query, - mock_stage_put, - mock_service_manager_class, - mock_object_manager_class, - mock_sleep, - runner, - temporary_directory, -): - """--skip-validation propagates SKIP_VALIDATION=true AND skips the post-push smoke test.""" - build_context = Path(temporary_directory) / "build_context" - build_context.mkdir() - (build_context / "Dockerfile").write_text("FROM python:3.10-alpine") - mock_stage_put.return_value = Mock(fetchall=lambda: []) - - mock_service_manager = Mock() - 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_service_manager_class.return_value = mock_service_manager - - mock_object_manager = Mock() - mock_object_manager_class.return_value = mock_object_manager - describe_cursor = Mock() - describe_cursor.fetchone.return_value = {"status": "RUNNING"} - mock_object_manager.describe.return_value = 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", - "--workload-type", - "ml_job", - "--skip-validation", - ], - catch_exceptions=False, - ) - - assert result.exit_code == 0, f"unexpected exit; output: {result.output}" - # The runner-side flag was passed to build_image: - build_image_kwargs = mock_service_manager.build_image.call_args.kwargs - assert build_image_kwargs["skip_validation"] is True - assert build_image_kwargs["workload_type"] == "ml_job" - # And the post-push smoke test was NOT invoked: - mock_service_manager.start_image_validation_service.assert_not_called() - mock_service_manager.cleanup_image_validation_service.assert_not_called() - - -@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") -def test_build_image_cli_validation_keyboard_interrupt( - mock_stage_execute_query, - mock_stage_put, - mock_service_manager_class, - mock_object_manager_class, - mock_sleep, - runner, - temporary_directory, -): - """Ctrl-C during log streaming still triggers cleanup and reports transient.""" - build_context = Path(temporary_directory) / "build_context" - build_context.mkdir() - (build_context / "Dockerfile").write_text("FROM python:3.10-alpine") - mock_stage_put.return_value = Mock(fetchall=lambda: []) - - mock_service_manager = Mock() - 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 - - def _stream_then_interrupt(): - # Build-image stream returns DONE quickly, then validation stream raises. - if not getattr(_stream_then_interrupt, "called_once", False): - _stream_then_interrupt.called_once = True - return iter([("__TERMINAL_STATUS__", "DONE")]) - - def gen(): - yield "validation log line" - raise KeyboardInterrupt() - - return gen() - - mock_service_manager.stream_logs.side_effect = lambda *a, **kw: _stream_then_interrupt() - mock_service_manager.build_image_registry_url.return_value = ( - "org-acct.registry-local.snowflakecomputing.com/db/sch/repo" - ) - mock_service_manager.start_image_validation_service.return_value = { - "status": "started", - "service_name": "IMG_INTR", - "schema_fqn": "DB.SCH", - "stage_fqn": "@DB.SCH._SNOWFLAKE_IMAGE_VALIDATION_PAYLOAD_IMG_INTR", - "workload_type": "ML_JOB", - } - mock_service_manager.cleanup_image_validation_service.return_value = { - "dropped_service": True, - "dropped_stage": True, - } - mock_service_manager_class.return_value = mock_service_manager - - mock_object_manager = Mock() - mock_object_manager_class.return_value = mock_object_manager - describe_cursor = Mock() - describe_cursor.fetchone.return_value = {"status": "RUNNING"} - mock_object_manager.describe.return_value = 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", - "--workload-type", - "ml_job", - ], - catch_exceptions=False, - ) - - # Ctrl-C => transient (exit 3); cleanup must still run exactly once. - assert result.exit_code == 3, f"unexpected exit; output: {result.output}" - mock_service_manager.cleanup_image_validation_service.assert_called_once_with( - service_name="IMG_INTR" - ) - assert "could not run to completion" in result.output