# 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 typing import TYPE_CHECKING, Sequence
from airflow.exceptions import AirflowException
from airflow.providers.amazon.aws.hooks.eventbridge import EventBridgeHook
from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
from airflow.utils.helpers import prune_dict
if TYPE_CHECKING:
from airflow.utils.context import Context
[docs]class EventBridgePutEventsOperator(AwsBaseOperator[EventBridgeHook]):
"""
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 Airflow connection used for AWS credentials.
If this is ``None`` or empty then the default boto3 behaviour is used. If
running Airflow in a distributed manner and aws_conn_id is None or
empty, then default boto3 configuration would be used (and must be
maintained on each worker node).
:param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
:param verify: Whether or not to verify SSL certificates. See:
https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html
:param botocore_config: Configuration dictionary (key-values) for botocore client. See:
https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.htmlt
"""
[docs] aws_hook_class = EventBridgeHook
[docs] template_fields: Sequence[str] = aws_template_fields("entries", "endpoint_id")
def __init__(self, *, entries: list[dict], endpoint_id: str | None = None, **kwargs):
super().__init__(**kwargs)
self.entries = entries
self.endpoint_id = endpoint_id
[docs] def execute(self, context: Context):
response = self.hook.conn.put_events(
**prune_dict(
{
"Entries": self.entries,
"EndpointId": self.endpoint_id,
}
)
)
self.log.info("Sent %d events to EventBridge.", len(self.entries))
if response.get("FailedEntryCount"):
for event in response["Entries"]:
if "ErrorCode" in event:
self.log.error(event)
raise AirflowException(
f"{response['FailedEntryCount']} entries in this request have failed to send."
)
if self.do_xcom_push:
return [e["EventId"] for e in response["Entries"]]
[docs]class EventBridgePutRuleOperator(AwsBaseOperator[EventBridgeHook]):
"""
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 aws_conn_id: The Airflow connection used for AWS credentials.
If this is ``None`` or empty then the default boto3 behaviour is used. If
running Airflow in a distributed manner and aws_conn_id is None or
empty, then default boto3 configuration would be used (and must be
maintained on each worker node).
:param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
:param verify: Whether or not to verify SSL certificates. See:
https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html
:param botocore_config: Configuration dictionary (key-values) for botocore client. See:
https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.htmlt
"""
[docs] aws_hook_class = EventBridgeHook
[docs] template_fields: Sequence[str] = aws_template_fields(
"name",
"description",
"event_bus_name",
"event_pattern",
"role_arn",
"schedule_expression",
"state",
"tags",
)
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,
**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.schedule_expression = schedule_expression
self.state = state
self.tags = tags
[docs] 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,
)
[docs]class EventBridgeEnableRuleOperator(AwsBaseOperator[EventBridgeHook]):
"""
Enable an EventBridge Rule.
.. seealso::
For more information on how to use this operator, take a look at the guide:
:ref:`howto/operator:EventBridgeEnableRuleOperator`
:param name: the name of the rule to enable
:param event_bus_name: the name or ARN of the event bus associated with the rule (default if omitted)
:param aws_conn_id: The Airflow connection used for AWS credentials.
If this is ``None`` or empty then the default boto3 behaviour is used. If
running Airflow in a distributed manner and aws_conn_id is None or
empty, then default boto3 configuration would be used (and must be
maintained on each worker node).
:param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
:param verify: Whether or not to verify SSL certificates. See:
https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html
:param botocore_config: Configuration dictionary (key-values) for botocore client. See:
https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.htmlt
"""
[docs] aws_hook_class = EventBridgeHook
[docs] template_fields: Sequence[str] = aws_template_fields("name", "event_bus_name")
def __init__(self, *, name: str, event_bus_name: str | None = None, **kwargs):
super().__init__(**kwargs)
self.name = name
self.event_bus_name = event_bus_name
[docs] def execute(self, context: Context):
self.hook.conn.enable_rule(
**prune_dict(
{
"Name": self.name,
"EventBusName": self.event_bus_name,
}
)
)
self.log.info('Enabled rule "%s"', self.name)
[docs]class EventBridgeDisableRuleOperator(AwsBaseOperator[EventBridgeHook]):
"""
Disable an EventBridge Rule.
.. seealso::
For more information on how to use this operator, take a look at the guide:
:ref:`howto/operator:EventBridgeDisableRuleOperator`
:param name: the name of the rule to disable
:param event_bus_name: the name or ARN of the event bus associated with the rule (default if omitted)
:param aws_conn_id: The Airflow connection used for AWS credentials.
If this is ``None`` or empty then the default boto3 behaviour is used. If
running Airflow in a distributed manner and aws_conn_id is None or
empty, then default boto3 configuration would be used (and must be
maintained on each worker node).
:param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
:param verify: Whether or not to verify SSL certificates. See:
https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html
:param botocore_config: Configuration dictionary (key-values) for botocore client. See:
https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.htmlt
"""
[docs] aws_hook_class = EventBridgeHook
[docs] template_fields: Sequence[str] = aws_template_fields("name", "event_bus_name")
def __init__(self, *, name: str, event_bus_name: str | None = None, **kwargs):
super().__init__(**kwargs)
self.name = name
self.event_bus_name = event_bus_name
[docs] def execute(self, context: Context):
self.hook.conn.disable_rule(
**prune_dict(
{
"Name": self.name,
"EventBusName": self.event_bus_name,
}
)
)
self.log.info('Disabled rule "%s"', self.name)