dayrize-usecase/dags/sustainability_score/__init__.py

72 lines
2.1 KiB
Python

"""
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