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..b22f6137ce 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 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 @@ -113,6 +114,8 @@ "TuningStep", # Step implementations "AutoMLStep", + "AthenaQueryConfig", + "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..1e113c0eff --- /dev/null +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/athena_step.py @@ -0,0 +1,250 @@ +# 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. Describe the query with an ``AthenaQueryConfig``; it assembles the +Athena ``StartQueryExecution`` request that becomes the step's ``Arguments``:: + + from sagemaker.mlops.workflow import AthenaQueryConfig, AthenaStep, ParameterString + + output_location = ParameterString(name="OutputLocation") + + step = AthenaStep( + name="RunQuery", + 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'"], + ), + ) + +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: + +* ``QueryString`` is literal-only. Pipeline expressions such as + ``{"Get": "Parameters.X"}`` inside it are rejected to prevent SQL + 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 + 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 PipelineVariable, 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 AthenaQueryConfig: + """Config for an Athena query run by an ``AthenaStep``. + + 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} + + +class AthenaStep(ConfigurableRetryStep): + """Runs an Amazon Athena query as part of a SageMaker Pipeline.""" + + def __init__( + self, + name: str, + 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, + cache_config: Optional[CacheConfig] = None, + retry_policies: Optional[List[RetryPolicy]] = None, + ): # pylint: disable=too-many-positional-arguments + """Construct an ``AthenaStep``. + + Args: + name: The name of the step. + 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 + ``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, + ) + + # Copy so caller-side mutations after construction do not leak into the + # compiled pipeline definition. + 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( + 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..ef8d07100b --- /dev/null +++ b/sagemaker-mlops/tests/integ/workflow/test_athena_step.py @@ -0,0 +1,257 @@ +# 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, 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 + 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 AthenaQueryConfig, 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", + query_config=AthenaQueryConfig( + query_string="SELECT 1 AS test_result", + output_location=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_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") + 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", + query_config=AthenaQueryConfig( + query_string="SELCT INVALID SYNTAX FROM nonexistent_table_xyz", + output_location=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_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", + 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( + 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..1d2b9c0b8e --- /dev/null +++ b/sagemaker-mlops/tests/unit/workflow/test_athena_step.py @@ -0,0 +1,308 @@ +# 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 and AthenaQueryConfig.""" + +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 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 + +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 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_KMS", + "KmsKey": "arn:aws:kms:us-west-2:111122223333:key/example", + }, + "ExpectedBucketOwner": "111122223333", + "AclConfiguration": {"S3AclOption": "BUCKET_OWNER_FULL_CONTROL"}, + }, + "ResultReuseConfiguration": { + "ResultReuseByAgeConfiguration": {"Enabled": True, "MaxAgeInMinutes": 60} + }, +} + + +# --------------------------------------------------------------------------- +# 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( + "kwargs, expected", + [ + ({"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=["context", "output-location", "encryption", "acl"], +) +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_request_dict_passed_through_unchanged(): + request = { + "QueryString": "SELECT 1", + "EngineConfiguration": {"MaxConcurrentDpus": 4}, + } + step = AthenaStep(name="RunQuery", query_config=request) + assert step.arguments == request + + +def test_query_config_is_required(): + with pytest.raises(TypeError): + AthenaStep(name="RunQuery") + + +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", query_config=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"], + } + + +# --------------------------------------------------------------------------- +# 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"), + sagemaker_config={}, + ) + output_location = ParameterString(name="OutputLocation") + step = AthenaStep( + name="RunQuery", + query_config=AthenaQueryConfig(query_string="SELECT 1", output_location=output_location), + ) + pipeline = Pipeline( + name="athena-unit", + parameters=[output_location], + steps=[step], + sagemaker_session=pipeline_session, + ) + + compiled_step = json.loads(pipeline.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(): + config = AthenaQueryConfig(query_string="SELECT 1") + query = AthenaStep(name="Query", query_config=config) + on_failure = AthenaStep( + name="OnFailure", + query_config=config, + depends_on=[StepDependency(query, TriggerCondition.FAILED)], + ) + on_any = AthenaStep( + name="OnAny", + query_config=config, + depends_on=[StepDependency("Query", TriggerCondition.ANY)], + ) + 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], + 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", query_config=AthenaQueryConfig(query_string="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", + query_config=FULL_QUERY_CONFIG, + 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", 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 + + +# --------------------------------------------------------------------------- +# Package wiring +# --------------------------------------------------------------------------- + + +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(): + 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", }