feat(mlops): Add AthenaStep pipeline step

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)
 

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",

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
+    <https://docs.aws.amazon.com/athena/latest/APIReference/API_StartQueryExecution.html>`_
+    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

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):