""" DAG IDs: sustainability_score """ import os from datetime import datetime import utils from airflow import DAG from airflow.operators.dummy_operator import DummyOperator from airflow.providers.apache.beam.operators.beam import BeamRunPythonPipelineOperator from airflow.providers.postgres.operators.postgres import PostgresOperator HOME = os.environ["HOME"] CSV_FNAME = ( "large_target_store_products_dataset_sample - " "large_target_store_products_dataset_sample.csv" ) CONFIG = { "input": f"{ HOME }/gcs/data/{ CSV_FNAME }", "beam_etl_path": "/etl/main.py", "products_table": "sustainability_score.products", "scored_table": "sustainability_score.scored_products", } with DAG( "sustainability_score", schedule_interval="0 * * * 1-5", catchup=False, max_active_runs=10, start_date=datetime(2023, 6, 21), doc_md=utils.load_docs(__file__), params=CONFIG, template_searchpath=["/sql"], ) as dag: create_products_table = PostgresOperator( task_id="create_products_table", sql="products_schema.sql", postgres_conn_id="pg_db", ) create_scores_table = PostgresOperator( task_id="create_scored_products_table", sql="scored_products_schema.sql", postgres_conn_id="pg_db", ) etl_pipeline = BeamRunPythonPipelineOperator( task_id="beam_etl", py_file="{{ params.beam_etl_path }}", pipeline_options={ "input": "{{ params.input }}", "pg_hostname": "{{ conn.get('pg_db').host }}", "pg_port": "{{ conn.get('pg_db').port }}", "pg_username": "{{ conn.get('pg_db').login }}", "pg_password": "{{ conn.get('pg_db').password }}", "pg_database": "{{ conn.get('pg_db').schema }}", "pg_table": "{{ params.products_table }}", }, ) calculate_score = PostgresOperator( task_id="calculate_score", sql="calculate_score.sql", postgres_conn_id="pg_db", ) create_products_table >> etl_pipeline [etl_pipeline, create_scores_table] >> calculate_score