apache-airflow-providers-amazon 9.0.0rc1__py3-none-any.whl → 9.1.0rc1__py3-none-any.whl

airflow/providers/amazon/__init__.py

@@ -29,7 +29,7 @@ from airflow import __version__ as airflow_version
 
 __all__ = ["__version__"]
 
-__version__ = "9.0.0"
+__version__ = "9.1.0"
 
 if packaging.version.parse(packaging.version.parse(airflow_version).base_version) < packaging.version.parse(
     "2.8.0"

airflow/providers/amazon/aws/assets/s3.py

@@ -19,16 +19,14 @@ from __future__ import annotations
 from typing import TYPE_CHECKING
 
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
-
-try:
-    from airflow.assets import Asset
-except ModuleNotFoundError:
-    from airflow.datasets import Dataset as Asset  # type: ignore[no-redef]
+from airflow.providers.common.compat.assets import Asset
 
 if TYPE_CHECKING:
     from urllib.parse import SplitResult
 
-    from airflow.providers.common.compat.openlineage.facet import Dataset as OpenLineageDataset
+    from airflow.providers.common.compat.openlineage.facet import (
+        Dataset as OpenLineageDataset,
+    )
 
 
 def create_asset(*, bucket: str, key: str, extra=None) -> Asset:
@@ -43,7 +41,9 @@ def sanitize_uri(uri: SplitResult) -> SplitResult:
 
 def convert_asset_to_openlineage(asset: Asset, lineage_context) -> OpenLineageDataset:
     """Translate Asset with valid AIP-60 uri to OpenLineage with assistance from the hook."""
-    from airflow.providers.common.compat.openlineage.facet import Dataset as OpenLineageDataset
+    from airflow.providers.common.compat.openlineage.facet import (
+        Dataset as OpenLineageDataset,
+    )
 
     bucket, key = S3Hook.parse_s3_url(asset.uri)
     return OpenLineageDataset(namespace=f"s3://{bucket}", name=key if key else "/")

airflow/providers/amazon/aws/hooks/appflow.py

@@ -117,9 +117,9 @@ class AppflowHook(AwsGenericHook["AppflowClient"]):
 
         self.conn.update_flow(
             flowName=response["flowName"],
-            destinationFlowConfigList=response["destinationFlowConfigList"],
-            sourceFlowConfig=response["sourceFlowConfig"],
-            triggerConfig=response["triggerConfig"],
+            destinationFlowConfigList=response["destinationFlowConfigList"],  # type: ignore[arg-type]
+            sourceFlowConfig=response["sourceFlowConfig"],  # type: ignore[arg-type]
+            triggerConfig=response["triggerConfig"],  # type: ignore[arg-type]
             description=response.get("description", "Flow description."),
-            tasks=tasks,
+            tasks=tasks,  # type: ignore[arg-type]
         )

airflow/providers/amazon/aws/hooks/redshift_sql.py

@@ -163,7 +163,7 @@ class RedshiftSQLHook(DbApiHook):
         # Compatibility: The 'create' factory method was added in SQLAlchemy 1.4
         # to replace calling the default URL constructor directly.
         create_url = getattr(URL, "create", URL)
-        return str(create_url(drivername="redshift+redshift_connector", **conn_params))
+        return str(create_url(drivername="postgresql", **conn_params))
 
     def get_sqlalchemy_engine(self, engine_kwargs=None):
         """Overridden to pass Redshift-specific arguments."""

airflow/providers/amazon/aws/hooks/s3.py

@@ -41,14 +41,16 @@ from urllib.parse import urlsplit
 from uuid import uuid4
 
 if TYPE_CHECKING:
-    from mypy_boto3_s3.service_resource import Bucket as S3Bucket, Object as S3ResourceObject
+    from mypy_boto3_s3.service_resource import (
+        Bucket as S3Bucket,
+        Object as S3ResourceObject,
+    )
 
     from airflow.utils.types import ArgNotSet
 
     with suppress(ImportError):
         from aiobotocore.client import AioBaseClient
 
-from importlib.util import find_spec
 
 from asgiref.sync import sync_to_async
 from boto3.s3.transfer import S3Transfer, TransferConfig
@@ -58,15 +60,9 @@ from airflow.exceptions import AirflowException, AirflowNotFoundException
 from airflow.providers.amazon.aws.exceptions import S3HookUriParseFailure
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
 from airflow.providers.amazon.aws.utils.tags import format_tags
+from airflow.providers.common.compat.lineage.hook import get_hook_lineage_collector
 from airflow.utils.helpers import chunks
 
-if find_spec("airflow.assets"):
-    from airflow.lineage.hook import get_hook_lineage_collector
-else:
-    # TODO: import from common.compat directly after common.compat providers with
-    # asset_compat_lineage_collector released
-    from airflow.providers.amazon.aws.utils.asset_compat_lineage_collector import get_hook_lineage_collector
-
 logger = logging.getLogger(__name__)
 
 
@@ -148,9 +144,10 @@ def unify_bucket_name_and_key(func: Callable) -> Callable:
 
         if "bucket_name" not in bound_args.arguments:
             with suppress(S3HookUriParseFailure):
-                bound_args.arguments["bucket_name"], bound_args.arguments[key_name] = S3Hook.parse_s3_url(
-                    bound_args.arguments[key_name]
-                )
+                (
+                    bound_args.arguments["bucket_name"],
+                    bound_args.arguments[key_name],
+                ) = S3Hook.parse_s3_url(bound_args.arguments[key_name])
 
         return func(*bound_args.args, **bound_args.kwargs)
 
@@ -318,7 +315,8 @@ class S3Hook(AwsBaseHook):
                 self.log.info('Bucket "%s" does not exist', bucket_name)
             elif return_code == 403:
                 self.log.error(
-                    'Access to bucket "%s" is forbidden or there was an error with the request', bucket_name
+                    'Access to bucket "%s" is forbidden or there was an error with the request',
+                    bucket_name,
                 )
                 self.log.error(e)
             return False
@@ -359,7 +357,8 @@ class S3Hook(AwsBaseHook):
             self.get_conn().create_bucket(Bucket=bucket_name)
         else:
             self.get_conn().create_bucket(
-                Bucket=bucket_name, CreateBucketConfiguration={"LocationConstraint": region_name}
+                Bucket=bucket_name,
+                CreateBucketConfiguration={"LocationConstraint": region_name},
             )
 
     @provide_bucket_name
@@ -410,7 +409,10 @@ class S3Hook(AwsBaseHook):
 
         paginator = self.get_conn().get_paginator("list_objects_v2")
         response = paginator.paginate(
-            Bucket=bucket_name, Prefix=prefix, Delimiter=delimiter, PaginationConfig=config
+            Bucket=bucket_name,
+            Prefix=prefix,
+            Delimiter=delimiter,
+            PaginationConfig=config,
         )
 
         prefixes: list[str] = []
@@ -471,7 +473,10 @@ class S3Hook(AwsBaseHook):
 
         paginator = client.get_paginator("list_objects_v2")
         response = paginator.paginate(
-            Bucket=bucket_name, Prefix=prefix, Delimiter=delimiter, PaginationConfig=config
+            Bucket=bucket_name,
+            Prefix=prefix,
+            Delimiter=delimiter,
+            PaginationConfig=config,
         )
 
         prefixes = []
@@ -569,7 +574,11 @@ class S3Hook(AwsBaseHook):
         return await self._check_key_async(client, bucket, wildcard_match, bucket_keys, use_regex)
 
     async def check_for_prefix_async(
-        self, client: AioBaseClient, prefix: str, delimiter: str, bucket_name: str | None = None
+        self,
+        client: AioBaseClient,
+        prefix: str,
+        delimiter: str,
+        bucket_name: str | None = None,
     ) -> bool:
         """
         Check that a prefix exists in a bucket.
@@ -587,7 +596,11 @@ class S3Hook(AwsBaseHook):
         return prefix in plist
 
     async def _check_for_prefix_async(
-        self, client: AioBaseClient, prefix: str, delimiter: str, bucket_name: str | None = None
+        self,
+        client: AioBaseClient,
+        prefix: str,
+        delimiter: str,
+        bucket_name: str | None = None,
     ) -> bool:
         return await self.check_for_prefix_async(
             client, prefix=prefix, delimiter=delimiter, bucket_name=bucket_name
@@ -643,7 +656,10 @@ class S3Hook(AwsBaseHook):
 
         paginator = client.get_paginator("list_objects_v2")
         response = paginator.paginate(
-            Bucket=bucket_name, Prefix=prefix, Delimiter=delimiter, PaginationConfig=config
+            Bucket=bucket_name,
+            Prefix=prefix,
+            Delimiter=delimiter,
+            PaginationConfig=config,
         )
 
         keys = []
@@ -655,7 +671,10 @@ class S3Hook(AwsBaseHook):
         return keys
 
     def _list_key_object_filter(
-        self, keys: list, from_datetime: datetime | None = None, to_datetime: datetime | None = None
+        self,
+        keys: list,
+        from_datetime: datetime | None = None,
+        to_datetime: datetime | None = None,
     ) -> list:
         def _is_in_period(input_date: datetime) -> bool:
             if from_datetime is not None and input_date <= from_datetime:
@@ -766,7 +785,10 @@ class S3Hook(AwsBaseHook):
                     "message": success_message,
                 }
 
-            self.log.error("FAILURE: Inactivity Period passed, not enough objects found in %s", path)
+            self.log.error(
+                "FAILURE: Inactivity Period passed, not enough objects found in %s",
+                path,
+            )
             return {
                 "status": "error",
                 "message": f"FAILURE: Inactivity Period passed, not enough objects found in {path}",
@@ -1109,7 +1131,13 @@ class S3Hook(AwsBaseHook):
             extra_args["ACL"] = acl_policy
 
         client = self.get_conn()
-        client.upload_file(filename, bucket_name, key, ExtraArgs=extra_args, Config=self.transfer_config)
+        client.upload_file(
+            filename,
+            bucket_name,
+            key,
+            ExtraArgs=extra_args,
+            Config=self.transfer_config,
+        )
         get_hook_lineage_collector().add_input_asset(
             context=self, scheme="file", asset_kwargs={"path": filename}
         )
@@ -1308,18 +1336,32 @@ class S3Hook(AwsBaseHook):
         )
 
         source_bucket_name, source_bucket_key = self.get_s3_bucket_key(
-            source_bucket_name, source_bucket_key, "source_bucket_name", "source_bucket_key"
+            source_bucket_name,
+            source_bucket_key,
+            "source_bucket_name",
+            "source_bucket_key",
         )
 
-        copy_source = {"Bucket": source_bucket_name, "Key": source_bucket_key, "VersionId": source_version_id}
+        copy_source = {
+            "Bucket": source_bucket_name,
+            "Key": source_bucket_key,
+            "VersionId": source_version_id,
+        }
         response = self.get_conn().copy_object(
-            Bucket=dest_bucket_name, Key=dest_bucket_key, CopySource=copy_source, **kwargs
+            Bucket=dest_bucket_name,
+            Key=dest_bucket_key,
+            CopySource=copy_source,
+            **kwargs,
         )
         get_hook_lineage_collector().add_input_asset(
-            context=self, scheme="s3", asset_kwargs={"bucket": source_bucket_name, "key": source_bucket_key}
+            context=self,
+            scheme="s3",
+            asset_kwargs={"bucket": source_bucket_name, "key": source_bucket_key},
         )
         get_hook_lineage_collector().add_output_asset(
-            context=self, scheme="s3", asset_kwargs={"bucket": dest_bucket_name, "key": dest_bucket_key}
+            context=self,
+            scheme="s3",
+            asset_kwargs={"bucket": dest_bucket_name, "key": dest_bucket_key},
         )
         return response
 
@@ -1435,7 +1477,10 @@ class S3Hook(AwsBaseHook):
             file_path = Path(local_dir, subdir, filename_in_s3)
 
             if file_path.is_file():
-                self.log.error("file '%s' already exists. Failing the task and not overwriting it", file_path)
+                self.log.error(
+                    "file '%s' already exists. Failing the task and not overwriting it",
+                    file_path,
+                )
                 raise FileExistsError
 
             file_path.parent.mkdir(exist_ok=True, parents=True)
@@ -1484,7 +1529,10 @@ class S3Hook(AwsBaseHook):
         s3_client = self.get_conn()
         try:
             return s3_client.generate_presigned_url(
-                ClientMethod=client_method, Params=params, ExpiresIn=expires_in, HttpMethod=http_method
+                ClientMethod=client_method,
+                Params=params,
+                ExpiresIn=expires_in,
+                HttpMethod=http_method,
             )
 
         except ClientError as e:

airflow/providers/amazon/aws/hooks/sagemaker.py

@@ -153,7 +153,9 @@ class SageMakerHook(AwsBaseHook):
     non_terminal_states = {"InProgress", "Stopping"}
     endpoint_non_terminal_states = {"Creating", "Updating", "SystemUpdating", "RollingBack", "Deleting"}
     pipeline_non_terminal_states = {"Executing", "Stopping"}
+    processing_job_non_terminal_states = {"InProgress", "Stopping"}
     failed_states = {"Failed"}
+    processing_job_failed_states = {*failed_states, "Stopped"}
     training_failed_states = {*failed_states, "Stopped"}
 
     def __init__(self, *args, **kwargs):

airflow/providers/amazon/aws/operators/redshift_data.py

@@ -28,7 +28,10 @@ from airflow.providers.amazon.aws.utils import validate_execute_complete_event
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
 
 if TYPE_CHECKING:
-    from mypy_boto3_redshift_data.type_defs import GetStatementResultResponseTypeDef
+    from mypy_boto3_redshift_data.type_defs import (
+        DescribeStatementResponseTypeDef,
+        GetStatementResultResponseTypeDef,
+    )
 
     from airflow.utils.context import Context
 
@@ -37,7 +40,7 @@ class RedshiftDataOperator(AwsBaseOperator[RedshiftDataHook]):
     """
     Executes SQL Statements against an Amazon Redshift cluster using Redshift Data.
 
-    .. seealso::
+    ... see also::
         For more information on how to use this operator, take a look at the guide:
         :ref:`howto/operator:RedshiftDataOperator`
 
@@ -84,7 +87,6 @@ class RedshiftDataOperator(AwsBaseOperator[RedshiftDataHook]):
     )
     template_ext = (".sql",)
     template_fields_renderers = {"sql": "sql"}
-    statement_id: str | None
 
     def __init__(
         self,
@@ -124,12 +126,11 @@ class RedshiftDataOperator(AwsBaseOperator[RedshiftDataHook]):
                 poll_interval,
             )
         self.return_sql_result = return_sql_result
-        self.statement_id: str | None = None
         self.deferrable = deferrable
         self.session_id = session_id
         self.session_keep_alive_seconds = session_keep_alive_seconds
 
-    def execute(self, context: Context) -> GetStatementResultResponseTypeDef | str:
+    def execute(self, context: Context) -> list[GetStatementResultResponseTypeDef] | list[str]:
         """Execute a statement against Amazon Redshift."""
         self.log.info("Executing statement: %s", self.sql)
 
@@ -154,13 +155,14 @@ class RedshiftDataOperator(AwsBaseOperator[RedshiftDataHook]):
             session_keep_alive_seconds=self.session_keep_alive_seconds,
         )
 
-        self.statement_id = query_execution_output.statement_id
+        # Pull the statement ID, session ID
+        self.statement_id: str = query_execution_output.statement_id
 
         if query_execution_output.session_id:
             self.xcom_push(context, key="session_id", value=query_execution_output.session_id)
 
         if self.deferrable and self.wait_for_completion:
-            is_finished = self.hook.check_query_is_finished(self.statement_id)
+            is_finished: bool = self.hook.check_query_is_finished(self.statement_id)
             if not is_finished:
                 self.defer(
                     timeout=self.execution_timeout,
@@ -176,16 +178,13 @@ class RedshiftDataOperator(AwsBaseOperator[RedshiftDataHook]):
                     method_name="execute_complete",
                 )
 
-        if self.return_sql_result:
-            result = self.hook.conn.get_statement_result(Id=self.statement_id)
-            self.log.debug("Statement result: %s", result)
-            return result
-        else:
-            return self.statement_id
+        # Use the get_sql_results method to return the results of the SQL query, or the statement_ids,
+        # depending on the value of self.return_sql_result
+        return self.get_sql_results(statement_id=self.statement_id, return_sql_result=self.return_sql_result)
 
     def execute_complete(
         self, context: Context, event: dict[str, Any] | None = None
-    ) -> GetStatementResultResponseTypeDef | str:
+    ) -> list[GetStatementResultResponseTypeDef] | list[str]:
         event = validate_execute_complete_event(event)
 
         if event["status"] == "error":
@@ -197,16 +196,40 @@ class RedshiftDataOperator(AwsBaseOperator[RedshiftDataHook]):
             raise AirflowException("statement_id should not be empty.")
 
         self.log.info("%s completed successfully.", self.task_id)
-        if self.return_sql_result:
-            result = self.hook.conn.get_statement_result(Id=statement_id)
-            self.log.debug("Statement result: %s", result)
-            return result
 
-        return statement_id
+        # Use the get_sql_results method to return the results of the SQL query, or the statement_ids,
+        # depending on the value of self.return_sql_result
+        return self.get_sql_results(statement_id=statement_id, return_sql_result=self.return_sql_result)
+
+    def get_sql_results(
+        self, statement_id: str, return_sql_result: bool
+    ) -> list[GetStatementResultResponseTypeDef] | list[str]:
+        """
+        Retrieve either the result of the SQL query, or the statement ID(s).
+
+        :param statement_id: Statement ID of the running queries
+        :param return_sql_result: Boolean, true if results should be returned
+        """
+        # ISSUE-40427: Pull the statement, and check to see if there are sub-statements. If that is the
+        # case, pull each of the sub-statement ID's, and grab the results. Otherwise, just use statement_id
+        statement: DescribeStatementResponseTypeDef = self.hook.conn.describe_statement(Id=statement_id)
+        statement_ids: list[str] = (
+            [sub_statement["Id"] for sub_statement in statement["SubStatements"]]
+            if len(statement.get("SubStatements", [])) > 0
+            else [statement_id]
+        )
+
+        # If returning the SQL result, use get_statement_result to return the records for each query
+        if return_sql_result:
+            results: list = [self.hook.conn.get_statement_result(Id=sid) for sid in statement_ids]
+            self.log.debug("Statement result(s): %s", results)
+            return results
+        else:
+            return statement_ids
 
     def on_kill(self) -> None:
         """Cancel the submitted redshift query."""
-        if self.statement_id:
+        if hasattr(self, "statement_id"):
             self.log.info("Received a kill signal.")
             self.log.info("Stopping Query with statementId - %s", self.statement_id)
 

airflow/providers/amazon/aws/sensors/sagemaker.py

@@ -330,3 +330,35 @@ class SageMakerAutoMLSensor(SageMakerBaseSensor):
 
     def state_from_response(self, response: dict) -> str:
         return response["AutoMLJobStatus"]
+
+
+class SageMakerProcessingSensor(SageMakerBaseSensor):
+    """
+    Poll the processing job until it reaches a terminal state; raise AirflowException with the failure reason.
+
+    .. seealso::
+        For more information on how to use this sensor, take a look at the guide:
+        :ref:`howto/sensor:SageMakerProcessingSensor`
+
+    :param job_name: Name of the processing job to watch.
+    """
+
+    template_fields: Sequence[str] = ("job_name",)
+    template_ext: Sequence[str] = ()
+
+    def __init__(self, *, job_name: str, **kwargs):
+        super().__init__(**kwargs)
+        self.job_name = job_name
+
+    def non_terminal_states(self) -> set[str]:
+        return SageMakerHook.processing_job_non_terminal_states
+
+    def failed_states(self) -> set[str]:
+        return SageMakerHook.processing_job_failed_states
+
+    def get_sagemaker_response(self) -> dict:
+        self.log.info("Poking Sagemaker ProcessingJob %s", self.job_name)
+        return self.hook.describe_processing_job(self.job_name)
+
+    def state_from_response(self, response: dict) -> str:
+        return response["ProcessingJobStatus"]

airflow/providers/amazon/aws/transfers/redshift_to_s3.py

@@ -152,6 +152,10 @@ class RedshiftToS3Operator(BaseOperator):
             table = self.table
         return f"SELECT * FROM {table}"
 
+    @property
+    def use_redshift_data(self):
+        return bool(self.redshift_data_api_kwargs)
+
     def execute(self, context: Context) -> None:
         if self.table and self.table_as_file_name:
             self.s3_key = f"{self.s3_key}/{self.table}_"
@@ -164,14 +168,13 @@ class RedshiftToS3Operator(BaseOperator):
         if self.include_header and "HEADER" not in [uo.upper().strip() for uo in self.unload_options]:
             self.unload_options = [*self.unload_options, "HEADER"]
 
-        redshift_hook: RedshiftDataHook | RedshiftSQLHook
-        if self.redshift_data_api_kwargs:
-            redshift_hook = RedshiftDataHook(aws_conn_id=self.redshift_conn_id)
+        if self.use_redshift_data:
+            redshift_data_hook = RedshiftDataHook(aws_conn_id=self.redshift_conn_id)
             for arg in ["sql", "parameters"]:
                 if arg in self.redshift_data_api_kwargs:
                     raise AirflowException(f"Cannot include param '{arg}' in Redshift Data API kwargs")
         else:
-            redshift_hook = RedshiftSQLHook(redshift_conn_id=self.redshift_conn_id)
+            redshift_sql_hook = RedshiftSQLHook(redshift_conn_id=self.redshift_conn_id)
         conn = S3Hook.get_connection(conn_id=self.aws_conn_id) if self.aws_conn_id else None
         if conn and conn.extra_dejson.get("role_arn", False):
             credentials_block = f"aws_iam_role={conn.extra_dejson['role_arn']}"
@@ -187,10 +190,106 @@ class RedshiftToS3Operator(BaseOperator):
         )
 
         self.log.info("Executing UNLOAD command...")
-        if isinstance(redshift_hook, RedshiftDataHook):
-            redshift_hook.execute_query(
+        if self.use_redshift_data:
+            redshift_data_hook.execute_query(
                 sql=unload_query, parameters=self.parameters, **self.redshift_data_api_kwargs
             )
         else:
-            redshift_hook.run(unload_query, self.autocommit, parameters=self.parameters)
+            redshift_sql_hook.run(unload_query, self.autocommit, parameters=self.parameters)
         self.log.info("UNLOAD command complete...")
+
+    def get_openlineage_facets_on_complete(self, task_instance):
+        """Implement on_complete as we may query for table details."""
+        from airflow.providers.amazon.aws.utils.openlineage import (
+            get_facets_from_redshift_table,
+            get_identity_column_lineage_facet,
+        )
+        from airflow.providers.common.compat.openlineage.facet import (
+            Dataset,
+            Error,
+            ExtractionErrorRunFacet,
+        )
+        from airflow.providers.openlineage.extractors import OperatorLineage
+
+        output_dataset = Dataset(
+            namespace=f"s3://{self.s3_bucket}",
+            name=self.s3_key,
+        )
+
+        if self.use_redshift_data:
+            redshift_data_hook = RedshiftDataHook(aws_conn_id=self.redshift_conn_id)
+            database = self.redshift_data_api_kwargs.get("database")
+            identifier = self.redshift_data_api_kwargs.get(
+                "cluster_identifier", self.redshift_data_api_kwargs.get("workgroup_name")
+            )
+            port = self.redshift_data_api_kwargs.get("port", "5439")
+            authority = f"{identifier}.{redshift_data_hook.region_name}:{port}"
+        else:
+            redshift_sql_hook = RedshiftSQLHook(redshift_conn_id=self.redshift_conn_id)
+            database = redshift_sql_hook.conn.schema
+            authority = redshift_sql_hook.get_openlineage_database_info(redshift_sql_hook.conn).authority
+
+        if self.select_query == self.default_select_query:
+            if self.use_redshift_data:
+                input_dataset_facets = get_facets_from_redshift_table(
+                    redshift_data_hook, self.table, self.redshift_data_api_kwargs, self.schema
+                )
+            else:
+                input_dataset_facets = get_facets_from_redshift_table(
+                    redshift_sql_hook, self.table, {}, self.schema
+                )
+
+            input_dataset = Dataset(
+                namespace=f"redshift://{authority}",
+                name=f"{database}.{self.schema}.{self.table}" if database else f"{self.schema}.{self.table}",
+                facets=input_dataset_facets,
+            )
+
+            # If default select query is used (SELECT *) output file matches the input table.
+            output_dataset.facets = {
+                "schema": input_dataset_facets["schema"],
+                "columnLineage": get_identity_column_lineage_facet(
+                    field_names=[field.name for field in input_dataset_facets["schema"].fields],
+                    input_datasets=[input_dataset],
+                ),
+            }
+
+            return OperatorLineage(inputs=[input_dataset], outputs=[output_dataset])
+
+        try:
+            from airflow.providers.openlineage.sqlparser import SQLParser, from_table_meta
+        except ImportError:
+            return OperatorLineage(outputs=[output_dataset])
+
+        run_facets = {}
+        parse_result = SQLParser(dialect="redshift", default_schema=self.schema).parse(self.select_query)
+        if parse_result.errors:
+            run_facets["extractionError"] = ExtractionErrorRunFacet(
+                totalTasks=1,
+                failedTasks=1,
+                errors=[
+                    Error(
+                        errorMessage=error.message,
+                        stackTrace=None,
+                        task=error.origin_statement,
+                        taskNumber=error.index,
+                    )
+                    for error in parse_result.errors
+                ],
+            )
+
+        input_datasets = []
+        for in_tb in parse_result.in_tables:
+            ds = from_table_meta(in_tb, database, f"redshift://{authority}", False)
+            schema, table = ds.name.split(".")[-2:]
+            if self.use_redshift_data:
+                input_dataset_facets = get_facets_from_redshift_table(
+                    redshift_data_hook, table, self.redshift_data_api_kwargs, schema
+                )
+            else:
+                input_dataset_facets = get_facets_from_redshift_table(redshift_sql_hook, table, {}, schema)
+
+            ds.facets = input_dataset_facets
+            input_datasets.append(ds)
+
+        return OperatorLineage(inputs=input_datasets, outputs=[output_dataset], run_facets=run_facets)

airflow/providers/amazon/aws/transfers/s3_to_redshift.py

@@ -241,7 +241,7 @@ class S3ToRedshiftOperator(BaseOperator):
 
         output_dataset = Dataset(
             namespace=f"redshift://{authority}",
-            name=f"{database}.{self.schema}.{self.table}",
+            name=f"{database}.{self.schema}.{self.table}" if database else f"{self.schema}.{self.table}",
             facets=output_dataset_facets,
         )
 

airflow/providers/amazon/aws/triggers/ecs.py

@@ -18,7 +18,8 @@
 from __future__ import annotations
 
 import asyncio
-from typing import TYPE_CHECKING, Any, AsyncIterator
+from collections.abc import AsyncIterator
+from typing import TYPE_CHECKING, Any
 
 from botocore.exceptions import ClientError, WaiterError
 
@@ -165,11 +166,10 @@ class TaskDoneTrigger(BaseTrigger):
         )
 
     async def run(self) -> AsyncIterator[TriggerEvent]:
-        async with EcsHook(
-            aws_conn_id=self.aws_conn_id, region_name=self.region
-        ).async_conn as ecs_client, AwsLogsHook(
-            aws_conn_id=self.aws_conn_id, region_name=self.region
-        ).async_conn as logs_client:
+        async with (
+            EcsHook(aws_conn_id=self.aws_conn_id, region_name=self.region).async_conn as ecs_client,
+            AwsLogsHook(aws_conn_id=self.aws_conn_id, region_name=self.region).async_conn as logs_client,
+        ):
             waiter = ecs_client.get_waiter("tasks_stopped")
             logs_token = None
             while self.waiter_max_attempts:

airflow/providers/amazon/get_provider_info.py

@@ -28,8 +28,9 @@ def get_provider_info():
         "name": "Amazon",
         "description": "Amazon integration (including `Amazon Web Services (AWS) <https://aws.amazon.com/>`__).\n",
         "state": "ready",
-        "source-date-epoch": 1728484310,
+        "source-date-epoch": 1730011042,
         "versions": [
+            "9.1.0",
             "9.0.0",
             "8.29.0",
             "8.28.0",
@@ -98,7 +99,7 @@ def get_provider_info():
         ],
         "dependencies": [
             "apache-airflow>=2.8.0",
-            "apache-airflow-providers-common-compat>=1.1.0",
+            "apache-airflow-providers-common-compat>=1.2.1",
             "apache-airflow-providers-common-sql>=1.3.1",
             "apache-airflow-providers-http",
             "boto3>=1.34.90",
@@ -107,7 +108,6 @@ def get_provider_info():
             "watchtower>=3.0.0,!=3.3.0,<4",
             "jsonpath_ng>=1.5.3",
             "redshift_connector>=2.0.918",
-            "sqlalchemy_redshift>=0.8.6",
             "asgiref>=2.3.0",
             "PyAthena>=3.0.10",
             "jmespath>=0.7.0",

{apache_airflow_providers_amazon-9.0.0rc1.dist-info → apache_airflow_providers_amazon-9.1.0rc1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: apache-airflow-providers-amazon
-Version: 9.0.0rc1
+Version: 9.1.0rc1
 Summary: Provider package apache-airflow-providers-amazon for Apache Airflow
 Keywords: airflow-provider,amazon,airflow,integration
 Author-email: Apache Software Foundation <dev@airflow.apache.org>
@@ -21,7 +21,7 @@ Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Topic :: System :: Monitoring
 Requires-Dist: PyAthena>=3.0.10
-Requires-Dist: apache-airflow-providers-common-compat>=1.1.0rc0
+Requires-Dist: apache-airflow-providers-common-compat>=1.2.1rc0
 Requires-Dist: apache-airflow-providers-common-sql>=1.3.1rc0
 Requires-Dist: apache-airflow-providers-http
 Requires-Dist: apache-airflow>=2.8.0rc0
@@ -33,7 +33,6 @@ Requires-Dist: jmespath>=0.7.0
 Requires-Dist: jsonpath_ng>=1.5.3
 Requires-Dist: python3-saml>=1.16.0
 Requires-Dist: redshift_connector>=2.0.918
-Requires-Dist: sqlalchemy_redshift>=0.8.6
 Requires-Dist: watchtower>=3.0.0,!=3.3.0,<4
 Requires-Dist: aiobotocore[boto3]>=2.13.0 ; extra == "aiobotocore"
 Requires-Dist: apache-airflow-providers-apache-hive ; extra == "apache.hive"
@@ -55,8 +54,8 @@ Requires-Dist: s3fs>=2023.10.0 ; extra == "s3fs"
 Requires-Dist: apache-airflow-providers-salesforce ; extra == "salesforce"
 Requires-Dist: apache-airflow-providers-ssh ; extra == "ssh"
 Project-URL: Bug Tracker, https://github.com/apache/airflow/issues
-Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-amazon/9.0.0/changelog.html
-Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-amazon/9.0.0
+Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-amazon/9.1.0/changelog.html
+Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-amazon/9.1.0
 Project-URL: Slack Chat, https://s.apache.org/airflow-slack
 Project-URL: Source Code, https://github.com/apache/airflow
 Project-URL: Twitter, https://twitter.com/ApacheAirflow
@@ -124,7 +123,7 @@ Provides-Extra: ssh
 
 Package ``apache-airflow-providers-amazon``
 
-Release: ``9.0.0.rc1``
+Release: ``9.1.0.rc1``
 
 
 Amazon integration (including `Amazon Web Services (AWS) <https://aws.amazon.com/>`__).
@@ -137,7 +136,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/9.0.0/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-amazon/9.1.0/>`_.
 
 Installation
 ------------
@@ -155,7 +154,7 @@ Requirements
 PIP package                                 Version required
 ==========================================  ======================
 ``apache-airflow``                          ``>=2.8.0``
-``apache-airflow-providers-common-compat``  ``>=1.1.0``
+``apache-airflow-providers-common-compat``  ``>=1.2.1``
 ``apache-airflow-providers-common-sql``     ``>=1.3.1``
 ``apache-airflow-providers-http``
 ``boto3``                                   ``>=1.34.90``
@@ -164,7 +163,6 @@ PIP package                                 Version required
 ``watchtower``                              ``>=3.0.0,!=3.3.0,<4``
 ``jsonpath_ng``                             ``>=1.5.3``
 ``redshift_connector``                      ``>=2.0.918``
-``sqlalchemy_redshift``                     ``>=0.8.6``
 ``asgiref``                                 ``>=2.3.0``
 ``PyAthena``                                ``>=3.0.10``
 ``jmespath``                                ``>=0.7.0``
@@ -204,4 +202,4 @@ Dependent package
 ======================================================================================================================  ===================
 
 The changelog for the provider package can be found in the
-`changelog <https://airflow.apache.org/docs/apache-airflow-providers-amazon/9.0.0/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-amazon/9.1.0/changelog.html>`_.

{apache_airflow_providers_amazon-9.0.0rc1.dist-info → apache_airflow_providers_amazon-9.1.0rc1.dist-info}/RECORD

@@ -1,10 +1,10 @@
 airflow/providers/amazon/LICENSE,sha256=FFb4jd2AXnOOf7XLP04pQW6jbdhG49TxlGY6fFpCV1Y,13609
-airflow/providers/amazon/__init__.py,sha256=Ga3sigSdEQEWuktzplPZ3O2thA2zlCfcNpooQtGfAUg,1493
-airflow/providers/amazon/get_provider_info.py,sha256=aHxj3H6b5zvUNgHHh1YJRzh1aey-Y9hNbd5eymWhuNU,68981
+airflow/providers/amazon/__init__.py,sha256=MuTYiz17zqZxTSbZ537dOIWjnfLeKguUmhr0_326bro,1493
+airflow/providers/amazon/get_provider_info.py,sha256=JnyQK8t0-7kcB2ed5tBqkxh0PtzYoFpe1yYseO_dU6I,68960
 airflow/providers/amazon/aws/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
 airflow/providers/amazon/aws/exceptions.py,sha256=uRGNMgXvgdzfphpOTiyj74lQhjzb70J-X8n6fsx5Jog,1864
 airflow/providers/amazon/aws/assets/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
-airflow/providers/amazon/aws/assets/s3.py,sha256=2wSSiuba_yMwcoQBNPruSNjlLsIrdZeRrkrgmPvqngY,1917
+airflow/providers/amazon/aws/assets/s3.py,sha256=wNaJiOM90-SCauD4EQneZVXMO54yDRjLPfI8D5o0-fw,1861
 airflow/providers/amazon/aws/auth_manager/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
 airflow/providers/amazon/aws/auth_manager/aws_auth_manager.py,sha256=0E38_x_EjV6uwMqkqe0vfmLGGCpcQKEAvnr7FPQwXTY,16625
 airflow/providers/amazon/aws/auth_manager/constants.py,sha256=Jdluo42InhyNGkYHB_dRtoFMpKanJLJdH0hyR9-5AZg,1050
@@ -38,7 +38,7 @@ airflow/providers/amazon/aws/executors/utils/exponential_backoff_retry.py,sha256
 airflow/providers/amazon/aws/fs/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
 airflow/providers/amazon/aws/fs/s3.py,sha256=Ty9XT9c1XArkUYcQkalvNZhuoTlEg3uKy-AIzNW9LgY,4797
 airflow/providers/amazon/aws/hooks/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
-airflow/providers/amazon/aws/hooks/appflow.py,sha256=-le6RsIMWIqTav7KGknsph9Td42znSm_eIYztxc_RsE,5263
+airflow/providers/amazon/aws/hooks/appflow.py,sha256=v7o6BgFDm8McE3JtB0oMkc80_nTP0e_u_uEDWFa0sVI,5367
 airflow/providers/amazon/aws/hooks/athena.py,sha256=tnck2Ts9QFVDU5MsY4_9mTULyMO9QHyMm_KTJPtCsk0,13001
 airflow/providers/amazon/aws/hooks/athena_sql.py,sha256=vFIUbMMTem3xvYAUTvW3h1ypjpKVLNck3VbrAlupVLA,6844
 airflow/providers/amazon/aws/hooks/base_aws.py,sha256=LDxXMKjlYFEhsZ0u0DsIjf7qt-Wr6aRn9NpfoNJc0nc,43630
@@ -74,9 +74,9 @@ airflow/providers/amazon/aws/hooks/quicksight.py,sha256=2Am_K-BcoqcfuWwLbWjW1Lsb
 airflow/providers/amazon/aws/hooks/rds.py,sha256=h7NF3GZ42RKeh70rlg2BQFVpa8vNadS37slj0MsAT3w,15211
 airflow/providers/amazon/aws/hooks/redshift_cluster.py,sha256=ywJxbcOy91-oGGkApo6_nRKucyaEhxXqEXikG7q9uZ4,7977
 airflow/providers/amazon/aws/hooks/redshift_data.py,sha256=5Kz7tsTbDwDDAqTud1--vyi74IksfHWBouIPRXYYFJk,11812
-airflow/providers/amazon/aws/hooks/redshift_sql.py,sha256=7Iqr_5IVREdsW9mdTZEm2_Ng13eIBvkqZKS-LtTMM-k,11318
-airflow/providers/amazon/aws/hooks/s3.py,sha256=fEAdQboiKsG9IYTSAAZq_ivWmFsNlYfa7Q4HvEgbSgc,61114
-airflow/providers/amazon/aws/hooks/sagemaker.py,sha256=wLmaKnuDv80f59oqGgLqG87ddwH_AAN8U386k2NCgng,60330
+airflow/providers/amazon/aws/hooks/redshift_sql.py,sha256=UTH2cyCfdGp5o2l94le5HLCF82KpsoGgRAgd7jrmeVw,11301
+airflow/providers/amazon/aws/hooks/s3.py,sha256=ChlYVLix_fQcZMl1-rpOx7ZxOUaxFGfGxhMPAiA-_lQ,61506
+airflow/providers/amazon/aws/hooks/sagemaker.py,sha256=w3_zqbuNznP47TBLioHuBenNv82blDul-rppyZMRaKg,60461
 airflow/providers/amazon/aws/hooks/secrets_manager.py,sha256=6srh3jUeSGoqyrSj1M6aSOaA9xT5kna0VGUC0kzH-q0,2690
 airflow/providers/amazon/aws/hooks/ses.py,sha256=uOTjyhb87jNyf2B11zH1wg5Oomnsx0nM4aHteP-mCHs,4147
 airflow/providers/amazon/aws/hooks/sns.py,sha256=VNbhqXZv-tw9IqWhgrAsJPgT9mm9LJ4L3EJIfp3c97w,3436
@@ -125,7 +125,7 @@ airflow/providers/amazon/aws/operators/neptune.py,sha256=on5oNX5K4yHfW1POE0eeZuj
 airflow/providers/amazon/aws/operators/quicksight.py,sha256=jc3Eof19UfLt5IqbQswRzaHaK8h0ACLY99i_1Prtq10,4089
 airflow/providers/amazon/aws/operators/rds.py,sha256=U2YLPx5MZCdDrLIyy-9K93W5aUtQ2VHUm6Tl0QMDomo,38396
 airflow/providers/amazon/aws/operators/redshift_cluster.py,sha256=rmBHCssxrYEJ8EnENY-AnzC004lbtHvxXHpy69sHtV0,36681
-airflow/providers/amazon/aws/operators/redshift_data.py,sha256=X93ZcyO3YgaGJ-xyMJzbBxuRE5xfWtG-II-PwIu87ng,9392
+airflow/providers/amazon/aws/operators/redshift_data.py,sha256=36MVojiezDyGZ_4aQuY8xvs9doQlz_SWpJEp6Kwkw0U,10832
 airflow/providers/amazon/aws/operators/s3.py,sha256=d_K2DDNXEXkoi-WZ02-bwCf244Ogiw1PBaHcbsX-8Sg,36272
 airflow/providers/amazon/aws/operators/sagemaker.py,sha256=nyGS6uLP3eUYPCwOXDhdlucSGvI2lrSV8PUJ1_1f_5w,82337
 airflow/providers/amazon/aws/operators/sns.py,sha256=Rttd015UhLo4pCplGybxtLhflyu_26IFzYP7WTmQFk8,3730
@@ -158,7 +158,7 @@ airflow/providers/amazon/aws/sensors/quicksight.py,sha256=ow5CqGVcqH5h93EuJe71WN
 airflow/providers/amazon/aws/sensors/rds.py,sha256=AB2dH7fLwAaQogj0NYRrOOftfeOk_INetsyVHr1_qfM,6476
 airflow/providers/amazon/aws/sensors/redshift_cluster.py,sha256=m2z8vU_OfJds-hLecaebrtO6DIsseqJ-MHr7NfoJm0k,4021
 airflow/providers/amazon/aws/sensors/s3.py,sha256=Fp8LnFGbAaXV_JOgvIyN_KLs5uNiV86j-YL2_oG3GMI,17306
-airflow/providers/amazon/aws/sensors/sagemaker.py,sha256=IxxmJvK1wl--yGtTAPmFBIOhTdakGv6x6C2Jt2Ehp7g,12525
+airflow/providers/amazon/aws/sensors/sagemaker.py,sha256=YfXSdpTVfIDW0Wrbe13wmnuyc4xiWcTW6EQ7KedscV4,13634
 airflow/providers/amazon/aws/sensors/sqs.py,sha256=m3uPnuLZm8XB8_wd6cGmzKnaRqi67ZSJFeonEWCRDmw,10569
 airflow/providers/amazon/aws/sensors/step_function.py,sha256=RqyiAelyPearI6wvkauDwfhGcAISNXvGVzwAzF1C02E,3569
 airflow/providers/amazon/aws/transfers/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
@@ -175,10 +175,10 @@ airflow/providers/amazon/aws/transfers/http_to_s3.py,sha256=J1HjIdGsd4Zl8kk-RJAX
 airflow/providers/amazon/aws/transfers/imap_attachment_to_s3.py,sha256=xYJ94xNDsadluJpUUv1fURLW7YhSgL-9GaebZ6l4RLU,4536
 airflow/providers/amazon/aws/transfers/local_to_s3.py,sha256=yp9m7aZuL6YgzYRsFcyZ1wcGTXZTMO0F0CuBfkH1eGo,4165
 airflow/providers/amazon/aws/transfers/mongo_to_s3.py,sha256=OU7Cge_0WQd7xEb38V-0hjSHbjZRCQ7Ay4xntcG9R28,6020
-airflow/providers/amazon/aws/transfers/redshift_to_s3.py,sha256=Ui25l0vxvb5fymN3-sX4NBoQVXDz5MkutJ34n5ZZ4-k,8718
+airflow/providers/amazon/aws/transfers/redshift_to_s3.py,sha256=-XMmYoHmQiZqZ5rqdYp4ZZbUNUe-Vu1z8TrkH9pxNHA,13001
 airflow/providers/amazon/aws/transfers/s3_to_dynamodb.py,sha256=BRYID2nCmjooVH4WVD7tcy5MnkGQPK4c9BwioeCsIsw,11658
 airflow/providers/amazon/aws/transfers/s3_to_ftp.py,sha256=cxyNRW_FJQNNluuYr5fVluGLYnNRUvN75iHSSEHrVnY,2966
-airflow/providers/amazon/aws/transfers/s3_to_redshift.py,sha256=J3udUq63RdNFlavMYitvcYrWSKpCPNyBbKe1ky0_rJo,11162
+airflow/providers/amazon/aws/transfers/s3_to_redshift.py,sha256=GXMJV-_7Vy0CvPBOgSGWw_L2xf25XqRF9ztDFdMpWLc,11209
 airflow/providers/amazon/aws/transfers/s3_to_sftp.py,sha256=bgHgKv7o8ueC_zkhzW5k2xZpFnMlBHMcDf0t4sQ7kHY,3488
 airflow/providers/amazon/aws/transfers/s3_to_sql.py,sha256=kUuHDDR2ATDBHoRVq_3DOXGe8MVH9gcLHDMLMtFe4GI,4949
 airflow/providers/amazon/aws/transfers/salesforce_to_s3.py,sha256=yt77guCyYqVwDLdpmLb_psAI8PpGvbCjKQO6pz1J564,5686
@@ -192,7 +192,7 @@ airflow/providers/amazon/aws/triggers/batch.py,sha256=GogZnPaSc1ms55_aNcSDtV4wIZ
 airflow/providers/amazon/aws/triggers/bedrock.py,sha256=IiKyl0UUax-ex4siLjZpQGDZWyAOOhvJ-9USyRi_r3c,7260
 airflow/providers/amazon/aws/triggers/comprehend.py,sha256=atK02t-G6e-Rgd-a-IHc4n-wGZ3oC4pKueOwNeaLCrI,4063
 airflow/providers/amazon/aws/triggers/ec2.py,sha256=gMY3EP4TmL6SodLw12FNSLttlHd7hRhOu-q3CiG7y2w,3245
-airflow/providers/amazon/aws/triggers/ecs.py,sha256=NFgFR9rkj5wiLN4X9t5itHf4ayOD7eoiWbnG7W13N8A,9202
+airflow/providers/amazon/aws/triggers/ecs.py,sha256=YXFXUpQ1ejvtMIwkiet4LTLdJSkG8nr_ZEUI5lpqRGA,9222
 airflow/providers/amazon/aws/triggers/eks.py,sha256=9elEgIErRQRndk4sgPZ2F9bjcaWsUoUG18Qv758ly1U,16151
 airflow/providers/amazon/aws/triggers/emr.py,sha256=og83L9BPUcqqVXHwfUTD9cA-276YDo3Fnc0e2svQfqE,16566
 airflow/providers/amazon/aws/triggers/glue.py,sha256=hv_nLzBRPG13MetjEfU_-KuTphLE-xyF6yW4uQJQuBc,9480
@@ -210,7 +210,6 @@ airflow/providers/amazon/aws/triggers/sagemaker.py,sha256=0h_JlleQwUYq0JY2QST681
 airflow/providers/amazon/aws/triggers/sqs.py,sha256=tVA1i8XzV5AqbVQAdWrgrLKoZF8ewqgEwV7ggk1hrQM,8257
 airflow/providers/amazon/aws/triggers/step_function.py,sha256=M1HGdrnxL_T9KSCBNy2t531xMNJaFc-Y792T9cSmLGM,2685
 airflow/providers/amazon/aws/utils/__init__.py,sha256=yUkoHb2LuqSyHvj-HAhc2r2s04Kv_PhdyLMq52KarO8,3878
-airflow/providers/amazon/aws/utils/asset_compat_lineage_collector.py,sha256=7W1T15iz-h5cf0wpW-oMroeIF882PGwqJck3elRJOIs,3606
 airflow/providers/amazon/aws/utils/connection_wrapper.py,sha256=sPS-h_DK9SL2koHTaYwGoaC13Sr6NAgQXRLmvkvgXZY,16122
 airflow/providers/amazon/aws/utils/eks_get_token.py,sha256=q4utFF2c02T2Lm6KIZLABOiXJeglVZKCOxq6gn14dsk,2342
 airflow/providers/amazon/aws/utils/emailer.py,sha256=y-bzg1BZzOQ8J9-ed-74LY3VMv6LrLfBDtw5S4t3Tv4,1855
@@ -250,7 +249,7 @@ airflow/providers/amazon/aws/waiters/rds.json,sha256=HNmNQm5J-VaFHzjWb1pE5P7-Ix-
 airflow/providers/amazon/aws/waiters/redshift.json,sha256=jOBotCgbkko1b_CHcGEbhhRvusgt0YSzVuFiZrqVP30,1742
 airflow/providers/amazon/aws/waiters/sagemaker.json,sha256=JPHuQtUFZ1B7EMLfVmCRevNZ9jgpB71LM0dva8ZEO9A,5254
 airflow/providers/amazon/aws/waiters/stepfunctions.json,sha256=GsOH-emGerKGBAUFmI5lpMfNGH4c0ol_PSiea25DCEY,1033
-apache_airflow_providers_amazon-9.0.0rc1.dist-info/entry_points.txt,sha256=vlc0ZzhBkMrav1maTRofgksnAw4SwoQLFX9cmnTgktk,102
-apache_airflow_providers_amazon-9.0.0rc1.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
-apache_airflow_providers_amazon-9.0.0rc1.dist-info/METADATA,sha256=sfyXfGDEJUD0jdrC2qVOvfjLPl_vvKM3cBFbmVYk09g,10785
-apache_airflow_providers_amazon-9.0.0rc1.dist-info/RECORD,,
+apache_airflow_providers_amazon-9.1.0rc1.dist-info/entry_points.txt,sha256=vlc0ZzhBkMrav1maTRofgksnAw4SwoQLFX9cmnTgktk,102
+apache_airflow_providers_amazon-9.1.0rc1.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
+apache_airflow_providers_amazon-9.1.0rc1.dist-info/METADATA,sha256=qBjKdJd9WgtPT5L-V5k2D8IBorkkQxcQwRTKdnEXuEg,10687
+apache_airflow_providers_amazon-9.1.0rc1.dist-info/RECORD,,

airflow/providers/amazon/aws/utils/asset_compat_lineage_collector.py

@@ -1,106 +0,0 @@
-# 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
-
-from importlib.util import find_spec
-
-
-def _get_asset_compat_hook_lineage_collector():
-    from airflow.lineage.hook import get_hook_lineage_collector
-
-    collector = get_hook_lineage_collector()
-
-    if all(
-        getattr(collector, asset_method_name, None)
-        for asset_method_name in ("add_input_asset", "add_output_asset", "collected_assets")
-    ):
-        return collector
-
-    # dataset is renamed as asset in Airflow 3.0
-
-    from functools import wraps
-
-    from airflow.lineage.hook import DatasetLineageInfo, HookLineage
-
-    DatasetLineageInfo.asset = DatasetLineageInfo.dataset
-
-    def rename_dataset_kwargs_as_assets_kwargs(function):
-        @wraps(function)
-        def wrapper(*args, **kwargs):
-            if "asset_kwargs" in kwargs:
-                kwargs["dataset_kwargs"] = kwargs.pop("asset_kwargs")
-
-            if "asset_extra" in kwargs:
-                kwargs["dataset_extra"] = kwargs.pop("asset_extra")
-
-            return function(*args, **kwargs)
-
-        return wrapper
-
-    collector.create_asset = rename_dataset_kwargs_as_assets_kwargs(collector.create_dataset)
-    collector.add_input_asset = rename_dataset_kwargs_as_assets_kwargs(collector.add_input_dataset)
-    collector.add_output_asset = rename_dataset_kwargs_as_assets_kwargs(collector.add_output_dataset)
-
-    def collected_assets_compat(collector) -> HookLineage:
-        """Get the collected hook lineage information."""
-        lineage = collector.collected_datasets
-        return HookLineage(
-            [
-                DatasetLineageInfo(dataset=item.dataset, count=item.count, context=item.context)
-                for item in lineage.inputs
-            ],
-            [
-                DatasetLineageInfo(dataset=item.dataset, count=item.count, context=item.context)
-                for item in lineage.outputs
-            ],
-        )
-
-    setattr(
-        collector.__class__,
-        "collected_assets",
-        property(lambda collector: collected_assets_compat(collector)),
-    )
-
-    return collector
-
-
-def get_hook_lineage_collector():
-    # HookLineageCollector added in 2.10
-    try:
-        if find_spec("airflow.assets"):
-            # Dataset has been renamed as Asset in 3.0
-            from airflow.lineage.hook import get_hook_lineage_collector
-
-            return get_hook_lineage_collector()
-
-        return _get_asset_compat_hook_lineage_collector()
-    except ImportError:
-
-        class NoOpCollector:
-            """
-            NoOpCollector is a hook lineage collector that does nothing.
-
-            It is used when you want to disable lineage collection.
-            """
-
-            def add_input_asset(self, *_, **__):
-                pass
-
-            def add_output_asset(self, *_, **__):
-                pass
-
-        return NoOpCollector()