Source code for tests.system.providers.amazon.aws.example_redshift
## 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__importannotationsfromdatetimeimportdatetimeimportboto3fromairflowimportDAG,settingsfromairflow.decoratorsimporttaskfromairflow.modelsimportConnectionfromairflow.models.baseoperatorimportchainfromairflow.providers.amazon.aws.hooks.redshift_clusterimportRedshiftHookfromairflow.providers.amazon.aws.operators.redshift_clusterimport(RedshiftCreateClusterOperator,RedshiftCreateClusterSnapshotOperator,RedshiftDeleteClusterOperator,RedshiftDeleteClusterSnapshotOperator,RedshiftPauseClusterOperator,RedshiftResumeClusterOperator,)fromairflow.providers.amazon.aws.operators.redshift_dataimportRedshiftDataOperatorfromairflow.providers.amazon.aws.operators.redshift_sqlimportRedshiftSQLOperatorfromairflow.providers.amazon.aws.sensors.redshift_clusterimportRedshiftClusterSensorfromairflow.utils.trigger_ruleimportTriggerRulefromtests.system.providers.amazon.aws.utilsimportENV_ID_KEY,SystemTestContextBuilder
env_id=test_context[ENV_ID_KEY]redshift_cluster_identifier=f"{env_id}-redshift-cluster"redshift_cluster_snapshot_identifier=f"{env_id}-snapshot"conn_id_name=f"{env_id}-conn-id"sg_name=f"{env_id}-sg"set_up_sg=setup_security_group(sec_group_name=sg_name,ip_permissions=[IP_PERMISSION])# [START howto_operator_redshift_cluster]create_cluster=RedshiftCreateClusterOperator(task_id="create_cluster",cluster_identifier=redshift_cluster_identifier,vpc_security_group_ids=[set_up_sg],publicly_accessible=True,cluster_type="single-node",node_type="dc2.large",master_username=DB_LOGIN,master_user_password=DB_PASS,)# [END howto_operator_redshift_cluster]# [START howto_sensor_redshift_cluster]wait_cluster_available=RedshiftClusterSensor(task_id="wait_cluster_available",cluster_identifier=redshift_cluster_identifier,target_status="available",poke_interval=15,timeout=60*15,)# [END howto_sensor_redshift_cluster]# [START howto_operator_redshift_create_cluster_snapshot]create_cluster_snapshot=RedshiftCreateClusterSnapshotOperator(task_id="create_cluster_snapshot",cluster_identifier=redshift_cluster_identifier,snapshot_identifier=redshift_cluster_snapshot_identifier,poll_interval=30,max_attempt=100,retention_period=1,wait_for_completion=True,)# [END howto_operator_redshift_create_cluster_snapshot]wait_cluster_available_before_pause=RedshiftClusterSensor(task_id="wait_cluster_available_before_pause",cluster_identifier=redshift_cluster_identifier,target_status="available",poke_interval=15,timeout=60*15,)# [START howto_operator_redshift_pause_cluster]pause_cluster=RedshiftPauseClusterOperator(task_id="pause_cluster",cluster_identifier=redshift_cluster_identifier,)# [END howto_operator_redshift_pause_cluster]wait_cluster_paused=RedshiftClusterSensor(task_id="wait_cluster_paused",cluster_identifier=redshift_cluster_identifier,target_status="paused",poke_interval=15,timeout=60*15,)# [START howto_operator_redshift_resume_cluster]resume_cluster=RedshiftResumeClusterOperator(task_id="resume_cluster",cluster_identifier=redshift_cluster_identifier,)# [END howto_operator_redshift_resume_cluster]wait_cluster_available_after_resume=RedshiftClusterSensor(task_id="wait_cluster_available_after_resume",cluster_identifier=redshift_cluster_identifier,target_status="available",poke_interval=15,timeout=60*15,)set_up_connection=create_connection(conn_id_name,cluster_id=redshift_cluster_identifier)# [START howto_operator_redshift_data]create_table_redshift_data=RedshiftDataOperator(task_id="create_table_redshift_data",cluster_identifier=redshift_cluster_identifier,database=DB_NAME,db_user=DB_LOGIN,sql=""" CREATE TABLE IF NOT EXISTS fruit ( fruit_id INTEGER, name VARCHAR NOT NULL, color VARCHAR NOT NULL ); """,poll_interval=POLL_INTERVAL,wait_for_completion=True,)# [END howto_operator_redshift_data]insert_data=RedshiftDataOperator(task_id="insert_data",cluster_identifier=redshift_cluster_identifier,database=DB_NAME,db_user=DB_LOGIN,sql=""" INSERT INTO fruit VALUES ( 1, 'Banana', 'Yellow'); INSERT INTO fruit VALUES ( 2, 'Apple', 'Red'); INSERT INTO fruit VALUES ( 3, 'Lemon', 'Yellow'); INSERT INTO fruit VALUES ( 4, 'Grape', 'Purple'); INSERT INTO fruit VALUES ( 5, 'Pear', 'Green'); INSERT INTO fruit VALUES ( 6, 'Strawberry', 'Red'); """,poll_interval=POLL_INTERVAL,wait_for_completion=True,)# [START howto_operator_redshift_sql]select_data=RedshiftSQLOperator(task_id="select_data",redshift_conn_id=conn_id_name,sql="""CREATE TABLE more_fruit AS SELECT * FROM fruit;""",)# [END howto_operator_redshift_sql]# [START howto_operator_redshift_sql_with_params]select_filtered_data=RedshiftSQLOperator(task_id="select_filtered_data",redshift_conn_id=conn_id_name,sql="""CREATE TABLE filtered_fruit AS SELECT * FROM fruit WHERE color = '{{ params.color }}';""",params={"color":"Red"},)# [END howto_operator_redshift_sql_with_params]drop_table=RedshiftSQLOperator(task_id="drop_table",redshift_conn_id=conn_id_name,sql="DROP TABLE IF EXISTS fruit",trigger_rule=TriggerRule.ALL_DONE,)# [START howto_operator_redshift_delete_cluster]delete_cluster=RedshiftDeleteClusterOperator(task_id="delete_cluster",cluster_identifier=redshift_cluster_identifier,)# [END howto_operator_redshift_delete_cluster]delete_cluster.trigger_rule=TriggerRule.ALL_DONE# [START howto_operator_redshift_delete_cluster_snapshot]delete_cluster_snapshot=RedshiftDeleteClusterSnapshotOperator(task_id="delete_cluster_snapshot",cluster_identifier=redshift_cluster_identifier,snapshot_identifier=redshift_cluster_snapshot_identifier,)# [END howto_operator_redshift_delete_cluster_snapshot]delete_sg=delete_security_group(sec_group_id=set_up_sg,sec_group_name=sg_name,)chain(# TEST SETUPtest_context,set_up_sg,# TEST BODYcreate_cluster,wait_cluster_available,create_cluster_snapshot,wait_cluster_available_before_pause,pause_cluster,wait_cluster_paused,resume_cluster,wait_cluster_available_after_resume,set_up_connection,create_table_redshift_data,insert_data,[select_data,select_filtered_data],drop_table,delete_cluster_snapshot,delete_cluster,# TEST TEARDOWNdelete_sg,)fromtests.system.utils.watcherimportwatcher# This test needs watcher in order to properly mark success/failure# when "tearDown" task with trigger rule is part of the DAGlist(dag.tasks)>>watcher()fromtests.system.utilsimportget_test_run# noqa: E402# Needed to run the example DAG with pytest (see: tests/system/README.md#run_via_pytest)