You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by po...@apache.org on 2023/08/08 20:29:58 UTC
[airflow] branch main updated: Implement EventBridge enable and disable rule operators (#33226)
This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 741fb27f50 Implement EventBridge enable and disable rule operators (#33226)
741fb27f50 is described below
commit 741fb27f50a9fee7fca776abf24c32e22783aa03
Author: Maham Ali <50...@users.noreply.github.com>
AuthorDate: Tue Aug 8 13:29:50 2023 -0700
Implement EventBridge enable and disable rule operators (#33226)
Co-authored-by: Maham Ali <ma...@amazon.com>
---
airflow/providers/amazon/aws/hooks/eventbridge.py | 4 +-
.../providers/amazon/aws/operators/eventbridge.py | 105 ++++++++++++++++++++-
.../operators/eventbridge.rst | 34 ++++++-
.../providers/amazon/aws/hooks/test_eventbridge.py | 12 ++-
.../amazon/aws/operators/test_eventbridge.py | 55 ++++++++++-
.../providers/amazon/aws/example_eventbridge.py | 17 +++-
6 files changed, 216 insertions(+), 11 deletions(-)
diff --git a/airflow/providers/amazon/aws/hooks/eventbridge.py b/airflow/providers/amazon/aws/hooks/eventbridge.py
index 41c3de871f..4cc4a1ddc4 100644
--- a/airflow/providers/amazon/aws/hooks/eventbridge.py
+++ b/airflow/providers/amazon/aws/hooks/eventbridge.py
@@ -19,7 +19,7 @@ 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
+from airflow.utils.helpers import prune_dict
def _validate_json(pattern: str) -> None:
@@ -73,7 +73,7 @@ class EventBridgeHook(AwsBaseHook):
_validate_json(event_pattern)
put_rule_kwargs: dict[str, str | list] = {
- **trim_none_values(
+ **prune_dict(
{
"Name": name,
"Description": description,
diff --git a/airflow/providers/amazon/aws/operators/eventbridge.py b/airflow/providers/amazon/aws/operators/eventbridge.py
index 5e5a8f15ce..88cf8b46e4 100644
--- a/airflow/providers/amazon/aws/operators/eventbridge.py
+++ b/airflow/providers/amazon/aws/operators/eventbridge.py
@@ -22,7 +22,7 @@ from typing import TYPE_CHECKING, Sequence
from airflow import AirflowException
from airflow.models import BaseOperator
from airflow.providers.amazon.aws.hooks.eventbridge import EventBridgeHook
-from airflow.providers.amazon.aws.utils import trim_none_values
+from airflow.utils.helpers import prune_dict
if TYPE_CHECKING:
from airflow.utils.context import Context
@@ -68,7 +68,7 @@ class EventBridgePutEventsOperator(BaseOperator):
def execute(self, context: Context):
response = self.hook.conn.put_events(
- **trim_none_values(
+ **prune_dict(
{
"Entries": self.entries,
"EndpointId": self.endpoint_id,
@@ -170,3 +170,104 @@ class EventBridgePutRuleOperator(BaseOperator):
state=self.state,
tags=self.tags,
)
+
+
+class EventBridgeEnableRuleOperator(BaseOperator):
+ """
+ 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 AWS connection to use
+ :param region_name: the region of the rule to be enabled
+
+ """
+
+ template_fields: Sequence[str] = ("name", "event_bus_name", "region_name", "aws_conn_id")
+
+ def __init__(
+ self,
+ *,
+ name: str,
+ event_bus_name: str | None = None,
+ region_name: str | None = None,
+ aws_conn_id: str = "aws_default",
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.name = name
+ self.event_bus_name = event_bus_name
+ self.region_name = region_name
+ 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.hook.conn.enable_rule(
+ **prune_dict(
+ {
+ "Name": self.name,
+ "EventBusName": self.event_bus_name,
+ }
+ )
+ )
+
+ self.log.info('Enabled rule "%s"', self.name)
+
+
+class EventBridgeDisableRuleOperator(BaseOperator):
+ """
+ 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 AWS connection to use
+ :param region_name: the region of the rule to be disabled
+
+ """
+
+ template_fields: Sequence[str] = ("name", "event_bus_name", "region_name", "aws_conn_id")
+
+ def __init__(
+ self,
+ *,
+ name: str,
+ event_bus_name: str | None = None,
+ region_name: str | None = None,
+ aws_conn_id: str = "aws_default",
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.name = name
+ self.event_bus_name = event_bus_name
+ self.region_name = region_name
+ 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.hook.conn.disable_rule(
+ **prune_dict(
+ {
+ "Name": self.name,
+ "EventBusName": self.event_bus_name,
+ }
+ )
+ )
+
+ self.log.info('Disabled rule "%s"', self.name)
diff --git a/docs/apache-airflow-providers-amazon/operators/eventbridge.rst b/docs/apache-airflow-providers-amazon/operators/eventbridge.rst
index 0e6a04d49e..302f3657ec 100644
--- a/docs/apache-airflow-providers-amazon/operators/eventbridge.rst
+++ b/docs/apache-airflow-providers-amazon/operators/eventbridge.rst
@@ -53,7 +53,7 @@ To send custom events to Amazon EventBridge, use
Create or update a rule on Amazon EventBridge
-======================================================
+=============================================
To create or update a rule on EventBridge, use
:class:`~airflow.providers.amazon.aws.operators.eventbridge.EventBridgePutRuleOperator`.
@@ -65,6 +65,38 @@ To create or update a rule on EventBridge, use
:end-before: [END howto_operator_eventbridge_put_rule]
+.. _howto/operator:EventBridgeEnableRuleOperator:
+
+Enable a rule on Amazon EventBridge
+===================================
+
+To enable an existing rule on EventBridge, use
+:class:`~airflow.providers.amazon.aws.operators.eventbridge.EventBridgeEnableRuleOperator`.
+
+.. exampleinclude:: /../../tests/system/providers/amazon/aws/example_eventbridge.py
+ :language: python
+ :dedent: 4
+ :start-after: [START howto_operator_eventbridge_enable_rule]
+ :end-before: [END howto_operator_eventbridge_enable_rule]
+
+
+.. _howto/operator:EventBridgeDisableRuleOperator:
+
+Disable a rule on Amazon EventBridge
+====================================
+
+To disable an existing rule on EventBridge, use
+:class:`~airflow.providers.amazon.aws.operators.eventbridge.EventBridgeDisableRuleOperator`.
+
+.. exampleinclude:: /../../tests/system/providers/amazon/aws/example_eventbridge.py
+ :language: python
+ :dedent: 4
+ :start-after: [START howto_operator_eventbridge_disable_rule]
+ :end-before: [END howto_operator_eventbridge_disable_rule]
+
+
+
+
Reference
---------
diff --git a/tests/providers/amazon/aws/hooks/test_eventbridge.py b/tests/providers/amazon/aws/hooks/test_eventbridge.py
index 1693f56811..1659ba1d8d 100644
--- a/tests/providers/amazon/aws/hooks/test_eventbridge.py
+++ b/tests/providers/amazon/aws/hooks/test_eventbridge.py
@@ -30,10 +30,18 @@ class TestEventBridgeHook:
def test_put_rule(self):
hook = EventBridgeHook(aws_conn_id="aws_default")
- response = hook.put_rule(name="test", event_pattern='{"source": ["aws.s3"]}', state="ENABLED")
+ response = hook.put_rule(
+ name="test",
+ event_pattern='{"source": ["aws.s3"]}',
+ state="ENABLED",
+ )
assert "RuleArn" in response
def test_put_rule_with_bad_json_fails(self):
hook = EventBridgeHook(aws_conn_id="aws_default")
with pytest.raises(ValueError):
- hook.put_rule(name="test", event_pattern="invalid json", state="ENABLED")
+ hook.put_rule(
+ name="test",
+ event_pattern="invalid json",
+ state="ENABLED",
+ )
diff --git a/tests/providers/amazon/aws/operators/test_eventbridge.py b/tests/providers/amazon/aws/operators/test_eventbridge.py
index c6a5b7e916..8ff5703d4f 100644
--- a/tests/providers/amazon/aws/operators/test_eventbridge.py
+++ b/tests/providers/amazon/aws/operators/test_eventbridge.py
@@ -24,6 +24,8 @@ import pytest
from airflow import AirflowException
from airflow.providers.amazon.aws.hooks.eventbridge import EventBridgeHook
from airflow.providers.amazon.aws.operators.eventbridge import (
+ EventBridgeDisableRuleOperator,
+ EventBridgeEnableRuleOperator,
EventBridgePutEventsOperator,
EventBridgePutRuleOperator,
)
@@ -31,6 +33,7 @@ from airflow.providers.amazon.aws.operators.eventbridge import (
ENTRIES = [{"Detail": "test-detail", "Source": "test-source", "DetailType": "test-detail-type"}]
FAILED_ENTRIES_RESPONSE = [{"ErrorCode": "test_code"}, {"ErrorCode": "test_code"}]
EVENT_PATTERN = '{"source": ["aws.s3"]}'
+RULE_NAME = "match_s3_events"
class TestEventBridgePutEventsOperator:
@@ -78,7 +81,7 @@ class TestEventBridgePutEventsOperator:
class TestEventBridgePutRuleOperator:
def test_init(self):
operator = EventBridgePutRuleOperator(
- task_id="events_put_rule_job", name="match_s3_events", event_pattern=EVENT_PATTERN
+ task_id="events_put_rule_job", name=RULE_NAME, event_pattern=EVENT_PATTERN
)
assert operator.event_pattern == EVENT_PATTERN
@@ -89,7 +92,9 @@ class TestEventBridgePutRuleOperator:
mock_conn.put_rule.return_value = hook_response
operator = EventBridgePutRuleOperator(
- task_id="events_put_rule_job", name="match_s3_events", event_pattern=EVENT_PATTERN
+ task_id="events_put_rule_job",
+ name=RULE_NAME,
+ event_pattern=EVENT_PATTERN,
)
result = operator.execute(None)
@@ -98,8 +103,52 @@ class TestEventBridgePutRuleOperator:
def test_put_rule_with_bad_json_fails(self):
operator = EventBridgePutRuleOperator(
- task_id="failed_put_rule_job", name="match_s3_events", event_pattern="invalid json"
+ task_id="failed_put_rule_job",
+ name=RULE_NAME,
+ event_pattern="invalid json",
)
with pytest.raises(ValueError):
operator.execute(None)
+
+
+class TestEventBridgeEnableRuleOperator:
+ def test_init(self):
+ operator = EventBridgeDisableRuleOperator(
+ task_id="enable_rule_task",
+ name=RULE_NAME,
+ )
+
+ assert operator.name == RULE_NAME
+
+ @mock.patch.object(EventBridgeHook, "conn")
+ def test_enable_rule(self, mock_conn: MagicMock):
+
+ enable_rule = EventBridgeEnableRuleOperator(
+ task_id="events_enable_rule_job",
+ name=RULE_NAME,
+ )
+
+ enable_rule.execute(None)
+ mock_conn.enable_rule.assert_called_with(Name=RULE_NAME)
+
+
+class TestEventBridgeDisableRuleOperator:
+ def test_init(self):
+ operator = EventBridgeDisableRuleOperator(
+ task_id="disable_rule_task",
+ name=RULE_NAME,
+ )
+
+ assert operator.name == RULE_NAME
+
+ @mock.patch.object(EventBridgeHook, "conn")
+ def test_disable_rule(self, mock_conn: MagicMock):
+
+ disable_rule = EventBridgeDisableRuleOperator(
+ task_id="events_disable_rule_job",
+ name=RULE_NAME,
+ )
+
+ disable_rule.execute(None)
+ mock_conn.disable_rule.assert_called_with(Name=RULE_NAME)
diff --git a/tests/system/providers/amazon/aws/example_eventbridge.py b/tests/system/providers/amazon/aws/example_eventbridge.py
index 4d052ee510..b1c4fb3f38 100644
--- a/tests/system/providers/amazon/aws/example_eventbridge.py
+++ b/tests/system/providers/amazon/aws/example_eventbridge.py
@@ -21,6 +21,8 @@ from datetime import datetime
from airflow import DAG
from airflow.models.baseoperator import chain
from airflow.providers.amazon.aws.operators.eventbridge import (
+ EventBridgeDisableRuleOperator,
+ EventBridgeEnableRuleOperator,
EventBridgePutEventsOperator,
EventBridgePutRuleOperator,
)
@@ -36,6 +38,7 @@ ENTRIES = [
}
]
+
sys_test_context_task = SystemTestContextBuilder().build()
with DAG(
@@ -59,10 +62,22 @@ with DAG(
name="example_rule",
event_pattern='{"source": ["example.myapp"]}',
description="This rule matches events from example.myapp.",
+ state="DISABLED",
)
# [END howto_operator_eventbridge_put_rule]
- chain(test_context, put_events, put_rule)
+ # [START howto_operator_eventbridge_enable_rule]
+ enable_rule = EventBridgeEnableRuleOperator(task_id="enable_rule_task", name="example_rule")
+ # [END howto_operator_eventbridge_enable_rule]
+
+ # [START howto_operator_eventbridge_disable_rule]
+ disable_rule = EventBridgeDisableRuleOperator(
+ task_id="disable_rule_task",
+ name="example_rule",
+ )
+ # [END howto_operator_eventbridge_disable_rule]
+
+ chain(test_context, put_events, put_rule, enable_rule, disable_rule)
from tests.system.utils import get_test_run # noqa: E402