使用 Airflow 编排查询

'''
使用 Apache Drill 来转换、加载和报告 COVID 案例。

数据来源引用：

疾病控制和预防中心，COVID-19 报告。COVID-19 病历监控公共数据访问，摘要和限制。

https://data.cdc.gov/Case-Surveillance/COVID-19-Case-Surveillance-Public-Use-Data/vbim-akqf
'''
from datetime import timedelta

from airflow import DAG
# 使用 PythonOperator 从 CDC 网站缓存 COVID-19 CSV 文件
from airflow.operators.python import PythonOperator
# 使用 DrillOperators 启动对 COVID-19 数据的查询
from airflow.providers.apache.drill.operators.drill import DrillOperator
from airflow.utils.dates import days_ago
# 假设请求是存在的，因为 sqlalchemy-drill 需要这些请求
import requests
# 传递这些参数
# 每个新任务可以传递新的参数
default_args = {
    'owner': 'Joe Public',
    'depends_on_past': False,
    'email': ['joe@public.com'],
    'email_on_failure': False,
    'email_on_retry': False,
    'retries': 1,
    'retry_delay': timedelta(minutes=5),
}


def stage_from_www(src_url, tgt_path):
    '''
    使用 Request 库将病例监控数据从 CDC 下载到本地。如果是分布式环境，用 HDFS、S3 等替换本地文件系统。
    另一种选择是使用 Drill 提供的 HTTP 存储插件来直接从数据源获取数据。
    请注意：避免在内存中缓存大量数据集（启用流式响应）
    '''
    resp = requests.get(
        src_url,
        stream=True  # don't buffer big datasets in memory
    )
    with open(tgt_path) as f:
        f.write(resp.content)


with DAG(
    'drill_tutorial',
    default_args=default_args,
    description='Drill tutorial that loads COVID-19 case data from the CDC.',
    schedule_interval=timedelta(weeks=2),  # 数据源每两周更新一次
    start_date=days_ago(0),
) as dag:

    # 将此模块的字符串用于 DAG 的文档 (在 Web UI 中可视化)。
    dag.doc_md = __doc__

    # 第一个任务是使用 PythonOperator 从 CDC 网站获取 CSV 数据。
    stage_from_www_task = PythonOperator(
        task_id='stage_from_www',
        python_callable=stage_from_www,
        op_kwargs= {
            'src_url': 'https://data.cdc.gov/api/views/vbim-akqf/rows.csv?accessType=DOWNLOAD',
            'tgt_path': '/tmp/cdc_covid_cases.csvh'
        }
    )

    stage_from_www.doc = 'Download COVID case CSV data from the CDC using ' \
        'an HTTP GET'

    # 第二个任务是通过 DrillOperator 的外部脚本执行 CTAS 的 ETL 流程。
    # 也可以指定 SQL，并跨任务拆分此多语句的 SQL 脚本。也就是将初始 DROP TABLE 设为单独的任务。
    ctas_etl_task = DrillOperator(
        drill_conn_id='drill_tutorial',
        task_id='ctas_etl',
        sql='cdc_covid_cases.drill.sql'
    )

    ctas_etl_task.doc = 'Recreate dfs.tmp.cdc_covid_cases using CTAS'

    # 第三个任务是通过 DrillOperator 生成每日案例计数报告。
    # 将报告由 dfs.tmp 格式转换为可读的 CSV 格式, 使用 Airflow 可以用多种方式实现。
    daily_count_report_task = DrillOperator(
        drill_conn_id='drill_tutorial',
        task_id='drill_report',
        sql='''
        set `store.format` = 'csv';

        drop table if exists dfs.tmp.cdc_daily_counts;

        create table dfs.tmp.cdc_daily_counts as
        select
            cdc_case_earliest_dt,
            count(*) as case_count
        from
            dfs.tmp.cdc_covid_cases
        group by
            cdc_case_earliest_dt
        order by
            cdc_case_earliest_dt;
        '''
    )

    daily_count_report_task.doc = 'Report daily case counts to CSV'

    # 指定 DAG 的依赖关系。
    stage_from_www_task >> ctas_etl_task >> daily_count_report_task
age_parse;