diff --git a/dags/configuration.py b/dags/configuration.py new file mode 100644 index 0000000..1f093f0 --- /dev/null +++ b/dags/configuration.py @@ -0,0 +1,85 @@ +""" +Module containing DAG for configuration data + +Outputs: + - configuration data in parquet_datasets bucket + - configuration data in digital land db +""" + +from datetime import timedelta + +import boto3 +from airflow import DAG +from airflow.models.param import Param +from airflow.operators.python import PythonOperator +from airflow.providers.amazon.aws.operators.emr import EmrServerlessStartJobOperator +from emr_dags_utils import get_secrets +from utils import dag_default_args, get_config + +config = get_config() + +with DAG( + dag_id="configuration", + default_args=dag_default_args, + description="run processes related to our configuration files", + schedule=None, + catchup=False, + params={ + "debug": Param(default=False, type="boolean", description="Enable debug logging for the Spark job"), + }, + render_template_as_native_obj=True, + is_paused_upon_creation=False, +) as dag: + + ENV = config["env"] + EXECUTION_ROLE_ARN = get_secrets("emr_execution_role", ENV) + S3_BUCKET = f"{ENV}-pd-batch-jobs-codepackage-bucket" + S3_LOG_BUCKET = f"{ENV}-pd-batch-jobs-logs-bucket" + S3_TASKS_ENTRY_POINT = f"s3://{S3_BUCKET}/pkg/entry_script/run_config.py" + S3_WHEEL_FILE = f"s3://{S3_BUCKET}/pkg/whl_pkg/pyspark_jobs-0.1.0-py3-none-any.whl" + S3_LOG_URI = f"s3://{S3_LOG_BUCKET}/" + S3_DATA_PATH = f"s3://{ENV}-collection-data/" + + def get_tasks_emr_application_id(**context): + app_name = f"{ENV}-pd-batch-emrsl-application" + client = boto3.client("emr-serverless", region_name="eu-west-2") + response = client.list_applications(maxResults=50) + for app in response.get("applications", []): + if app["name"] == app_name: + app_id = app["id"] + context["ti"].xcom_push(key="application_id", value=app_id) + return app_id + raise ValueError(f"EMR application '{app_name}' not found") + + get_tasks_app_id = PythonOperator( + task_id="get-tasks-emr-app-id", + python_callable=get_tasks_emr_application_id, + execution_timeout=timedelta(minutes=2), + dag=dag, + ) + + assemble_tasks_emr_task = EmrServerlessStartJobOperator( + task_id="assemble-tasks", + application_id='{{ task_instance.xcom_pull(task_ids="get-tasks-emr-app-id", key="application_id") }}', + execution_role_arn=EXECUTION_ROLE_ARN, + job_driver={ + "sparkSubmit": { + "entryPoint": S3_TASKS_ENTRY_POINT, + "entryPointArguments": ( + f"{{{{ ['--env', '{ENV}', '--collection-data-path', '{S3_DATA_PATH}', " + f"'--parquet-datasets-path', 's3://{ENV}-parquet-datasets'] " + f"+ (['--debug'] if params.debug else []) }}}}" + ), + "sparkSubmitParameters": f"--py-files {S3_WHEEL_FILE} " "--conf spark.serializer=org.apache.spark.serializer.KryoSerializer", + } + }, + configuration_overrides={"monitoringConfiguration": {"s3MonitoringConfiguration": {"logUri": S3_LOG_URI}}}, + name="assemble-tasks-job", + wait_for_completion=True, + aws_conn_id="aws_default", + waiter_max_attempts=180, + waiter_delay=60, + execution_timeout=timedelta(hours=3), + ) + + get_tasks_app_id >> assemble_tasks_emr_task