diff --git a/airflow/providers/amazon/aws/sensors/rds.py b/airflow/providers/amazon/aws/sensors/rds.py index 3fdc37543d11b..6240f10a91602 100644 --- a/airflow/providers/amazon/aws/sensors/rds.py +++ b/airflow/providers/amazon/aws/sensors/rds.py @@ -174,6 +174,12 @@ class RdsDbSensor(RdsBaseSensor): :param target_statuses: Target status of DB """ + template_fields: Sequence[str] = ( + 'db_identifier', + 'db_type', + 'target_statuses', + ) + def __init__( self, *, diff --git a/docs/apache-airflow-providers-amazon/operators/rds.rst b/docs/apache-airflow-providers-amazon/operators/rds.rst index e5f46ad66be42..d0752a014d0d3 100644 --- a/docs/apache-airflow-providers-amazon/operators/rds.rst +++ b/docs/apache-airflow-providers-amazon/operators/rds.rst @@ -146,7 +146,7 @@ Create a database instance To create a AWS DB instance you can use :class:`~airflow.providers.amazon.aws.operators.rds.RdsCreateDbInstanceOperator`. -.. exampleinclude:: /../../tests/system/providers/amazon/aws/rds/example_rds_instance.py +.. exampleinclude:: /../../tests/system/providers/amazon/aws/example_rds_instance.py :language: python :dedent: 4 :start-after: [START howto_operator_rds_create_db_instance] @@ -160,7 +160,7 @@ Delete a database instance To delete a AWS DB instance you can use :class:`~airflow.providers.amazon.aws.operators.rds.RDSDeleteDbInstanceOperator`. -.. exampleinclude:: /../../tests/system/providers/amazon/aws/rds/example_rds_instance.py +.. exampleinclude:: /../../tests/system/providers/amazon/aws/example_rds_instance.py :language: python :dedent: 4 :start-after: [START howto_operator_rds_delete_db_instance] @@ -178,7 +178,7 @@ To wait for an Amazon RDS instance or cluster to reach a specific status you can :class:`~airflow.providers.amazon.aws.sensors.rds.RdsDbSensor`. By default, the sensor waits for a database instance to reach the ``available`` state. -.. exampleinclude:: /../../tests/system/providers/amazon/aws/rds/example_rds_instance.py +.. exampleinclude:: /../../tests/system/providers/amazon/aws/example_rds_instance.py :language: python :dedent: 4 :start-after: [START howto_sensor_rds_instance] diff --git a/tests/system/providers/amazon/aws/rds/example_rds_instance.py b/tests/system/providers/amazon/aws/example_rds_instance.py similarity index 60% rename from tests/system/providers/amazon/aws/rds/example_rds_instance.py rename to tests/system/providers/amazon/aws/example_rds_instance.py index ec6cba2d2550c..c27f7f99eeedb 100644 --- a/tests/system/providers/amazon/aws/rds/example_rds_instance.py +++ b/tests/system/providers/amazon/aws/example_rds_instance.py @@ -1,4 +1,3 @@ -# # 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 @@ -26,12 +25,13 @@ RdsDeleteDbInstanceOperator, ) from airflow.providers.amazon.aws.sensors.rds import RdsDbSensor -from tests.system.providers.amazon.aws.utils import set_env_id +from airflow.utils.trigger_rule import TriggerRule +from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder + +sys_test_context_task = SystemTestContextBuilder().build() -ENV_ID = set_env_id() -DAG_ID = "example_rds_instance" +DAG_ID = 'example_rds_instance' -RDS_DB_IDENTIFIER = f'{ENV_ID}-database' RDS_USERNAME = 'database_username' # NEVER store your production password in plaintext in a DAG like this. # Use Airflow Secrets or a secret manager for this in production. @@ -39,43 +39,64 @@ with DAG( dag_id=DAG_ID, - schedule=None, + schedule='@once', start_date=datetime(2021, 1, 1), tags=['example'], catchup=False, ) as dag: + test_context = sys_test_context_task() + rds_db_identifier = f'{test_context[ENV_ID_KEY]}-database' + # [START howto_operator_rds_create_db_instance] create_db_instance = RdsCreateDbInstanceOperator( task_id='create_db_instance', - db_instance_identifier=RDS_DB_IDENTIFIER, - db_instance_class="db.t4g.micro", - engine="postgres", + db_instance_identifier=rds_db_identifier, + db_instance_class='db.t4g.micro', + engine='postgres', rds_kwargs={ - "MasterUsername": RDS_USERNAME, - "MasterUserPassword": RDS_PASSWORD, - "AllocatedStorage": 20, + 'MasterUsername': RDS_USERNAME, + 'MasterUserPassword': RDS_PASSWORD, + 'AllocatedStorage': 20, }, ) # [END howto_operator_rds_create_db_instance] + # RdsCreateDbInstanceOperator waits by default, setting as False to test the Sensor below. + create_db_instance.wait_for_completion = False + # [START howto_sensor_rds_instance] - db_instance_available = RdsDbSensor( - task_id="db_instance_available", - db_identifier=RDS_DB_IDENTIFIER, + await_db_instance = RdsDbSensor( + task_id='await_db_instance', + db_identifier=rds_db_identifier, ) # [END howto_sensor_rds_instance] # [START howto_operator_rds_delete_db_instance] delete_db_instance = RdsDeleteDbInstanceOperator( task_id='delete_db_instance', - db_instance_identifier=RDS_DB_IDENTIFIER, + db_instance_identifier=rds_db_identifier, rds_kwargs={ - "SkipFinalSnapshot": True, + 'SkipFinalSnapshot': True, }, ) # [END howto_operator_rds_delete_db_instance] + delete_db_instance.trigger_rule = TriggerRule.ALL_DONE + + chain( + # TEST SETUP + test_context, + # TEST BODY + create_db_instance, + await_db_instance, + delete_db_instance, + ) + + from tests.system.utils.watcher import watcher + + # This test needs watcher in order to properly mark success/failure + # when "tearDown" task with trigger rule is part of the DAG + list(dag.tasks) >> watcher() - chain(create_db_instance, db_instance_available, delete_db_instance) from tests.system.utils import get_test_run # noqa: E402 diff --git a/tests/system/providers/amazon/aws/rds/__init__.py b/tests/system/providers/amazon/aws/rds/__init__.py deleted file mode 100644 index 217e5db960782..0000000000000 --- a/tests/system/providers/amazon/aws/rds/__init__.py +++ /dev/null @@ -1,17 +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.