Airflow operators

1
import hashlib
2
import json
3
4
from airflow.exceptions import AirflowException
5
from airflow.decorators import dag, task
6
from airflow.models import Variable
7
from airflow.models.baseoperator import chain
8
from airflow.operators.empty import EmptyOperator
9
from airflow.utils.dates import datetime
10
from airflow.providers.amazon.aws.hooks.s3 import S3Hook
11
from airflow.providers.amazon.aws.transfers.local_to_s3 import (
12
    LocalFilesystemToS3Operator,
13
)
14
from airflow.providers.amazon.aws.transfers.s3_to_redshift import S3ToRedshiftOperator
15
from airflow.providers.postgres.operators.postgres import PostgresOperator
16
from airflow.operators.sql import SQLCheckOperator
17
from airflow.utils.task_group import TaskGroup
18
19
20
# The file(s) to upload shouldn't be hardcoded in a production setting,
21
# this is just for demo purposes.
22
CSV_FILE_NAME = "forestfires.csv"
23
CSV_FILE_PATH = f"include/sample_data/forestfire_data/{CSV_FILE_NAME}"
24
25
26
@dag(
27
    "simple_redshift_3",
28
    start_date=datetime(2021, 7, 7),
29
    description="""A sample Airflow DAG to load data from csv files to S3
30
                 and then Redshift, with data integrity and quality checks.""",
31
    schedule=None,
32
    template_searchpath="/usr/local/airflow/include/sql/redshift_examples/",
33
    catchup=False,
34
)
35
def simple_redshift_3():
36
    """
37
    Before running the DAG, set the following in an Airflow
38
    or Environment Variable:
39
    - key: aws_configs
40
    - value: { "s3_bucket": [bucket_name], "s3_key_prefix": [key_prefix],
41
             "redshift_table": [table_name]}
42
    Fully replacing [bucket_name], [key_prefix], and [table_name].
43
    """
44
45
    upload_file = LocalFilesystemToS3Operator(
46
        task_id="upload_to_s3",
47
        filename=CSV_FILE_PATH,
48
        dest_key="{{ var.json.aws_configs.s3_key_prefix }}/" + CSV_FILE_PATH,
49
        dest_bucket="{{ var.json.aws_configs.s3_bucket }}",
50
        aws_conn_id="aws_default",
51
        replace=True,
52
    )
53
54
    @task
55
    def validate_etag():
56
        """
57
        #### Validation task
58
        Check the destination ETag against the local MD5 hash to ensure
59
        the file was uploaded without errors.
60
        """
61
        s3 = S3Hook()
62
        aws_configs = Variable.get("aws_configs", deserialize_json=True)
63
        obj = s3.get_key(
64
            key=f"{aws_configs.get('s3_key_prefix')}/{CSV_FILE_PATH}",
65
            bucket_name=aws_configs.get("s3_bucket"),
66
        )
67
        obj_etag = obj.e_tag.strip('"')
68
        # Change `CSV_FILE_PATH` to `CSV_CORRUPT_FILE_PATH` for the "sad path".
69
        file_hash = hashlib.md5(open(CSV_FILE_PATH).read().encode("utf-8")).hexdigest()
70
        if obj_etag != file_hash:
71
            raise AirflowException(
72
                """Upload Error: Object ETag in S3 did not match
73
                hash of local file."""
74
            )
75
76
    # Tasks that were created using decorators have to be called to be used
77
    validate_file = validate_etag()
78
79
    # --- Create Redshift Table --- #
80
    create_redshift_table = PostgresOperator(
81
        task_id="create_table",
82
        sql="create_redshift_forestfire_table.sql",
83
        postgres_conn_id="redshift_default",
84
    )
85
86
    # --- Second load task --- #
87
    load_to_redshift = S3ToRedshiftOperator(
88
        task_id="load_to_redshift",
89
        s3_bucket="{{ var.json.aws_configs.s3_bucket }}",
90
        s3_key="{{ var.json.aws_configs.s3_key_prefix }}" + f"/{CSV_FILE_PATH}",
91
        schema="PUBLIC",
92
        table="{{ var.json.aws_configs.redshift_table }}",
93
        copy_options=["csv"],
94
    )
95
96
    # --- Redshift row validation task --- #
97
    validate_redshift = SQLCheckOperator(
98
        task_id="validate_redshift",
99
        conn_id="redshift_default",
100
        sql="validate_redshift_forestfire_load.sql",
101
        params={"filename": CSV_FILE_NAME},
102
    )
103
104
    # --- Row-level data quality check --- #
105
    with open("include/validation/forestfire_validation.json") as ffv:
106
        with TaskGroup(group_id="row_quality_checks") as quality_check_group:
107
            ffv_json = json.load(ffv)
108
            for id, values in ffv_json.items():
109
                values["id"] = id
110
                SQLCheckOperator(
111
                    task_id=f"forestfire_row_quality_check_{id}",
112
                    conn_id="redshift_default",
113
                    sql="row_quality_redshift_forestfire_check.sql",
114
                    params=values,
115
                )
116
117
    # --- Drop Redshift table --- #
118
    drop_redshift_table = PostgresOperator(
119
        task_id="drop_table",
120
        sql="drop_redshift_forestfire_table.sql",
121
        postgres_conn_id="redshift_default",
122
    )
123
124
    begin = EmptyOperator(task_id="begin")
125
    end = EmptyOperator(task_id="end")
126
127
    # --- Define task dependencies --- #
128
    chain(
129
        begin,
130
        upload_file,
131
        validate_file,
132
        create_redshift_table,
133
        load_to_redshift,
134
        validate_redshift,
135
        quality_check_group,
136
        drop_redshift_table,
137
        end,
138
    )
139
140
141
simple_redshift_3()