apache-airflow-providers-amazon 8.4.0__tar.gz → 8.5.0rc1__tar.gz

{apache-airflow-providers-amazon-8.4.0/apache_airflow_providers_amazon.egg-info → apache-airflow-providers-amazon-8.5.0rc1}/PKG-INFO

@@ -1,14 +1,14 @@
 Metadata-Version: 2.1
 Name: apache-airflow-providers-amazon
-Version: 8.4.0
+Version: 8.5.0rc1
 Summary: Provider for Apache Airflow. Implements apache-airflow-providers-amazon package
 Home-page: https://airflow.apache.org/
 Download-URL: https://archive.apache.org/dist/airflow/providers
 Author: Apache Software Foundation
 Author-email: dev@airflow.apache.org
 License: Apache License 2.0
-Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-amazon/8.4.0/
-Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-amazon/8.4.0/changelog.html
+Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-amazon/8.5.0/
+Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-amazon/8.5.0/changelog.html
 Project-URL: Bug Tracker, https://github.com/apache/airflow/issues
 Project-URL: Source Code, https://github.com/apache/airflow
 Project-URL: Slack Chat, https://s.apache.org/airflow-slack
@@ -39,6 +39,7 @@ Provides-Extra: http
 Provides-Extra: imap
 Provides-Extra: microsoft.azure
 Provides-Extra: mongo
+Provides-Extra: openlineage
 Provides-Extra: salesforce
 Provides-Extra: ssh
 Provides-Extra: pandas
@@ -84,7 +85,7 @@ License-File: NOTICE
 
 Package ``apache-airflow-providers-amazon``
 
-Release: ``8.4.0``
+Release: ``8.5.0rc1``
 
 
 Amazon integration (including `Amazon Web Services (AWS) <https://aws.amazon.com/>`__).
@@ -97,7 +98,7 @@ This is a provider package for ``amazon`` provider. All classes for this provide
 are in ``airflow.providers.amazon`` python package.
 
 You can find package information and changelog for the provider
-in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-amazon/8.4.0/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-amazon/8.5.0/>`_.
 
 
 Installation
@@ -157,9 +158,10 @@ Dependent package
 `apache-airflow-providers-imap <https://airflow.apache.org/docs/apache-airflow-providers-imap>`_                        ``imap``
 `apache-airflow-providers-microsoft-azure <https://airflow.apache.org/docs/apache-airflow-providers-microsoft-azure>`_  ``microsoft.azure``
 `apache-airflow-providers-mongo <https://airflow.apache.org/docs/apache-airflow-providers-mongo>`_                      ``mongo``
+`apache-airflow-providers-openlineage <https://airflow.apache.org/docs/apache-airflow-providers-openlineage>`_          ``openlineage``
 `apache-airflow-providers-salesforce <https://airflow.apache.org/docs/apache-airflow-providers-salesforce>`_            ``salesforce``
 `apache-airflow-providers-ssh <https://airflow.apache.org/docs/apache-airflow-providers-ssh>`_                          ``ssh``
 ======================================================================================================================  ===================
 
 The changelog for the provider package can be found in the
-`changelog <https://airflow.apache.org/docs/apache-airflow-providers-amazon/8.4.0/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-amazon/8.5.0/changelog.html>`_.

{apache-airflow-providers-amazon-8.4.0 → apache-airflow-providers-amazon-8.5.0rc1}/README.rst

@@ -36,7 +36,7 @@
 
 Package ``apache-airflow-providers-amazon``
 
-Release: ``8.4.0``
+Release: ``8.5.0rc1``
 
 
 Amazon integration (including `Amazon Web Services (AWS) <https://aws.amazon.com/>`__).
@@ -49,7 +49,7 @@ This is a provider package for ``amazon`` provider. All classes for this provide
 are in ``airflow.providers.amazon`` python package.
 
 You can find package information and changelog for the provider
-in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-amazon/8.4.0/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-amazon/8.5.0/>`_.
 
 
 Installation
@@ -109,9 +109,10 @@ Dependent package
 `apache-airflow-providers-imap <https://airflow.apache.org/docs/apache-airflow-providers-imap>`_                        ``imap``
 `apache-airflow-providers-microsoft-azure <https://airflow.apache.org/docs/apache-airflow-providers-microsoft-azure>`_  ``microsoft.azure``
 `apache-airflow-providers-mongo <https://airflow.apache.org/docs/apache-airflow-providers-mongo>`_                      ``mongo``
+`apache-airflow-providers-openlineage <https://airflow.apache.org/docs/apache-airflow-providers-openlineage>`_          ``openlineage``
 `apache-airflow-providers-salesforce <https://airflow.apache.org/docs/apache-airflow-providers-salesforce>`_            ``salesforce``
 `apache-airflow-providers-ssh <https://airflow.apache.org/docs/apache-airflow-providers-ssh>`_                          ``ssh``
 ======================================================================================================================  ===================
 
 The changelog for the provider package can be found in the
-`changelog <https://airflow.apache.org/docs/apache-airflow-providers-amazon/8.4.0/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-amazon/8.5.0/changelog.html>`_.

{apache-airflow-providers-amazon-8.4.0 → apache-airflow-providers-amazon-8.5.0rc1}/airflow/providers/amazon/__init__.py

@@ -28,7 +28,7 @@ import packaging.version
 
 __all__ = ["__version__"]
 
-__version__ = "8.4.0"
+__version__ = "8.5.0"
 
 try:
     from airflow import __version__ as airflow_version

apache-airflow-providers-amazon-8.5.0rc1/airflow/providers/amazon/aws/hooks/eventbridge.py

@@ -0,0 +1,90 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License 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.
+from __future__ import annotations
+
+import json
+
+from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+from airflow.providers.amazon.aws.utils import trim_none_values
+
+
+def _validate_json(pattern: str) -> None:
+    try:
+        json.loads(pattern)
+    except ValueError:
+        raise ValueError("`event_pattern` must be a valid JSON string.")
+
+
+class EventBridgeHook(AwsBaseHook):
+    """Amazon EventBridge Hook."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(client_type="events", *args, **kwargs)
+
+    def put_rule(
+        self,
+        name: str,
+        description: str | None = None,
+        event_bus_name: str | None = None,
+        event_pattern: str | None = None,
+        role_arn: str | None = None,
+        schedule_expression: str | None = None,
+        state: str | None = None,
+        tags: list[dict] | None = None,
+        **kwargs,
+    ):
+        """
+        Create or update an EventBridge rule.
+
+        :param name: name of the rule to create or update (required)
+        :param description: description of the rule
+        :param event_bus_name: name or ARN of the event bus to associate with this rule
+        :param event_pattern: pattern of events to be matched to this rule
+        :param role_arn: the Amazon Resource Name of the IAM role associated with the rule
+        :param schedule_expression: the scheduling expression (for example, a cron or rate expression)
+        :param state: indicates whether rule is set to be "ENABLED" or "DISABLED"
+        :param tags: list of key-value pairs to associate with the rule
+
+        """
+        if not (event_pattern or schedule_expression):
+            raise ValueError(
+                "One of `event_pattern` or `schedule_expression` are required in order to "
+                "put or update your rule."
+            )
+
+        if state and state not in ["ENABLED", "DISABLED"]:
+            raise ValueError("`state` must be specified as ENABLED or DISABLED.")
+
+        if event_pattern:
+            _validate_json(event_pattern)
+
+        put_rule_kwargs: dict[str, str | list] = {
+            **trim_none_values(
+                {
+                    "Name": name,
+                    "Description": description,
+                    "EventBusName": event_bus_name,
+                    "EventPattern": event_pattern,
+                    "RoleArn": role_arn,
+                    "ScheduleExpression": schedule_expression,
+                    "State": state,
+                    "Tags": tags,
+                }
+            )
+        }
+
+        return self.conn.put_rule(**put_rule_kwargs)

{apache-airflow-providers-amazon-8.4.0 → apache-airflow-providers-amazon-8.5.0rc1}/airflow/providers/amazon/aws/hooks/s3.py

@@ -57,7 +57,7 @@ from airflow.providers.amazon.aws.utils.tags import format_tags
 from airflow.utils.helpers import chunks
 
 if TYPE_CHECKING:
-    from mypy_boto3_s3.service_resource import Object as S3ResourceObject
+    from mypy_boto3_s3.service_resource import Bucket as S3Bucket, Object as S3ResourceObject
 
 T = TypeVar("T", bound=Callable)
 
@@ -298,7 +298,7 @@ class S3Hook(AwsBaseHook):
             return False
 
     @provide_bucket_name
-    def get_bucket(self, bucket_name: str | None = None) -> object:
+    def get_bucket(self, bucket_name: str | None = None) -> S3Bucket:
         """
         Returns a :py:class:`S3.Bucket` object.
 

{apache-airflow-providers-amazon-8.4.0 → apache-airflow-providers-amazon-8.5.0rc1}/airflow/providers/amazon/aws/operators/batch.py

@@ -149,7 +149,7 @@ class BatchOperator(BaseOperator):
         parameters: dict | None = None,
         job_id: str | None = None,
         waiters: Any | None = None,
-        max_retries: int | None = None,
+        max_retries: int = 4200,
         status_retries: int | None = None,
         aws_conn_id: str | None = None,
         region_name: str | None = None,
@@ -223,7 +223,7 @@ class BatchOperator(BaseOperator):
                 timeout=self.execution_timeout,
                 trigger=BatchJobTrigger(
                     job_id=self.job_id,
-                    waiter_max_attempts=self.max_retries or 10,
+                    waiter_max_attempts=self.max_retries,
                     aws_conn_id=self.aws_conn_id,
                     region_name=self.region_name,
                     waiter_delay=self.poll_interval,

{apache-airflow-providers-amazon-8.4.0 → apache-airflow-providers-amazon-8.5.0rc1}/airflow/providers/amazon/aws/operators/eventbridge.py

@@ -32,6 +32,10 @@ class EventBridgePutEventsOperator(BaseOperator):
     """
     Put Events onto Amazon EventBridge.
 
+    .. seealso::
+        For more information on how to use this operator, take a look at the guide:
+        :ref:`howto/operator:EventBridgePutEventsOperator`
+
     :param entries: the list of events to be put onto EventBridge, each event is a dict (required)
     :param endpoint_id: the URL subdomain of the endpoint
     :param aws_conn_id: the AWS connection to use
@@ -85,3 +89,84 @@ class EventBridgePutEventsOperator(BaseOperator):
 
         if self.do_xcom_push:
             return [e["EventId"] for e in response["Entries"]]
+
+
+class EventBridgePutRuleOperator(BaseOperator):
+    """
+    Create or update a specified EventBridge rule.
+
+    .. seealso::
+        For more information on how to use this operator, take a look at the guide:
+        :ref:`howto/operator:EventBridgePutRuleOperator`
+
+    :param name: name of the rule to create or update (required)
+    :param description: description of the rule
+    :param event_bus_name: name or ARN of the event bus to associate with this rule
+    :param event_pattern: pattern of events to be matched to this rule
+    :param role_arn: the Amazon Resource Name of the IAM role associated with the rule
+    :param schedule_expression: the scheduling expression (for example, a cron or rate expression)
+    :param state: indicates whether rule is set to be "ENABLED" or "DISABLED"
+    :param tags: list of key-value pairs to associate with the rule
+    :param region: the region where rule is to be created or updated
+
+    """
+
+    template_fields: Sequence[str] = (
+        "aws_conn_id",
+        "name",
+        "description",
+        "event_bus_name",
+        "event_pattern",
+        "role_arn",
+        "schedule_expression",
+        "state",
+        "tags",
+        "region_name",
+    )
+
+    def __init__(
+        self,
+        *,
+        name: str,
+        description: str | None = None,
+        event_bus_name: str | None = None,
+        event_pattern: str | None = None,
+        role_arn: str | None = None,
+        schedule_expression: str | None = None,
+        state: str | None = None,
+        tags: list | None = None,
+        region_name: str | None = None,
+        aws_conn_id: str = "aws_default",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.name = name
+        self.description = description
+        self.event_bus_name = event_bus_name
+        self.event_pattern = event_pattern
+        self.role_arn = role_arn
+        self.region_name = region_name
+        self.schedule_expression = schedule_expression
+        self.state = state
+        self.tags = tags
+        self.aws_conn_id = aws_conn_id
+
+    @cached_property
+    def hook(self) -> EventBridgeHook:
+        """Create and return an EventBridgeHook."""
+        return EventBridgeHook(aws_conn_id=self.aws_conn_id, region_name=self.region_name)
+
+    def execute(self, context: Context):
+
+        self.log.info('Sending rule "%s" to EventBridge.', self.name)
+
+        return self.hook.put_rule(
+            name=self.name,
+            description=self.description,
+            event_bus_name=self.event_bus_name,
+            event_pattern=self.event_pattern,
+            role_arn=self.role_arn,
+            schedule_expression=self.schedule_expression,
+            state=self.state,
+            tags=self.tags,
+        )

{apache-airflow-providers-amazon-8.4.0 → apache-airflow-providers-amazon-8.5.0rc1}/airflow/providers/amazon/aws/operators/rds.py

@@ -392,6 +392,8 @@ class RdsCancelExportTaskOperator(RdsBaseOperator):
 
     :param export_task_identifier: The identifier of the snapshot export task to cancel
     :param wait_for_completion:  If True, waits for DB snapshot export to cancel. (default: True)
+    :param check_interval: The amount of time in seconds to wait between attempts
+    :param max_attempts: The maximum number of attempts to be made
     """
 
     template_fields = ("export_task_identifier",)
@@ -402,6 +404,7 @@ class RdsCancelExportTaskOperator(RdsBaseOperator):
         export_task_identifier: str,
         wait_for_completion: bool = True,
         check_interval: int = 30,
+        max_attempts: int = 40,
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -409,6 +412,7 @@ class RdsCancelExportTaskOperator(RdsBaseOperator):
         self.export_task_identifier = export_task_identifier
         self.wait_for_completion = wait_for_completion
         self.check_interval = check_interval
+        self.max_attempts = max_attempts
 
     def execute(self, context: Context) -> str:
         self.log.info("Canceling export task %s", self.export_task_identifier)
@@ -419,7 +423,10 @@ class RdsCancelExportTaskOperator(RdsBaseOperator):
 
         if self.wait_for_completion:
             self.hook.wait_for_export_task_state(
-                self.export_task_identifier, target_state="canceled", check_interval=self.check_interval
+                self.export_task_identifier,
+                target_state="canceled",
+                check_interval=self.check_interval,
+                max_attempts=self.max_attempts,
             )
         return json.dumps(cancel_export, default=str)
 

{apache-airflow-providers-amazon-8.4.0 → apache-airflow-providers-amazon-8.5.0rc1}/airflow/providers/amazon/aws/operators/sagemaker.py

@@ -40,13 +40,16 @@ from airflow.providers.amazon.aws.utils.tags import format_tags
 from airflow.utils.json import AirflowJsonEncoder
 
 if TYPE_CHECKING:
+    from openlineage.client.run import Dataset
+
+    from airflow.providers.openlineage.extractors.base import OperatorLineage
     from airflow.utils.context import Context
 
 DEFAULT_CONN_ID: str = "aws_default"
 CHECK_INTERVAL_SECOND: int = 30
 
 
-def serialize(result: dict) -> str:
+def serialize(result: dict) -> dict:
     return json.loads(json.dumps(result, cls=AirflowJsonEncoder))
 
 
@@ -158,6 +161,14 @@ class SageMakerBaseOperator(BaseOperator):
         """Return SageMakerHook."""
         return SageMakerHook(aws_conn_id=self.aws_conn_id)
 
+    @staticmethod
+    def path_to_s3_dataset(path) -> Dataset:
+        from openlineage.client.run import Dataset
+
+        path = path.replace("s3://", "")
+        split_path = path.split("/")
+        return Dataset(namespace=f"s3://{split_path[0]}", name="/".join(split_path[1:]), facets={})
+
 
 class SageMakerProcessingOperator(SageMakerBaseOperator):
     """
@@ -225,6 +236,7 @@ class SageMakerProcessingOperator(SageMakerBaseOperator):
         self.max_attempts = max_attempts or 60
         self.max_ingestion_time = max_ingestion_time
         self.deferrable = deferrable
+        self.serialized_job: dict
 
     def _create_integer_fields(self) -> None:
         """Set fields which should be cast to integers."""
@@ -282,14 +294,48 @@ class SageMakerProcessingOperator(SageMakerBaseOperator):
                 method_name="execute_complete",
             )
 
-        return {"Processing": serialize(self.hook.describe_processing_job(self.config["ProcessingJobName"]))}
+        self.serialized_job = serialize(self.hook.describe_processing_job(self.config["ProcessingJobName"]))
+        return {"Processing": self.serialized_job}
 
     def execute_complete(self, context, event=None):
         if event["status"] != "success":
             raise AirflowException(f"Error while running job: {event}")
         else:
             self.log.info(event["message"])
-        return {"Processing": serialize(self.hook.describe_processing_job(self.config["ProcessingJobName"]))}
+        self.serialized_job = serialize(self.hook.describe_processing_job(self.config["ProcessingJobName"]))
+        return {"Processing": self.serialized_job}
+
+    def get_openlineage_facets_on_complete(self, task_instance) -> OperatorLineage:
+        """Returns OpenLineage data gathered from SageMaker's API response saved by processing job."""
+        from airflow.providers.openlineage.extractors.base import OperatorLineage
+
+        inputs = []
+        outputs = []
+        try:
+            inputs, outputs = self._extract_s3_dataset_identifiers(
+                processing_inputs=self.serialized_job["ProcessingInputs"],
+                processing_outputs=self.serialized_job["ProcessingOutputConfig"]["Outputs"],
+            )
+        except KeyError:
+            self.log.exception("Could not find input/output information in Xcom.")
+
+        return OperatorLineage(inputs=inputs, outputs=outputs)
+
+    def _extract_s3_dataset_identifiers(self, processing_inputs, processing_outputs):
+        inputs = []
+        outputs = []
+        try:
+            for processing_input in processing_inputs:
+                inputs.append(self.path_to_s3_dataset(processing_input["S3Input"]["S3Uri"]))
+        except KeyError:
+            self.log.exception("Cannot find S3 input details", exc_info=True)
+
+        try:
+            for processing_output in processing_outputs:
+                outputs.append(self.path_to_s3_dataset(processing_output["S3Output"]["S3Uri"]))
+        except KeyError:
+            self.log.exception("Cannot find S3 output details.", exc_info=True)
+        return inputs, outputs
 
 
 class SageMakerEndpointConfigOperator(SageMakerBaseOperator):
@@ -579,6 +625,8 @@ class SageMakerTransformOperator(SageMakerBaseOperator):
                 Provided value: '{action_if_job_exists}'."
             )
         self.deferrable = deferrable
+        self.serialized_model: dict
+        self.serialized_tranform: dict
 
     def _create_integer_fields(self) -> None:
         """Set fields which should be cast to integers."""
@@ -650,10 +698,11 @@ class SageMakerTransformOperator(SageMakerBaseOperator):
                 method_name="execute_complete",
             )
 
-        return {
-            "Model": serialize(self.hook.describe_model(transform_config["ModelName"])),
-            "Transform": serialize(self.hook.describe_transform_job(transform_config["TransformJobName"])),
-        }
+        self.serialized_model = serialize(self.hook.describe_model(transform_config["ModelName"]))
+        self.serialized_tranform = serialize(
+            self.hook.describe_transform_job(transform_config["TransformJobName"])
+        )
+        return {"Model": self.serialized_model, "Transform": self.serialized_tranform}
 
     def execute_complete(self, context, event=None):
         if event["status"] != "success":
@@ -661,10 +710,62 @@ class SageMakerTransformOperator(SageMakerBaseOperator):
         else:
             self.log.info(event["message"])
         transform_config = self.config.get("Transform", self.config)
-        return {
-            "Model": serialize(self.hook.describe_model(transform_config["ModelName"])),
-            "Transform": serialize(self.hook.describe_transform_job(transform_config["TransformJobName"])),
-        }
+        self.serialized_model = serialize(self.hook.describe_model(transform_config["ModelName"]))
+        self.serialized_tranform = serialize(
+            self.hook.describe_transform_job(transform_config["TransformJobName"])
+        )
+        return {"Model": self.serialized_model, "Transform": self.serialized_tranform}
+
+    def get_openlineage_facets_on_complete(self, task_instance) -> OperatorLineage:
+        """Returns OpenLineage data gathered from SageMaker's API response saved by transform job."""
+        from airflow.providers.openlineage.extractors import OperatorLineage
+
+        model_package_arn = None
+        transform_input = None
+        transform_output = None
+
+        try:
+            model_package_arn = self.serialized_model["PrimaryContainer"]["ModelPackageName"]
+        except KeyError:
+            self.log.error("Cannot find Model Package Name.", exc_info=True)
+
+        try:
+            transform_input = self.serialized_tranform["TransformInput"]["DataSource"]["S3DataSource"][
+                "S3Uri"
+            ]
+            transform_output = self.serialized_tranform["TransformOutput"]["S3OutputPath"]
+        except KeyError:
+            self.log.error("Cannot find some required input/output details.", exc_info=True)
+
+        inputs = []
+
+        if transform_input is not None:
+            inputs.append(self.path_to_s3_dataset(transform_input))
+
+        if model_package_arn is not None:
+            model_data_urls = self._get_model_data_urls(model_package_arn)
+            for model_data_url in model_data_urls:
+                inputs.append(self.path_to_s3_dataset(model_data_url))
+
+        outputs = []
+        if transform_output is not None:
+            outputs.append(self.path_to_s3_dataset(transform_output))
+
+        return OperatorLineage(inputs=inputs, outputs=outputs)
+
+    def _get_model_data_urls(self, model_package_arn) -> list:
+        model_data_urls = []
+        try:
+            model_containers = self.hook.get_conn().describe_model_package(
+                ModelPackageName=model_package_arn
+            )["InferenceSpecification"]["Containers"]
+
+            for container in model_containers:
+                model_data_urls.append(container["ModelDataUrl"])
+        except KeyError:
+            self.log.exception("Cannot retrieve model details.", exc_info=True)
+
+        return model_data_urls
 
 
 class SageMakerTuningOperator(SageMakerBaseOperator):
@@ -891,6 +992,7 @@ class SageMakerTrainingOperator(SageMakerBaseOperator):
                 Provided value: '{action_if_job_exists}'."
             )
         self.deferrable = deferrable
+        self.serialized_training_data: dict
 
     def expand_role(self) -> None:
         """Expands an IAM role name into an ARN."""
@@ -951,16 +1053,40 @@ class SageMakerTrainingOperator(SageMakerBaseOperator):
                 method_name="execute_complete",
             )
 
-        result = {"Training": serialize(self.hook.describe_training_job(self.config["TrainingJobName"]))}
-        return result
+        self.serialized_training_data = serialize(
+            self.hook.describe_training_job(self.config["TrainingJobName"])
+        )
+        return {"Training": self.serialized_training_data}
 
     def execute_complete(self, context, event=None):
         if event["status"] != "success":
             raise AirflowException(f"Error while running job: {event}")
         else:
             self.log.info(event["message"])
-        result = {"Training": serialize(self.hook.describe_training_job(self.config["TrainingJobName"]))}
-        return result
+        self.serialized_training_data = serialize(
+            self.hook.describe_training_job(self.config["TrainingJobName"])
+        )
+        return {"Training": self.serialized_training_data}
+
+    def get_openlineage_facets_on_complete(self, task_instance) -> OperatorLineage:
+        """Returns OpenLineage data gathered from SageMaker's API response saved by training job."""
+        from airflow.providers.openlineage.extractors import OperatorLineage
+
+        inputs = []
+        outputs = []
+        try:
+            for input_data in self.serialized_training_data["InputDataConfig"]:
+                inputs.append(self.path_to_s3_dataset(input_data["DataSource"]["S3DataSource"]["S3Uri"]))
+        except KeyError:
+            self.log.exception("Issues extracting inputs.")
+
+        try:
+            outputs.append(
+                self.path_to_s3_dataset(self.serialized_training_data["ModelArtifacts"]["S3ModelArtifacts"])
+            )
+        except KeyError:
+            self.log.exception("Issues extracting inputs.")
+        return OperatorLineage(inputs=inputs, outputs=outputs)
 
 
 class SageMakerDeleteModelOperator(SageMakerBaseOperator):