From 149e05e6b59f37aae9b1d613210d47a0dff57fe3 Mon Sep 17 00:00:00 2001 From: nayan3107 Date: Tue, 20 Oct 2026 15:07:56 -0700 Subject: [PATCH 1/4] feat(mlops): Add AthenaStep pipeline step Add AthenaStep, which takes the Athena StartQueryExecution request as a dict and emits it unchanged as the step's Arguments. The service validates the request at CreatePipeline, so the SDK adds no field vocabulary of its own. Properties follow the GetQueryExecutionOutput shape, so downstream references use the QueryExecution prefix. Integ tests passed against a preprod stage. --- X-AI-Prompt: Port the AthenaStep to the latest master of the PySDK and commit it on feat/athena-step X-AI-Tool: kiro-cli --- .../src/sagemaker/core/workflow/properties.py | 4 +- .../src/sagemaker/mlops/workflow/__init__.py | 2 + .../sagemaker/mlops/workflow/athena_step.py | 147 +++++++++++ .../src/sagemaker/mlops/workflow/steps.py | 1 + .../tests/integ/workflow/test_athena_step.py | 218 +++++++++++++++ .../tests/unit/workflow/test_athena_step.py | 248 ++++++++++++++++++ .../tests/unit/workflow/test_job_step.py | 1 + 7 files changed, 620 insertions(+), 1 deletion(-) create mode 100644 sagemaker-mlops/src/sagemaker/mlops/workflow/athena_step.py create mode 100644 sagemaker-mlops/tests/integ/workflow/test_athena_step.py create mode 100644 sagemaker-mlops/tests/unit/workflow/test_athena_step.py diff --git a/sagemaker-core/src/sagemaker/core/workflow/properties.py b/sagemaker-core/src/sagemaker/core/workflow/properties.py index 6c51cb6ced..d4c3551d02 100644 --- a/sagemaker-core/src/sagemaker/core/workflow/properties.py +++ b/sagemaker-core/src/sagemaker/core/workflow/properties.py @@ -31,7 +31,7 @@ class PropertiesMeta(ABCMeta): """Load an internal shapes attribute from the botocore service model - for sagemaker and emr service. + for sagemaker, emr, and athena services. """ _shapes_map = dict() @@ -44,8 +44,10 @@ def __new__(mcs, *args, **kwargs): sagemaker_model = loader.load_service_model("sagemaker", "service-2") emr_model = loader.load_service_model("emr", "service-2") + athena_model = loader.load_service_model("athena", "service-2") mcs._shapes_map["sagemaker"] = sagemaker_model["shapes"] mcs._shapes_map["emr"] = emr_model["shapes"] + mcs._shapes_map["athena"] = athena_model["shapes"] return super().__new__(mcs, *args, **kwargs) diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/__init__.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/__init__.py index dfd1f4a8bb..f5cf65493f 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/__init__.py +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/__init__.py @@ -44,6 +44,7 @@ # Step implementations from sagemaker.mlops.workflow.automl_step import AutoMLStep +from sagemaker.mlops.workflow.athena_step import AthenaStep from sagemaker.mlops.workflow.callback_step import CallbackStep, CallbackOutput from sagemaker.mlops.workflow.clarify_check_step import ClarifyCheckStep from sagemaker.mlops.workflow.cleanup_step import CleanupStep, S3CleanupSource @@ -113,6 +114,7 @@ "TuningStep", # Step implementations "AutoMLStep", + "AthenaStep", "CallbackStep", "CallbackOutput", "ClarifyCheckStep", diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/athena_step.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/athena_step.py new file mode 100644 index 0000000000..d1ecd9aeda --- /dev/null +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/athena_step.py @@ -0,0 +1,147 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Step definition for the Athena step type in SageMaker Pipelines. + +``AthenaStep`` runs an Amazon Athena query as part of a SageMaker Pipelines +execution. The step is a pure passthrough: ``arguments`` is the Athena +``StartQueryExecution`` request expressed as a dict, and it is emitted +verbatim as the step's ``Arguments`` in the pipeline definition. The +SageMaker Pipelines service validates it against the Athena API shape at +``CreatePipeline`` time, so the SDK adds no vocabulary of its own -- use +the same PascalCase field names as the Athena API / boto3. + +Example:: + + from sagemaker.mlops.workflow import AthenaStep, ParameterString + + output_location = ParameterString(name="OutputLocation") + + step = AthenaStep( + name="RunQuery", + arguments={ + "QueryString": "SELECT ? AS region", + "QueryExecutionContext": {"Database": "default"}, + "ResultConfiguration": {"OutputLocation": output_location}, + "WorkGroup": "primary", + "ExecutionParameters": ["'us-west-2'"], + }, + ) + +Notes on service-side rules (enforced at ``CreatePipeline`` / execution, +not by this class): + +* ``QueryString`` is literal-only. Pipeline expressions such as + ``{"Get": "Parameters.X"}`` inside it are rejected to prevent SQL + injection. Bind runtime values with ``ExecutionParameters`` (``?`` + placeholders) instead; pipeline variables are allowed there and in every + other field. +* ``ClientRequestToken`` is generated by the service at execution time and + should not be supplied. + +Downstream steps reference this step's output through ``properties``, which +follow the Athena ``GetQueryExecutionOutput`` shape:: + + step.properties.QueryExecution.QueryExecutionId + step.properties.QueryExecution.ResultConfiguration.OutputLocation + step.properties.QueryExecution.Status.State +""" + +from __future__ import absolute_import + +import copy +from typing import Any, Dict, List, Optional, Union + +from sagemaker.core.helper.pipeline_variable import RequestType +from sagemaker.core.workflow.properties import Properties + +from sagemaker.mlops.workflow.retry import RetryPolicy +from sagemaker.mlops.workflow.step_collections import StepCollection +from sagemaker.mlops.workflow.steps import ( + CacheConfig, + ConfigurableRetryStep, + Step, + StepTypeEnum, +) + + +class AthenaStep(ConfigurableRetryStep): + """Runs an Amazon Athena query as part of a SageMaker Pipeline. + + ``arguments`` is the Athena ``StartQueryExecution`` request as a dict and + is passed through unchanged as the step's ``Arguments``. Field names, + nesting, and which fields are required are all defined by the Athena API + and validated by the SageMaker Pipelines service, not by this class. + """ + + def __init__( + self, + name: str, + arguments: Dict[str, Any], + display_name: Optional[str] = None, + description: Optional[str] = None, + depends_on: Optional[List[Union[str, Step, StepCollection]]] = None, + cache_config: Optional[CacheConfig] = None, + retry_policies: Optional[List[RetryPolicy]] = None, + ): + """Construct an ``AthenaStep``. + + Args: + name: The name of the step. + arguments: The Athena ``StartQueryExecution`` request as a dict + (PascalCase keys, e.g. ``QueryString``, ``WorkGroup``, + ``QueryExecutionContext``, ``ResultConfiguration``, + ``ExecutionParameters``). Values may be literals or pipeline + variables, except ``QueryString`` which must be a literal. + display_name: Optional display name. + description: Optional description. + depends_on: Optional explicit step dependencies. Wrap an entry in + ``StepDependency`` to run this step when the upstream step fails, + stops, or reaches any terminal state. + cache_config: Optional cache configuration. + retry_policies: Optional list of ``StepRetryPolicy`` instances. + """ + super(AthenaStep, self).__init__( + name=name, + step_type=StepTypeEnum.ATHENA, + display_name=display_name, + description=description, + depends_on=depends_on, + retry_policies=retry_policies, + ) + # Defensive deep-copy so caller-side mutations after construction do + # not leak into the compiled pipeline definition. + self.args: Dict[str, Any] = copy.deepcopy(arguments) + self.cache_config = cache_config + self._properties = Properties( + step_name=name, + step=self, + shape_name="GetQueryExecutionOutput", + service_name="athena", + ) + + @property + def arguments(self) -> RequestType: + """The Athena ``StartQueryExecution`` request dict, emitted as ``Arguments``.""" + return self.args + + @property + def properties(self): + """A ``Properties`` object shaped like Athena's ``GetQueryExecutionOutput``.""" + return self._properties + + def to_request(self) -> RequestType: + """Get the request structure for workflow service calls.""" + request_dict = super().to_request() + if self.cache_config: + request_dict.update(self.cache_config.config) + return request_dict diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/steps.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/steps.py index 6073cf4715..8005363301 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/steps.py +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/steps.py @@ -71,6 +71,7 @@ class StepTypeEnum(Enum): JOB = "Job" CLEANUP = "Cleanup" FOR_EACH = "ForEach" + ATHENA = "Athena" class Step(Entity): diff --git a/sagemaker-mlops/tests/integ/workflow/test_athena_step.py b/sagemaker-mlops/tests/integ/workflow/test_athena_step.py new file mode 100644 index 0000000000..e44453fe2d --- /dev/null +++ b/sagemaker-mlops/tests/integ/workflow/test_athena_step.py @@ -0,0 +1,218 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Integration tests for AthenaStep. + +Runs Athena queries end-to-end through SageMaker Pipelines and checks that: + +* a valid query drives the execution to ``Succeeded`` and the result CSV is + written to the configured S3 output location, +* an invalid query drives the execution to ``Failed`` with a ``FailureReason``, +* a full-shape ``arguments`` dict, together with ``CacheConfig`` and + ``RetryPolicies``, is accepted by ``CreatePipeline`` and round-trips + through ``DescribePipeline`` unchanged. + +Requirements on top of the shared workflow integ fixtures: + +* The pipeline execution role (``role`` fixture) needs + ``athena:StartQueryExecution``, ``athena:GetQueryExecution`` and + ``athena:StopQueryExecution``, plus ``s3:GetBucketLocation`` and + ``s3:PutObject`` on the session default bucket, where Athena writes the + query results (under ``//``). +* No Athena workgroup, database or table setup: the queries run in the + account's default ``primary`` workgroup and do not read any table. +* The Athena step type must be available to the account and Region; + otherwise ``CreatePipeline`` rejects the pipeline definition. +""" + +from __future__ import absolute_import + +import json + +from sagemaker.core.common_utils import unique_name_from_base +from sagemaker.core.workflow.parameters import ParameterString +from sagemaker.mlops.workflow.athena_step import AthenaStep +from sagemaker.mlops.workflow.pipeline import Pipeline +from sagemaker.mlops.workflow.retry import StepExceptionTypeEnum, StepRetryPolicy +from sagemaker.mlops.workflow.steps import CacheConfig + +from tests.integ.workflow.helpers import wait_pipeline_execution + +_OUTPUT_LOCATION_PARAMETER = "OutputLocation" + + +def _result_location(sagemaker_session, pipeline_name): + """Return ``(bucket, key_prefix, s3_uri)`` for this pipeline's Athena results.""" + bucket = sagemaker_session.default_bucket() + key_prefix = "/".join( + part for part in (sagemaker_session.default_bucket_prefix, pipeline_name) if part + ) + return bucket, key_prefix, "s3://{}/{}/".format(bucket, key_prefix) + + +def _create_and_execute(pipeline, role, output_location): + """Create ``pipeline``, start one execution and wait for it to finish. + + Returns ``(execution, execution_status, step)`` for the pipeline's single step. + """ + pipeline.create(role_arn=role) + execution = pipeline.start(parameters={_OUTPUT_LOCATION_PARAMETER: output_location}) + wait_pipeline_execution(execution=execution, delay=20, max_attempts=30) + + status = execution.describe()["PipelineExecutionStatus"] + steps = execution.list_steps() + assert len(steps) == 1, "Expected exactly one step, got {}".format(len(steps)) + return execution, status, steps[0] + + +def _cleanup(pipeline, execution, sagemaker_session, bucket, key_prefix): + """Best-effort teardown: stop a still-running execution, delete the pipeline + and delete the Athena result files.""" + if execution is not None: + try: + execution.stop() + except Exception: + pass + try: + pipeline.delete() + except Exception: + pass + if bucket and key_prefix: + try: + s3_resource = sagemaker_session.boto_session.resource( + "s3", region_name=sagemaker_session.boto_region_name + ) + s3_resource.Bucket(bucket).objects.filter(Prefix=key_prefix).delete() + except Exception: + pass + + +def test_athena_step_query_succeeds_and_writes_result(sagemaker_session_for_pipeline, role): + """A well-formed ``SELECT 1`` succeeds and its result CSV lands in S3.""" + pipeline_name = unique_name_from_base("athena-step-integ") + bucket, key_prefix, output_location = _result_location( + sagemaker_session_for_pipeline, pipeline_name + ) + output_location_parameter = ParameterString(name=_OUTPUT_LOCATION_PARAMETER) + + step = AthenaStep( + name="AthenaQuery", + arguments={ + "QueryString": "SELECT 1 AS test_result", + "ResultConfiguration": {"OutputLocation": output_location_parameter}, + }, + ) + pipeline = Pipeline( + name=pipeline_name, + parameters=[output_location_parameter], + steps=[step], + sagemaker_session=sagemaker_session_for_pipeline, + ) + + execution = None + try: + execution, status, athena_step = _create_and_execute(pipeline, role, output_location) + assert status == "Succeeded", "Execution status was {}".format(status) + assert athena_step["StepName"] == "AthenaQuery" + assert athena_step["StepStatus"] == "Succeeded", athena_step.get("FailureReason") + + query_execution_id = athena_step["Metadata"]["Athena"]["QueryExecutionId"] + result_key = "{}/{}.csv".format(key_prefix, query_execution_id) + result_files = sagemaker_session_for_pipeline.list_s3_files(bucket, key_prefix) + assert result_key in result_files, "Athena result file not found: {}".format(result_files) + finally: + _cleanup(pipeline, execution, sagemaker_session_for_pipeline, bucket, key_prefix) + + +def test_athena_step_invalid_query_fails(sagemaker_session_for_pipeline, role): + """An invalid SQL statement drives the step, and the execution, to ``Failed``.""" + pipeline_name = unique_name_from_base("athena-step-integ") + bucket, key_prefix, output_location = _result_location( + sagemaker_session_for_pipeline, pipeline_name + ) + output_location_parameter = ParameterString(name=_OUTPUT_LOCATION_PARAMETER) + + step = AthenaStep( + name="AthenaInvalidQuery", + arguments={ + "QueryString": "SELCT INVALID SYNTAX FROM nonexistent_table_xyz", + "ResultConfiguration": {"OutputLocation": output_location_parameter}, + }, + ) + pipeline = Pipeline( + name=pipeline_name, + parameters=[output_location_parameter], + steps=[step], + sagemaker_session=sagemaker_session_for_pipeline, + ) + + execution = None + try: + execution, status, athena_step = _create_and_execute(pipeline, role, output_location) + assert status == "Failed", "Execution status was {}".format(status) + assert athena_step["StepStatus"] == "Failed" + assert athena_step.get("FailureReason"), "Step is Failed but no FailureReason was set" + finally: + _cleanup(pipeline, execution, sagemaker_session_for_pipeline, bucket, key_prefix) + + +def test_athena_step_full_arguments_roundtrip(sagemaker_session_for_pipeline, role): + """Full-shape ``arguments`` plus ``CacheConfig``/``RetryPolicies`` are accepted at + ``CreatePipeline`` and come back unchanged from ``DescribePipeline``.""" + pipeline_name = unique_name_from_base("athena-step-integ") + output_location_parameter = ParameterString(name=_OUTPUT_LOCATION_PARAMETER) + + step = AthenaStep( + name="AthenaFullQuery", + arguments={ + "QueryString": "SELECT 1", + "QueryExecutionContext": {"Database": "default", "Catalog": "AwsDataCatalog"}, + "ResultConfiguration": { + "OutputLocation": output_location_parameter, + "EncryptionConfiguration": {"EncryptionOption": "SSE_S3"}, + }, + "WorkGroup": "primary", + }, + cache_config=CacheConfig(enable_caching=True, expire_after="PT1H"), + retry_policies=[ + StepRetryPolicy( + exception_types=[StepExceptionTypeEnum.SERVICE_FAULT], + max_attempts=3, + interval_seconds=30, + backoff_rate=2.0, + ) + ], + ) + pipeline = Pipeline( + name=pipeline_name, + parameters=[output_location_parameter], + steps=[step], + sagemaker_session=sagemaker_session_for_pipeline, + ) + + try: + pipeline.create(role_arn=role) + + described_definition = json.loads(pipeline.describe()["PipelineDefinition"]) + sent_definition = json.loads(pipeline.definition()) + assert described_definition["Steps"] == sent_definition["Steps"] + + (athena_step,) = described_definition["Steps"] + assert athena_step["Type"] == "Athena" + assert athena_step["Arguments"]["WorkGroup"] == "primary" + assert athena_step["Arguments"]["ResultConfiguration"]["OutputLocation"] == { + "Get": "Parameters.{}".format(_OUTPUT_LOCATION_PARAMETER) + } + assert athena_step["CacheConfig"] == {"Enabled": True, "ExpireAfter": "PT1H"} + assert athena_step["RetryPolicies"][0]["ExceptionType"] == ["Step.SERVICE_FAULT"] + finally: + _cleanup(pipeline, None, sagemaker_session_for_pipeline, None, None) diff --git a/sagemaker-mlops/tests/unit/workflow/test_athena_step.py b/sagemaker-mlops/tests/unit/workflow/test_athena_step.py new file mode 100644 index 0000000000..3dfa5b23a3 --- /dev/null +++ b/sagemaker-mlops/tests/unit/workflow/test_athena_step.py @@ -0,0 +1,248 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Unit tests for AthenaStep -- pure passthrough of the StartQueryExecution request.""" + +from __future__ import absolute_import + +import json +from unittest.mock import MagicMock + +import boto3 +import pytest + +from sagemaker.core.workflow.parameters import ParameterString +from sagemaker.core.workflow.pipeline_context import PipelineSession +from sagemaker.core.workflow.properties import Properties +from sagemaker.mlops.workflow.athena_step import AthenaStep +from sagemaker.mlops.workflow.pipeline import Pipeline +from sagemaker.mlops.workflow.retry import StepExceptionTypeEnum, StepRetryPolicy +from sagemaker.mlops.workflow.step_dependency import StepDependency, TriggerCondition +from sagemaker.mlops.workflow.steps import CacheConfig, StepTypeEnum + +# A full-shape Athena StartQueryExecution request, as a customer would write it. +FULL_ARGUMENTS = { + "QueryString": "SELECT * FROM daily_txn WHERE txn_date >= '2026-01-01'", + "QueryExecutionContext": {"Catalog": "unity-catalog", "Database": "prod_catalog$transactions"}, + "ResultConfiguration": { + "OutputLocation": "s3://bucket/sagemaker-athena-staging/", + "EncryptionConfiguration": {"EncryptionOption": "SSE_S3"}, + }, + "WorkGroup": "primary", + "ResultReuseConfiguration": { + "ResultReuseByAgeConfiguration": {"Enabled": True, "MaxAgeInMinutes": 60} + }, + "ExecutionParameters": ["us-east-1", "2026-01-01"], +} + + +# --------------------------------------------------------------------------- +# Passthrough contract +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "arguments", + [ + FULL_ARGUMENTS, + # The service accepts QueryString alone; everything else is optional. + {"QueryString": "SELECT 1"}, + # The SDK holds no opinion on required fields -- the service decides. + {}, + ], + ids=["full", "minimal", "empty"], +) +def test_arguments_passed_through_unchanged(arguments): + step = AthenaStep(name="RunQuery", arguments=arguments) + assert step.arguments == arguments + assert step.to_request()["Arguments"] == arguments + + +def test_arguments_is_required(): + with pytest.raises(TypeError): + AthenaStep(name="RunQuery") # pylint: disable=no-value-for-parameter + + +def test_caller_mutations_do_not_leak_into_step_state(): + arguments = { + "QueryString": "SELECT 1", + "ResultConfiguration": {"EncryptionConfiguration": {"EncryptionOption": "SSE_S3"}}, + "ExecutionParameters": ["1"], + } + step = AthenaStep(name="RunQuery", arguments=arguments) + + arguments["QueryString"] = "MUTATED" + arguments["ResultConfiguration"]["EncryptionConfiguration"]["EncryptionOption"] = "MUTATED" + arguments["ExecutionParameters"].append("2") + + assert step.arguments == { + "QueryString": "SELECT 1", + "ResultConfiguration": {"EncryptionConfiguration": {"EncryptionOption": "SSE_S3"}}, + "ExecutionParameters": ["1"], + } + + +def test_pipeline_variable_inside_arguments_compiles_to_get_expression(): + """A pipeline variable nested anywhere in ``arguments`` serializes to a + ``{"Get": ...}`` expression in the compiled definition, and the step + compiles unchanged under a ``PipelineSession``.""" + pipeline_session = PipelineSession( + boto_session=boto3.Session(region_name="us-west-2"), + sagemaker_client=MagicMock(name="sagemaker_client"), + sagemaker_config={}, + ) + output_location = ParameterString(name="OutputLocation") + step = AthenaStep( + name="RunQuery", + arguments={ + "QueryString": "SELECT 1", + "ResultConfiguration": {"OutputLocation": output_location}, + }, + ) + pipeline = Pipeline( + name="athena-unit", + parameters=[output_location], + steps=[step], + sagemaker_session=pipeline_session, + ) + + definition = json.loads(pipeline.definition()) + + compiled_step = definition["Steps"][0] + assert compiled_step["Type"] == "Athena" + assert compiled_step["Arguments"]["ResultConfiguration"]["OutputLocation"] == { + "Get": "Parameters.OutputLocation" + } + # Nothing was intercepted: AthenaStep never calls a Session helper. + assert pipeline_session.context is None + + +def test_run_conditions_compile_to_depends_on_triggers(): + query = AthenaStep(name="Query", arguments={"QueryString": "SELECT 1"}) + on_failure = AthenaStep( + name="OnFailure", + arguments={"QueryString": "SELECT 1"}, + depends_on=[StepDependency(query, TriggerCondition.FAILED)], + ) + on_any = AthenaStep( + name="OnAny", + arguments={"QueryString": "SELECT 1"}, + depends_on=[StepDependency("Query", TriggerCondition.ANY)], + ) + on_success = AthenaStep( + name="OnSuccess", arguments={"QueryString": "SELECT 1"}, depends_on=[query] + ) + pipeline = Pipeline( + name="athena-run-conditions", + steps=[query, on_failure, on_any, on_success], + sagemaker_session=MagicMock(), + ) + + depends_on = { + step["Name"]: step.get("DependsOn") for step in json.loads(pipeline.definition())["Steps"] + } + + assert depends_on == { + "Query": None, + "OnFailure": [{"Step": "Query", "Trigger": "Failed"}], + "OnAny": [{"Step": "Query", "Trigger": "Any"}], + "OnSuccess": ["Query"], + } + + +# --------------------------------------------------------------------------- +# Properties shape +# --------------------------------------------------------------------------- + + +def test_properties_shape_and_reference_paths(): + step = AthenaStep(name="RunQuery", arguments={"QueryString": "SELECT 1"}) + + assert isinstance(step.properties, Properties) + # Loaded from botocore's GetQueryExecutionOutput shape (service_name="athena"); + # the QueryExecution envelope matches the service-side property paths. + assert ( + step.properties.QueryExecution.QueryExecutionId.expr["Get"] + == "Steps.RunQuery.QueryExecution.QueryExecutionId" + ) + assert ( + step.properties.QueryExecution.Status.State.expr["Get"] + == "Steps.RunQuery.QueryExecution.Status.State" + ) + assert ( + step.properties.QueryExecution.ResultConfiguration.OutputLocation.expr["Get"] + == "Steps.RunQuery.QueryExecution.ResultConfiguration.OutputLocation" + ) + + +# --------------------------------------------------------------------------- +# to_request() +# --------------------------------------------------------------------------- + + +def test_to_request_includes_step_fields_cache_config_and_retry_policies(): + step = AthenaStep( + name="QueryUCData", + arguments=FULL_ARGUMENTS, + display_name="Query UC Transactions", + description="Runs Athena query against Unity Catalog", + depends_on=["Upstream"], + cache_config=CacheConfig(enable_caching=True, expire_after="PT1H"), + retry_policies=[ + StepRetryPolicy( + exception_types=[StepExceptionTypeEnum.SERVICE_FAULT], + max_attempts=3, + interval_seconds=30, + backoff_rate=2.0, + ) + ], + ) + + assert step.to_request() == { + "Name": "QueryUCData", + "Type": "Athena", + "DisplayName": "Query UC Transactions", + "Description": "Runs Athena query against Unity Catalog", + "DependsOn": ["Upstream"], + "Arguments": FULL_ARGUMENTS, + "CacheConfig": {"Enabled": True, "ExpireAfter": "PT1H"}, + "RetryPolicies": [ + { + "ExceptionType": ["Step.SERVICE_FAULT"], + "MaxAttempts": 3, + "IntervalSeconds": 30, + "BackoffRate": 2.0, + } + ], + } + + +def test_to_request_omits_optional_sections_when_not_set(): + request_dict = AthenaStep(name="RunQuery", arguments={"QueryString": "SELECT 1"}).to_request() + assert "CacheConfig" not in request_dict + assert "RetryPolicies" not in request_dict + assert "DependsOn" not in request_dict + + +# --------------------------------------------------------------------------- +# Package wiring +# --------------------------------------------------------------------------- + + +def test_athena_step_exported_from_workflow_package(): + from sagemaker.mlops.workflow import AthenaStep as ExportedAthenaStep + + assert ExportedAthenaStep is AthenaStep + + +def test_step_type_enum_has_athena_value(): + assert StepTypeEnum.ATHENA.value == "Athena" diff --git a/sagemaker-mlops/tests/unit/workflow/test_job_step.py b/sagemaker-mlops/tests/unit/workflow/test_job_step.py index b49e752dfe..c208a79cd8 100644 --- a/sagemaker-mlops/tests/unit/workflow/test_job_step.py +++ b/sagemaker-mlops/tests/unit/workflow/test_job_step.py @@ -52,6 +52,7 @@ "JOB": "Job", "CLEANUP": "Cleanup", "FOR_EACH": "ForEach", + "ATHENA": "Athena", } From c0c42b56fec14332c53fd69bbca519f7af9ee10c Mon Sep 17 00:00:00 2001 From: nayan3107 Date: Wed, 21 Oct 2026 08:51:42 -0700 Subject: [PATCH 2/4] feat(mlops): Add AthenaQueryConfig for AthenaStep Customers should not hand-build the nested StartQueryExecution dict. AthenaQueryConfig takes named params and assembles the request, like EMRServerlessJobConfig; the service still validates it at CreatePipeline. A raw arguments dict remains as an escape hatch for fields the config does not cover. --- X-AI-Prompt: Address review asking for named params instead of a raw dict on AthenaStep X-AI-Tool: kiro-cli --- .../src/sagemaker/mlops/workflow/__init__.py | 3 +- .../sagemaker/mlops/workflow/athena_step.py | 188 ++++++++++++++---- .../tests/integ/workflow/test_athena_step.py | 42 ++-- .../tests/unit/workflow/test_athena_step.py | 151 ++++++++++---- 4 files changed, 285 insertions(+), 99 deletions(-) diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/__init__.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/__init__.py index f5cf65493f..b22f6137ce 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/__init__.py +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/__init__.py @@ -44,7 +44,7 @@ # Step implementations from sagemaker.mlops.workflow.automl_step import AutoMLStep -from sagemaker.mlops.workflow.athena_step import AthenaStep +from sagemaker.mlops.workflow.athena_step import AthenaQueryConfig, AthenaStep from sagemaker.mlops.workflow.callback_step import CallbackStep, CallbackOutput from sagemaker.mlops.workflow.clarify_check_step import ClarifyCheckStep from sagemaker.mlops.workflow.cleanup_step import CleanupStep, S3CleanupSource @@ -114,6 +114,7 @@ "TuningStep", # Step implementations "AutoMLStep", + "AthenaQueryConfig", "AthenaStep", "CallbackStep", "CallbackOutput", diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/athena_step.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/athena_step.py index d1ecd9aeda..585460275f 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/athena_step.py +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/athena_step.py @@ -13,36 +13,35 @@ """Step definition for the Athena step type in SageMaker Pipelines. ``AthenaStep`` runs an Amazon Athena query as part of a SageMaker Pipelines -execution. The step is a pure passthrough: ``arguments`` is the Athena -``StartQueryExecution`` request expressed as a dict, and it is emitted -verbatim as the step's ``Arguments`` in the pipeline definition. The -SageMaker Pipelines service validates it against the Athena API shape at -``CreatePipeline`` time, so the SDK adds no vocabulary of its own -- use -the same PascalCase field names as the Athena API / boto3. +execution. Describe the query with an ``AthenaQueryConfig``; it assembles the +Athena ``StartQueryExecution`` request that becomes the step's ``Arguments``:: -Example:: - - from sagemaker.mlops.workflow import AthenaStep, ParameterString + from sagemaker.mlops.workflow import AthenaQueryConfig, AthenaStep, ParameterString output_location = ParameterString(name="OutputLocation") step = AthenaStep( name="RunQuery", - arguments={ - "QueryString": "SELECT ? AS region", - "QueryExecutionContext": {"Database": "default"}, - "ResultConfiguration": {"OutputLocation": output_location}, - "WorkGroup": "primary", - "ExecutionParameters": ["'us-west-2'"], - }, + query_config=AthenaQueryConfig( + query_string="SELECT ? AS region", + database="default", + work_group="primary", + output_location=output_location, + encryption_option="SSE_S3", + execution_parameters=["'us-west-2'"], + ), ) -Notes on service-side rules (enforced at ``CreatePipeline`` / execution, -not by this class): +For ``StartQueryExecution`` fields that ``AthenaQueryConfig`` does not cover, +pass the full request as a dict through ``arguments`` instead; it is emitted +unchanged. + +The SageMaker Pipelines service validates the request at ``CreatePipeline`` +time, not this module: * ``QueryString`` is literal-only. Pipeline expressions such as ``{"Get": "Parameters.X"}`` inside it are rejected to prevent SQL - injection. Bind runtime values with ``ExecutionParameters`` (``?`` + injection. Bind runtime values with ``execution_parameters`` (``?`` placeholders) instead; pipeline variables are allowed there and in every other field. * ``ClientRequestToken`` is generated by the service at execution time and @@ -61,7 +60,7 @@ import copy from typing import Any, Dict, List, Optional, Union -from sagemaker.core.helper.pipeline_variable import RequestType +from sagemaker.core.helper.pipeline_variable import PipelineVariable, RequestType from sagemaker.core.workflow.properties import Properties from sagemaker.mlops.workflow.retry import RetryPolicy @@ -74,34 +73,144 @@ ) -class AthenaStep(ConfigurableRetryStep): - """Runs an Amazon Athena query as part of a SageMaker Pipeline. +class AthenaQueryConfig: + """Config for an Athena query run by an ``AthenaStep``. - ``arguments`` is the Athena ``StartQueryExecution`` request as a dict and - is passed through unchanged as the step's ``Arguments``. Field names, - nesting, and which fields are required are all defined by the Athena API - and validated by the SageMaker Pipelines service, not by this class. + See the Athena `StartQueryExecution + `_ + API for the meaning and allowed values of each field. """ + def __init__( + self, + query_string: str, + database: Optional[Union[str, PipelineVariable]] = None, + catalog: Optional[Union[str, PipelineVariable]] = None, + work_group: Optional[Union[str, PipelineVariable]] = None, + execution_parameters: Optional[List[Union[str, PipelineVariable]]] = None, + output_location: Optional[Union[str, PipelineVariable]] = None, + encryption_option: Optional[Union[str, PipelineVariable]] = None, + kms_key: Optional[Union[str, PipelineVariable]] = None, + expected_bucket_owner: Optional[Union[str, PipelineVariable]] = None, + s3_acl_option: Optional[Union[str, PipelineVariable]] = None, + result_reuse_max_age_in_minutes: Optional[Union[int, PipelineVariable]] = None, + ): # pylint: disable=too-many-positional-arguments + """Create a definition for an Athena query. + + Args: + query_string (str): The SQL query. Must be a literal string; bind runtime + values with ``execution_parameters``. + database (str, optional): The database the query runs in + (``QueryExecutionContext.Database``). + catalog (str, optional): The data catalog the query runs in + (``QueryExecutionContext.Catalog``). + work_group (str, optional): The Athena workgroup. + execution_parameters (List[str], optional): Values for the ``?`` + placeholders in ``query_string``, in order. + output_location (str, optional): S3 URI for the query results + (``ResultConfiguration.OutputLocation``). May be omitted when the + workgroup sets it. + encryption_option (str, optional): ``SSE_S3``, ``SSE_KMS`` or ``CSE_KMS`` + (``ResultConfiguration.EncryptionConfiguration.EncryptionOption``). + kms_key (str, optional): KMS key for ``SSE_KMS``/``CSE_KMS`` + (``ResultConfiguration.EncryptionConfiguration.KmsKey``). + expected_bucket_owner (str, optional): Account ID expected to own the + output bucket (``ResultConfiguration.ExpectedBucketOwner``). + s3_acl_option (str, optional): Canned ACL for the results, e.g. + ``BUCKET_OWNER_FULL_CONTROL`` + (``ResultConfiguration.AclConfiguration.S3AclOption``). + result_reuse_max_age_in_minutes (int, optional): When set, enables Athena + query result reuse with this maximum age + (``ResultReuseConfiguration.ResultReuseByAgeConfiguration``). + """ + self.query_string = query_string + self.database = database + self.catalog = catalog + self.work_group = work_group + self.execution_parameters = execution_parameters + self.output_location = output_location + self.encryption_option = encryption_option + self.kms_key = kms_key + self.expected_bucket_owner = expected_bucket_owner + self.s3_acl_option = s3_acl_option + self.result_reuse_max_age_in_minutes = result_reuse_max_age_in_minutes + + def to_request(self) -> RequestType: + """Convert this config to an Athena ``StartQueryExecution`` request dict.""" + request = {"QueryString": self.query_string} + + context = _drop_none({"Database": self.database, "Catalog": self.catalog}) + if context: + request["QueryExecutionContext"] = context + + if self.work_group is not None: + request["WorkGroup"] = self.work_group + if self.execution_parameters is not None: + request["ExecutionParameters"] = list(self.execution_parameters) + + encryption = _drop_none( + {"EncryptionOption": self.encryption_option, "KmsKey": self.kms_key} + ) + acl = _drop_none({"S3AclOption": self.s3_acl_option}) + result_configuration = _drop_none( + { + "OutputLocation": self.output_location, + "EncryptionConfiguration": encryption or None, + "ExpectedBucketOwner": self.expected_bucket_owner, + "AclConfiguration": acl or None, + } + ) + if result_configuration: + request["ResultConfiguration"] = result_configuration + + if self.result_reuse_max_age_in_minutes is not None: + request["ResultReuseConfiguration"] = { + "ResultReuseByAgeConfiguration": { + "Enabled": True, + "MaxAgeInMinutes": self.result_reuse_max_age_in_minutes, + } + } + + return request + + +def _drop_none(values: Dict[str, Any]) -> Dict[str, Any]: + return {key: value for key, value in values.items() if value is not None} + + +ERR_STR_WITH_BOTH_QUERY_CONFIG_AND_ARGUMENTS = ( + "AthenaStep {step_name} cannot have both query_config and arguments." +) + +ERR_STR_WITHOUT_QUERY_CONFIG_AND_ARGUMENTS = ( + "AthenaStep {step_name} must have either query_config or arguments." +) + + +class AthenaStep(ConfigurableRetryStep): + """Runs an Amazon Athena query as part of a SageMaker Pipeline.""" + def __init__( self, name: str, - arguments: Dict[str, Any], + query_config: Optional[AthenaQueryConfig] = None, + arguments: Optional[Dict[str, Any]] = None, display_name: Optional[str] = None, description: Optional[str] = None, depends_on: Optional[List[Union[str, Step, StepCollection]]] = None, cache_config: Optional[CacheConfig] = None, retry_policies: Optional[List[RetryPolicy]] = None, - ): + ): # pylint: disable=too-many-positional-arguments """Construct an ``AthenaStep``. + Exactly one of ``query_config`` and ``arguments`` must be given. + Args: name: The name of the step. - arguments: The Athena ``StartQueryExecution`` request as a dict - (PascalCase keys, e.g. ``QueryString``, ``WorkGroup``, - ``QueryExecutionContext``, ``ResultConfiguration``, - ``ExecutionParameters``). Values may be literals or pipeline - variables, except ``QueryString`` which must be a literal. + query_config: The Athena query to run. + arguments: The full Athena ``StartQueryExecution`` request as a dict + (PascalCase keys), emitted unchanged. Use it for request fields + that ``AthenaQueryConfig`` does not cover. display_name: Optional display name. description: Optional description. depends_on: Optional explicit step dependencies. Wrap an entry in @@ -118,9 +227,16 @@ def __init__( depends_on=depends_on, retry_policies=retry_policies, ) - # Defensive deep-copy so caller-side mutations after construction do - # not leak into the compiled pipeline definition. - self.args: Dict[str, Any] = copy.deepcopy(arguments) + + if query_config is None and arguments is None: + raise ValueError(ERR_STR_WITHOUT_QUERY_CONFIG_AND_ARGUMENTS.format(step_name=name)) + if query_config is not None and arguments is not None: + raise ValueError(ERR_STR_WITH_BOTH_QUERY_CONFIG_AND_ARGUMENTS.format(step_name=name)) + + # Copy so caller-side mutations after construction do not leak into the + # compiled pipeline definition. + request = query_config.to_request() if query_config is not None else arguments + self.args: Dict[str, Any] = copy.deepcopy(request) self.cache_config = cache_config self._properties = Properties( step_name=name, diff --git a/sagemaker-mlops/tests/integ/workflow/test_athena_step.py b/sagemaker-mlops/tests/integ/workflow/test_athena_step.py index e44453fe2d..01939eb3bd 100644 --- a/sagemaker-mlops/tests/integ/workflow/test_athena_step.py +++ b/sagemaker-mlops/tests/integ/workflow/test_athena_step.py @@ -17,7 +17,7 @@ * a valid query drives the execution to ``Succeeded`` and the result CSV is written to the configured S3 output location, * an invalid query drives the execution to ``Failed`` with a ``FailureReason``, -* a full-shape ``arguments`` dict, together with ``CacheConfig`` and +* a full ``AthenaQueryConfig``, together with ``CacheConfig`` and ``RetryPolicies``, is accepted by ``CreatePipeline`` and round-trips through ``DescribePipeline`` unchanged. @@ -40,7 +40,7 @@ from sagemaker.core.common_utils import unique_name_from_base from sagemaker.core.workflow.parameters import ParameterString -from sagemaker.mlops.workflow.athena_step import AthenaStep +from sagemaker.mlops.workflow.athena_step import AthenaQueryConfig, AthenaStep from sagemaker.mlops.workflow.pipeline import Pipeline from sagemaker.mlops.workflow.retry import StepExceptionTypeEnum, StepRetryPolicy from sagemaker.mlops.workflow.steps import CacheConfig @@ -106,10 +106,10 @@ def test_athena_step_query_succeeds_and_writes_result(sagemaker_session_for_pipe step = AthenaStep( name="AthenaQuery", - arguments={ - "QueryString": "SELECT 1 AS test_result", - "ResultConfiguration": {"OutputLocation": output_location_parameter}, - }, + query_config=AthenaQueryConfig( + query_string="SELECT 1 AS test_result", + output_location=output_location_parameter, + ), ) pipeline = Pipeline( name=pipeline_name, @@ -143,10 +143,10 @@ def test_athena_step_invalid_query_fails(sagemaker_session_for_pipeline, role): step = AthenaStep( name="AthenaInvalidQuery", - arguments={ - "QueryString": "SELCT INVALID SYNTAX FROM nonexistent_table_xyz", - "ResultConfiguration": {"OutputLocation": output_location_parameter}, - }, + query_config=AthenaQueryConfig( + query_string="SELCT INVALID SYNTAX FROM nonexistent_table_xyz", + output_location=output_location_parameter, + ), ) pipeline = Pipeline( name=pipeline_name, @@ -165,23 +165,23 @@ def test_athena_step_invalid_query_fails(sagemaker_session_for_pipeline, role): _cleanup(pipeline, execution, sagemaker_session_for_pipeline, bucket, key_prefix) -def test_athena_step_full_arguments_roundtrip(sagemaker_session_for_pipeline, role): - """Full-shape ``arguments`` plus ``CacheConfig``/``RetryPolicies`` are accepted at +def test_athena_step_full_query_config_roundtrip(sagemaker_session_for_pipeline, role): + """A full ``AthenaQueryConfig`` plus ``CacheConfig``/``RetryPolicies`` are accepted at ``CreatePipeline`` and come back unchanged from ``DescribePipeline``.""" pipeline_name = unique_name_from_base("athena-step-integ") output_location_parameter = ParameterString(name=_OUTPUT_LOCATION_PARAMETER) step = AthenaStep( name="AthenaFullQuery", - arguments={ - "QueryString": "SELECT 1", - "QueryExecutionContext": {"Database": "default", "Catalog": "AwsDataCatalog"}, - "ResultConfiguration": { - "OutputLocation": output_location_parameter, - "EncryptionConfiguration": {"EncryptionOption": "SSE_S3"}, - }, - "WorkGroup": "primary", - }, + query_config=AthenaQueryConfig( + query_string="SELECT 1", + database="default", + catalog="AwsDataCatalog", + work_group="primary", + output_location=output_location_parameter, + encryption_option="SSE_S3", + result_reuse_max_age_in_minutes=60, + ), cache_config=CacheConfig(enable_caching=True, expire_after="PT1H"), retry_policies=[ StepRetryPolicy( diff --git a/sagemaker-mlops/tests/unit/workflow/test_athena_step.py b/sagemaker-mlops/tests/unit/workflow/test_athena_step.py index 3dfa5b23a3..972980bc87 100644 --- a/sagemaker-mlops/tests/unit/workflow/test_athena_step.py +++ b/sagemaker-mlops/tests/unit/workflow/test_athena_step.py @@ -10,7 +10,7 @@ # distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF # ANY KIND, either express or implied. See the License for the specific # language governing permissions and limitations under the License. -"""Unit tests for AthenaStep -- pure passthrough of the StartQueryExecution request.""" +"""Unit tests for AthenaStep and AthenaQueryConfig.""" from __future__ import absolute_import @@ -23,53 +23,119 @@ from sagemaker.core.workflow.parameters import ParameterString from sagemaker.core.workflow.pipeline_context import PipelineSession from sagemaker.core.workflow.properties import Properties -from sagemaker.mlops.workflow.athena_step import AthenaStep +from sagemaker.mlops.workflow.athena_step import AthenaQueryConfig, AthenaStep from sagemaker.mlops.workflow.pipeline import Pipeline from sagemaker.mlops.workflow.retry import StepExceptionTypeEnum, StepRetryPolicy from sagemaker.mlops.workflow.step_dependency import StepDependency, TriggerCondition from sagemaker.mlops.workflow.steps import CacheConfig, StepTypeEnum -# A full-shape Athena StartQueryExecution request, as a customer would write it. +FULL_QUERY_CONFIG = AthenaQueryConfig( + query_string="SELECT * FROM daily_txn WHERE region = ? AND txn_date >= ?", + database="prod_catalog$transactions", + catalog="unity-catalog", + work_group="primary", + execution_parameters=["'us-east-1'", "'2026-01-01'"], + output_location="s3://bucket/sagemaker-athena-staging/", + encryption_option="SSE_KMS", + kms_key="arn:aws:kms:us-west-2:111122223333:key/example", + expected_bucket_owner="111122223333", + s3_acl_option="BUCKET_OWNER_FULL_CONTROL", + result_reuse_max_age_in_minutes=60, +) + FULL_ARGUMENTS = { - "QueryString": "SELECT * FROM daily_txn WHERE txn_date >= '2026-01-01'", - "QueryExecutionContext": {"Catalog": "unity-catalog", "Database": "prod_catalog$transactions"}, + "QueryString": "SELECT * FROM daily_txn WHERE region = ? AND txn_date >= ?", + "QueryExecutionContext": {"Database": "prod_catalog$transactions", "Catalog": "unity-catalog"}, + "WorkGroup": "primary", + "ExecutionParameters": ["'us-east-1'", "'2026-01-01'"], "ResultConfiguration": { "OutputLocation": "s3://bucket/sagemaker-athena-staging/", - "EncryptionConfiguration": {"EncryptionOption": "SSE_S3"}, + "EncryptionConfiguration": { + "EncryptionOption": "SSE_KMS", + "KmsKey": "arn:aws:kms:us-west-2:111122223333:key/example", + }, + "ExpectedBucketOwner": "111122223333", + "AclConfiguration": {"S3AclOption": "BUCKET_OWNER_FULL_CONTROL"}, }, - "WorkGroup": "primary", "ResultReuseConfiguration": { "ResultReuseByAgeConfiguration": {"Enabled": True, "MaxAgeInMinutes": 60} }, - "ExecutionParameters": ["us-east-1", "2026-01-01"], } # --------------------------------------------------------------------------- -# Passthrough contract +# AthenaQueryConfig # --------------------------------------------------------------------------- +def test_query_config_builds_full_request(): + assert FULL_QUERY_CONFIG.to_request() == FULL_ARGUMENTS + + +def test_query_config_minimal_emits_only_query_string(): + assert AthenaQueryConfig(query_string="SELECT 1").to_request() == {"QueryString": "SELECT 1"} + + @pytest.mark.parametrize( - "arguments", + "kwargs, expected", [ - FULL_ARGUMENTS, - # The service accepts QueryString alone; everything else is optional. - {"QueryString": "SELECT 1"}, - # The SDK holds no opinion on required fields -- the service decides. - {}, + ({"database": "db"}, {"QueryExecutionContext": {"Database": "db"}}), + ( + {"output_location": "s3://b/p/"}, + {"ResultConfiguration": {"OutputLocation": "s3://b/p/"}}, + ), + ( + {"encryption_option": "SSE_S3"}, + {"ResultConfiguration": {"EncryptionConfiguration": {"EncryptionOption": "SSE_S3"}}}, + ), + ( + {"s3_acl_option": "BUCKET_OWNER_FULL_CONTROL"}, + { + "ResultConfiguration": { + "AclConfiguration": {"S3AclOption": "BUCKET_OWNER_FULL_CONTROL"} + } + }, + ), ], - ids=["full", "minimal", "empty"], + ids=["context", "output-location", "encryption", "acl"], ) -def test_arguments_passed_through_unchanged(arguments): +def test_query_config_nests_only_the_fields_that_are_set(kwargs, expected): + request = AthenaQueryConfig(query_string="SELECT 1", **kwargs).to_request() + assert request == {"QueryString": "SELECT 1", **expected} + + +def test_step_with_query_config_emits_built_request(): + step = AthenaStep(name="RunQuery", query_config=FULL_QUERY_CONFIG) + assert step.arguments == FULL_ARGUMENTS + assert step.to_request()["Arguments"] == FULL_ARGUMENTS + + +# --------------------------------------------------------------------------- +# Raw arguments escape hatch +# --------------------------------------------------------------------------- + + +def test_raw_arguments_passed_through_unchanged(): + arguments = { + "QueryString": "SELECT 1", + "EngineConfiguration": {"MaxConcurrentDpus": 4}, + } step = AthenaStep(name="RunQuery", arguments=arguments) assert step.arguments == arguments - assert step.to_request()["Arguments"] == arguments -def test_arguments_is_required(): - with pytest.raises(TypeError): - AthenaStep(name="RunQuery") # pylint: disable=no-value-for-parameter +def test_requires_query_config_or_arguments(): + with pytest.raises(ValueError, match="must have either query_config or arguments"): + AthenaStep(name="RunQuery") + + +def test_rejects_both_query_config_and_arguments(): + with pytest.raises(ValueError, match="cannot have both query_config and arguments"): + AthenaStep( + name="RunQuery", + query_config=AthenaQueryConfig(query_string="SELECT 1"), + arguments={"QueryString": "SELECT 1"}, + ) def test_caller_mutations_do_not_leak_into_step_state(): @@ -91,10 +157,14 @@ def test_caller_mutations_do_not_leak_into_step_state(): } -def test_pipeline_variable_inside_arguments_compiles_to_get_expression(): - """A pipeline variable nested anywhere in ``arguments`` serializes to a - ``{"Get": ...}`` expression in the compiled definition, and the step - compiles unchanged under a ``PipelineSession``.""" +# --------------------------------------------------------------------------- +# Compilation +# --------------------------------------------------------------------------- + + +def test_pipeline_variable_in_query_config_compiles_to_get_expression(): + """A pipeline variable in a config field serializes to a ``{"Get": ...}`` + expression, and the step compiles unchanged under a ``PipelineSession``.""" pipeline_session = PipelineSession( boto_session=boto3.Session(region_name="us-west-2"), sagemaker_client=MagicMock(name="sagemaker_client"), @@ -103,10 +173,7 @@ def test_pipeline_variable_inside_arguments_compiles_to_get_expression(): output_location = ParameterString(name="OutputLocation") step = AthenaStep( name="RunQuery", - arguments={ - "QueryString": "SELECT 1", - "ResultConfiguration": {"OutputLocation": output_location}, - }, + query_config=AthenaQueryConfig(query_string="SELECT 1", output_location=output_location), ) pipeline = Pipeline( name="athena-unit", @@ -115,9 +182,8 @@ def test_pipeline_variable_inside_arguments_compiles_to_get_expression(): sagemaker_session=pipeline_session, ) - definition = json.loads(pipeline.definition()) + compiled_step = json.loads(pipeline.definition())["Steps"][0] - compiled_step = definition["Steps"][0] assert compiled_step["Type"] == "Athena" assert compiled_step["Arguments"]["ResultConfiguration"]["OutputLocation"] == { "Get": "Parameters.OutputLocation" @@ -127,20 +193,19 @@ def test_pipeline_variable_inside_arguments_compiles_to_get_expression(): def test_run_conditions_compile_to_depends_on_triggers(): - query = AthenaStep(name="Query", arguments={"QueryString": "SELECT 1"}) + config = AthenaQueryConfig(query_string="SELECT 1") + query = AthenaStep(name="Query", query_config=config) on_failure = AthenaStep( name="OnFailure", - arguments={"QueryString": "SELECT 1"}, + query_config=config, depends_on=[StepDependency(query, TriggerCondition.FAILED)], ) on_any = AthenaStep( name="OnAny", - arguments={"QueryString": "SELECT 1"}, + query_config=config, depends_on=[StepDependency("Query", TriggerCondition.ANY)], ) - on_success = AthenaStep( - name="OnSuccess", arguments={"QueryString": "SELECT 1"}, depends_on=[query] - ) + on_success = AthenaStep(name="OnSuccess", query_config=config, depends_on=[query]) pipeline = Pipeline( name="athena-run-conditions", steps=[query, on_failure, on_any, on_success], @@ -165,7 +230,7 @@ def test_run_conditions_compile_to_depends_on_triggers(): def test_properties_shape_and_reference_paths(): - step = AthenaStep(name="RunQuery", arguments={"QueryString": "SELECT 1"}) + step = AthenaStep(name="RunQuery", query_config=AthenaQueryConfig(query_string="SELECT 1")) assert isinstance(step.properties, Properties) # Loaded from botocore's GetQueryExecutionOutput shape (service_name="athena"); @@ -192,7 +257,7 @@ def test_properties_shape_and_reference_paths(): def test_to_request_includes_step_fields_cache_config_and_retry_policies(): step = AthenaStep( name="QueryUCData", - arguments=FULL_ARGUMENTS, + query_config=FULL_QUERY_CONFIG, display_name="Query UC Transactions", description="Runs Athena query against Unity Catalog", depends_on=["Upstream"], @@ -227,7 +292,9 @@ def test_to_request_includes_step_fields_cache_config_and_retry_policies(): def test_to_request_omits_optional_sections_when_not_set(): - request_dict = AthenaStep(name="RunQuery", arguments={"QueryString": "SELECT 1"}).to_request() + request_dict = AthenaStep( + name="RunQuery", query_config=AthenaQueryConfig(query_string="SELECT 1") + ).to_request() assert "CacheConfig" not in request_dict assert "RetryPolicies" not in request_dict assert "DependsOn" not in request_dict @@ -238,10 +305,12 @@ def test_to_request_omits_optional_sections_when_not_set(): # --------------------------------------------------------------------------- -def test_athena_step_exported_from_workflow_package(): +def test_athena_classes_exported_from_workflow_package(): + from sagemaker.mlops.workflow import AthenaQueryConfig as ExportedConfig from sagemaker.mlops.workflow import AthenaStep as ExportedAthenaStep assert ExportedAthenaStep is AthenaStep + assert ExportedConfig is AthenaQueryConfig def test_step_type_enum_has_athena_value(): From 47e9d5a9b0484291de6e56d8e20aa89b908a4a1b Mon Sep 17 00:00:00 2001 From: nayan3107 Date: Thu, 22 Oct 2026 09:33:03 -0700 Subject: [PATCH 3/4] refactor(mlops): Collapse AthenaStep query_config and arguments Take a single query_config that is either an AthenaQueryConfig or a raw StartQueryExecution dict, so customers have one field to set. The docstring prefers AthenaQueryConfig and keeps the dict for fields it does not cover. --- X-AI-Prompt: Collapse query_config and arguments into one union param X-AI-Tool: kiro-cli --- .../sagemaker/mlops/workflow/athena_step.py | 39 +++++++------------ .../tests/unit/workflow/test_athena_step.py | 23 ++++------- 2 files changed, 20 insertions(+), 42 deletions(-) diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/athena_step.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/athena_step.py index 585460275f..1e113c0eff 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/athena_step.py +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/athena_step.py @@ -32,9 +32,9 @@ ), ) -For ``StartQueryExecution`` fields that ``AthenaQueryConfig`` does not cover, -pass the full request as a dict through ``arguments`` instead; it is emitted -unchanged. +Prefer ``AthenaQueryConfig``. For ``StartQueryExecution`` fields it does not +cover, pass the full request as a dict to ``query_config`` instead; the dict +is emitted unchanged. The SageMaker Pipelines service validates the request at ``CreatePipeline`` time, not this module: @@ -178,23 +178,13 @@ def _drop_none(values: Dict[str, Any]) -> Dict[str, Any]: return {key: value for key, value in values.items() if value is not None} -ERR_STR_WITH_BOTH_QUERY_CONFIG_AND_ARGUMENTS = ( - "AthenaStep {step_name} cannot have both query_config and arguments." -) - -ERR_STR_WITHOUT_QUERY_CONFIG_AND_ARGUMENTS = ( - "AthenaStep {step_name} must have either query_config or arguments." -) - - class AthenaStep(ConfigurableRetryStep): """Runs an Amazon Athena query as part of a SageMaker Pipeline.""" def __init__( self, name: str, - query_config: Optional[AthenaQueryConfig] = None, - arguments: Optional[Dict[str, Any]] = None, + query_config: Union[AthenaQueryConfig, Dict[str, Any]], display_name: Optional[str] = None, description: Optional[str] = None, depends_on: Optional[List[Union[str, Step, StepCollection]]] = None, @@ -203,14 +193,13 @@ def __init__( ): # pylint: disable=too-many-positional-arguments """Construct an ``AthenaStep``. - Exactly one of ``query_config`` and ``arguments`` must be given. - Args: name: The name of the step. - query_config: The Athena query to run. - arguments: The full Athena ``StartQueryExecution`` request as a dict - (PascalCase keys), emitted unchanged. Use it for request fields - that ``AthenaQueryConfig`` does not cover. + query_config: The Athena query to run. Prefer an + ``AthenaQueryConfig``. Pass the full ``StartQueryExecution`` + request as a dict (PascalCase keys) only when you need a field + ``AthenaQueryConfig`` does not cover; the dict is emitted + unchanged. display_name: Optional display name. description: Optional description. depends_on: Optional explicit step dependencies. Wrap an entry in @@ -228,14 +217,12 @@ def __init__( retry_policies=retry_policies, ) - if query_config is None and arguments is None: - raise ValueError(ERR_STR_WITHOUT_QUERY_CONFIG_AND_ARGUMENTS.format(step_name=name)) - if query_config is not None and arguments is not None: - raise ValueError(ERR_STR_WITH_BOTH_QUERY_CONFIG_AND_ARGUMENTS.format(step_name=name)) - # Copy so caller-side mutations after construction do not leak into the # compiled pipeline definition. - request = query_config.to_request() if query_config is not None else arguments + if isinstance(query_config, AthenaQueryConfig): + request = query_config.to_request() + else: + request = query_config self.args: Dict[str, Any] = copy.deepcopy(request) self.cache_config = cache_config self._properties = Properties( diff --git a/sagemaker-mlops/tests/unit/workflow/test_athena_step.py b/sagemaker-mlops/tests/unit/workflow/test_athena_step.py index 972980bc87..1d2b9c0b8e 100644 --- a/sagemaker-mlops/tests/unit/workflow/test_athena_step.py +++ b/sagemaker-mlops/tests/unit/workflow/test_athena_step.py @@ -115,36 +115,27 @@ def test_step_with_query_config_emits_built_request(): # --------------------------------------------------------------------------- -def test_raw_arguments_passed_through_unchanged(): - arguments = { +def test_raw_request_dict_passed_through_unchanged(): + request = { "QueryString": "SELECT 1", "EngineConfiguration": {"MaxConcurrentDpus": 4}, } - step = AthenaStep(name="RunQuery", arguments=arguments) - assert step.arguments == arguments + step = AthenaStep(name="RunQuery", query_config=request) + assert step.arguments == request -def test_requires_query_config_or_arguments(): - with pytest.raises(ValueError, match="must have either query_config or arguments"): +def test_query_config_is_required(): + with pytest.raises(TypeError): AthenaStep(name="RunQuery") -def test_rejects_both_query_config_and_arguments(): - with pytest.raises(ValueError, match="cannot have both query_config and arguments"): - AthenaStep( - name="RunQuery", - query_config=AthenaQueryConfig(query_string="SELECT 1"), - arguments={"QueryString": "SELECT 1"}, - ) - - def test_caller_mutations_do_not_leak_into_step_state(): arguments = { "QueryString": "SELECT 1", "ResultConfiguration": {"EncryptionConfiguration": {"EncryptionOption": "SSE_S3"}}, "ExecutionParameters": ["1"], } - step = AthenaStep(name="RunQuery", arguments=arguments) + step = AthenaStep(name="RunQuery", query_config=arguments) arguments["QueryString"] = "MUTATED" arguments["ResultConfiguration"]["EncryptionConfiguration"]["EncryptionOption"] = "MUTATED" From 0efd6c2a202e2cf43aa2d055bb3130a02313a6e8 Mon Sep 17 00:00:00 2001 From: nayan3107 Date: Thu, 22 Oct 2026 09:54:22 -0700 Subject: [PATCH 4/4] test(mlops): Add AthenaStep integ test for raw request dict --- X-AI-Prompt: Add an integ test for the raw dict path X-AI-Tool: kiro-cli --- .../tests/integ/workflow/test_athena_step.py | 41 ++++++++++++++++++- 1 file changed, 40 insertions(+), 1 deletion(-) diff --git a/sagemaker-mlops/tests/integ/workflow/test_athena_step.py b/sagemaker-mlops/tests/integ/workflow/test_athena_step.py index 01939eb3bd..ef8d07100b 100644 --- a/sagemaker-mlops/tests/integ/workflow/test_athena_step.py +++ b/sagemaker-mlops/tests/integ/workflow/test_athena_step.py @@ -15,7 +15,8 @@ Runs Athena queries end-to-end through SageMaker Pipelines and checks that: * a valid query drives the execution to ``Succeeded`` and the result CSV is - written to the configured S3 output location, + written to the configured S3 output location, both when the query is built + with ``AthenaQueryConfig`` and when it is passed as a raw request dict, * an invalid query drives the execution to ``Failed`` with a ``FailureReason``, * a full ``AthenaQueryConfig``, together with ``CacheConfig`` and ``RetryPolicies``, is accepted by ``CreatePipeline`` and round-trips @@ -133,6 +134,44 @@ def test_athena_step_query_succeeds_and_writes_result(sagemaker_session_for_pipe _cleanup(pipeline, execution, sagemaker_session_for_pipeline, bucket, key_prefix) +def test_athena_step_raw_request_dict_succeeds(sagemaker_session_for_pipeline, role): + """A raw ``StartQueryExecution`` dict passed as ``query_config`` runs to ``Succeeded``.""" + pipeline_name = unique_name_from_base("athena-step-integ") + bucket, key_prefix, output_location = _result_location( + sagemaker_session_for_pipeline, pipeline_name + ) + output_location_parameter = ParameterString(name=_OUTPUT_LOCATION_PARAMETER) + + step = AthenaStep( + name="AthenaRawQuery", + query_config={ + "QueryString": "SELECT 1 AS test_result", + "QueryExecutionContext": {"Database": "default"}, + "WorkGroup": "primary", + "ResultConfiguration": {"OutputLocation": output_location_parameter}, + }, + ) + pipeline = Pipeline( + name=pipeline_name, + parameters=[output_location_parameter], + steps=[step], + sagemaker_session=sagemaker_session_for_pipeline, + ) + + execution = None + try: + execution, status, athena_step = _create_and_execute(pipeline, role, output_location) + assert status == "Succeeded", "Execution status was {}".format(status) + assert athena_step["StepStatus"] == "Succeeded", athena_step.get("FailureReason") + + query_execution_id = athena_step["Metadata"]["Athena"]["QueryExecutionId"] + result_key = "{}/{}.csv".format(key_prefix, query_execution_id) + result_files = sagemaker_session_for_pipeline.list_s3_files(bucket, key_prefix) + assert result_key in result_files, "Athena result file not found: {}".format(result_files) + finally: + _cleanup(pipeline, execution, sagemaker_session_for_pipeline, bucket, key_prefix) + + def test_athena_step_invalid_query_fails(sagemaker_session_for_pipeline, role): """An invalid SQL statement drives the step, and the execution, to ``Failed``.""" pipeline_name = unique_name_from_base("athena-step-integ")