Simplify/improve several DAGs

dags/data_quality_checks_dag.py

@@ -20,7 +20,6 @@
 ACCESS_KEY_ID=<your_aws_access_key>
 SECRET_ACCESS_KEY=<your_aws_secret_key>
 """
-
 import os
 import pendulum
 from airflow.decorators import dag, task, task_group
@@ -57,13 +56,15 @@ def slack_failure_notification(context):
     log_url = context.get("task_instance").log_url
     slack_msg = f"""
             :red_circle: Task Failed. 
-            *Task*: {task_id}  
-            *Dag*: {dag_id} 
+            *Task*: {task_id}
+            *DAG*: {dag_id}
             *Execution Time*: {exec_date}  
-            *Log Url*: {log_url} 
+            *Log URL*: {log_url}
             """
     failed_alert = SlackWebhookOperator(
-        task_id="slack_notification", http_conn_id="slack_webhook", message=slack_msg
+        task_id="slack_notification",
+        http_conn_id="slack_webhook",
+        message=slack_msg,
     )
     return failed_alert.execute(context=context)
 
@@ -78,22 +79,21 @@ def get_files_from_s3(bucket, prefix_value):
     return list(filter(lambda element: element.endswith(".csv"), paths))
 
 
-@task
-def get_import_statements(files):
-    statements = []
-    for path in files:
-        sql = f"""
-                COPY {TEMP_TABLE} FROM 's3://{ACCESS_KEY_ID}:{SECRET_ACCESS_KEY}@{S3_BUCKET}/{path}' WITH (format='csv');
-              """
-        statements.append(sql)
-    return statements
-
-
 @task
 def list_local_files(directory):
     return list(filter(lambda entry: entry.endswith(".csv"), os.listdir(directory)))
 
 
+def copy_file_kwargs(file):
+    return {
+        "temp_table": TEMP_TABLE,
+        "access_key_id": ACCESS_KEY_ID,
+        "secret_access_key": SECRET_ACCESS_KEY,
+        "s3_bucket": S3_BUCKET,
+        "path": file,
+    }
+
+
 def upload_kwargs(file):
     return {
         "filename": f"{FILE_DIR}/{file}",
@@ -181,41 +181,41 @@ def move_incoming_files(s3_files):
 def data_quality_checks():
     upload = upload_local_files()
     s3_files = get_files_from_s3(S3_BUCKET, INCOMING_DATA_PREFIX)
-    import_stmt = get_import_statements(s3_files)
 
+    # pylint: disable=E1101
     import_data = SQLExecuteQueryOperator.partial(
         task_id="import_data_to_cratedb",
         conn_id="cratedb_connection",
-    ).expand(sql=import_stmt)
+        sql="""
+            COPY {{params.temp_table}}
+            FROM 's3://{{params.acces_key_id}}:{{params.secret_access_key}}@{{params.s3_bucket}}/{{params.path}}'
+            WITH (format = 'csv');
+            """,
+    ).expand(params=s3_files.map(upload_kwargs))
 
     refresh = SQLExecuteQueryOperator(
         task_id="refresh_table",
         conn_id="cratedb_connection",
-        sql="""
-                REFRESH TABLE {{params.temp_table}};   
-            """,
-        params={
-            "temp_table": TEMP_TABLE,
-        },
+        sql="REFRESH TABLE {{params.temp_table}};",
+        params={"temp_table": TEMP_TABLE},
     )
 
     checks = home_data_checks()
 
     move_data = SQLExecuteQueryOperator(
         task_id="move_to_table",
         conn_id="cratedb_connection",
-        sql="""
-                INSERT INTO {{params.table}} SELECT * FROM {{params.temp_table}};   
-            """,
-        params={"table": TABLE, "temp_table": TEMP_TABLE},
+        sql="INSERT INTO {{params.table}} SELECT * FROM {{params.temp_table}};",
+        params={
+            "table": TABLE,
+            "temp_table": TEMP_TABLE,
+        },
     )
 
     delete_data = SQLExecuteQueryOperator(
         task_id="delete_from_temp_table",
         conn_id="cratedb_connection",
-        sql="""
-                DELETE FROM {{params.temp_table}};   
-            """,
+        sql="DELETE FROM {{params.temp_table}};",
         params={"temp_table": TEMP_TABLE},
         trigger_rule="all_done",
     )
@@ -231,7 +231,6 @@ def data_quality_checks():
     chain(
         upload,
         s3_files,
-        import_stmt,
         import_data,
         refresh,
         checks,

dags/data_retention_delete_dag.py

@@ -15,18 +15,12 @@
 from airflow.decorators import dag, task
 
 
-@task
-def generate_sql(policy):
-    """Generate DROP statment for a given partition"""
-    return (
-        Path("include/data_retention_delete.sql")
-        .read_text(encoding="utf-8")
-        .format(
-            table_fqn=policy[0],
-            column=policy[1],
-            value=policy[2],
-        )
-    )
+def map_policy(policy):
+    return {
+        "table_fqn": policy[0],
+        "column": policy[1],
+        "value": policy[2],
+    }
 
 
 @task
@@ -35,7 +29,8 @@ def get_policies(ds=None):
     pg_hook = PostgresHook(postgres_conn_id="cratedb_connection")
     sql = Path("include/data_retention_retrieve_delete_policies.sql")
     return pg_hook.get_records(
-        sql=sql.read_text(encoding="utf-8"), parameters={"day": ds}
+        sql=sql.read_text(encoding="utf-8"),
+        parameters={"day": ds},
     )
 
 
@@ -45,12 +40,11 @@ def get_policies(ds=None):
     catchup=False,
 )
 def data_retention_delete():
-    sql_statements = generate_sql.expand(policy=get_policies())
-
     SQLExecuteQueryOperator.partial(
         task_id="delete_partition",
         conn_id="cratedb_connection",
-    ).expand(sql=sql_statements)
+        sql="DELETE FROM {{params.table_fqn}} WHERE {{params.column}} = {{params.value}};",
+    ).expand(params=get_policies().map(map_policy))
 
 
 data_retention_delete()

dags/data_retention_reallocate_dag.py

@@ -5,8 +5,9 @@
 
 Prerequisites
 -------------
-In CrateDB, tables for storing retention policies need to be created once manually.
-See the file setup/data_retention_schema.sql in this repository.
+- CrateDB 5.2.0 or later
+- Tables for storing retention policies need to be created once manually in
+  CrateDB. See the file setup/data_retention_schema.sql in this repository.
 """
 from pathlib import Path
 import pendulum
@@ -21,11 +22,11 @@ def get_policies(ds=None):
     pg_hook = PostgresHook(postgres_conn_id="cratedb_connection")
     sql = Path("include/data_retention_retrieve_reallocate_policies.sql")
     return pg_hook.get_records(
-        sql=sql.read_text(encoding="utf-8"), parameters={"day": ds}
+        sql=sql.read_text(encoding="utf-8"),
+        parameters={"day": ds},
     )
 
 
-@task
 def map_policy(policy):
     """Map index-based policy to readable dict structure"""
     return {
@@ -39,37 +40,21 @@ def map_policy(policy):
     }
 
 
-@task
-def generate_sql_reallocate(policy):
-    """Generate SQL for reallocation"""
-    return (
-        Path("include/data_retention_reallocate.sql")
-        .read_text(encoding="utf-8")
-        .format(**policy)
-    )
-
-
 @dag(
     start_date=pendulum.datetime(2021, 11, 19, tz="UTC"),
     schedule="@daily",
     catchup=False,
     template_searchpath=["include"],
 )
 def data_retention_reallocate():
-    policies = map_policy.expand(policy=get_policies())
-
-    reallocate = SQLExecuteQueryOperator.partial(
+    SQLExecuteQueryOperator.partial(
         task_id="reallocate_partitions",
         conn_id="cratedb_connection",
-    ).expand(sql=generate_sql_reallocate.expand(policy=policies))
-
-    track = SQLExecuteQueryOperator.partial(
-        task_id="add_tracking_information",
-        conn_id="cratedb_connection",
-        sql="data_retention_reallocate_tracking.sql",
-    ).expand(parameters=policies)
-
-    reallocate >> track
+        sql="""
+            ALTER TABLE {{params.table_fqn}} PARTITION ({{params.column}} = {{params.value}})
+            SET ("routing.allocation.require.{{params.attribute_name}}" = '{{params.attribute_value}}');
+            """,
+    ).expand(params=get_policies().map(map_policy))
 
 
 data_retention_reallocate()

dags/data_retention_snapshot_dag.py

@@ -25,7 +25,6 @@ def get_policies(ds=None):
     )
 
 
-@task
 def map_policy(policy):
     """Map index-based policy to readable dict structure"""
     return {
@@ -38,35 +37,29 @@ def map_policy(policy):
     }
 
 
-@task
-def generate_sql(query_file, policy):
-    """Generate SQL statment for a given partition"""
-    return Path(query_file).read_text(encoding="utf-8").format(**policy)
-
-
 @dag(
     start_date=pendulum.datetime(2021, 11, 19, tz="UTC"),
     schedule="@daily",
     catchup=False,
 )
 def data_retention_snapshot():
-    policies = map_policy.expand(policy=get_policies())
-    sql_statements_snapshot = generate_sql.partial(
-        query_file="include/data_retention_snapshot.sql"
-    ).expand(policy=policies)
-    sql_statements_delete = generate_sql.partial(
-        query_file="include/data_retention_delete.sql"
-    ).expand(policy=policies)
+    policies = get_policies().map(map_policy)
 
     reallocate = SQLExecuteQueryOperator.partial(
         task_id="snapshot_partitions",
         conn_id="cratedb_connection",
-    ).expand(sql=sql_statements_snapshot)
+        sql="""
+            CREATE SNAPSHOT {{params.target_repository_name}}."{{params.schema}}.{{params.table}}-{{params.value}}"
+            TABLE {{params.table_fqn}} PARTITION ({{params.column}} = {{params.value}})
+            WITH ("wait_for_completion" = true);
+            """,
+    ).expand(params=policies)
 
     delete = SQLExecuteQueryOperator.partial(
         task_id="delete_partitions",
         conn_id="cratedb_connection",
-    ).expand(sql=sql_statements_delete)
+        sql="DELETE FROM {{params.table_fqn}} WHERE {{params.column}} = {{params.value}};",
+    ).expand(params=policies)
 
     reallocate >> delete