composer-dags/demo_pipeline_composer.py

# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# --------------------------------------------------------------------------------
# Load The Dependencies
# --------------------------------------------------------------------------------

import json
import uuid
import re
from airflow import models
import google.auth
from airflow.models.variable import Variable
from airflow.providers.google.cloud.operators.dataflow import DataflowStartFlexTemplateOperator
from airflow.operators import empty
from airflow.providers.google.cloud.hooks.gcs import GCSHook
from airflow.utils.task_group import TaskGroup
from airflow.providers.google.cloud.operators.bigquery import  BigQueryInsertJobOperator
from airflow.operators.empty import EmptyOperator
from airflow.operators.python import PythonOperator
from airflow.providers.google.cloud.sensors.dataform import DataformWorkflowInvocationStateSensor
from airflow.providers.google.cloud.operators.dataform import (
    DataformCreateCompilationResultOperator,
    DataformCreateWorkflowInvocationOperator,
)
from airflow.providers.google.cloud.operators.dataproc import (
    DataprocCreateBatchOperator,
    DataprocGetBatchOperator
)
from airflow.providers.google.cloud.sensors.dataproc import DataprocBatchSensor
from google.cloud.dataform_v1beta1 import WorkflowInvocation
from google.cloud import storage
from datetime import datetime, timedelta

# --------------------------------------------------------------------------------
# Read variables from GCS parameters file for the job
# --------------------------------------------------------------------------------
storage_client = storage.Client()
jobs_bucket = Variable.get("DATA_TRANSFORMATION_GCS_BUCKET")

batch_id = f"aef-{str(uuid.uuid4())}"

def extract_job_params(job_name, function_name, encoding='utf-8'):
    """Extracts parameters from a JSON job file.

    Args:
        bucket_name: Bucket containing the JSON parameters file .

    Returns:
        A dictionary containing the extracted parameters.
    """

    json_file_path = f'gs://{jobs_bucket}/{function_name}/{job_name}.json'

    parts = json_file_path.replace("gs://", "").split("/")
    bucket_name = parts[0]
    object_name = "/".join(parts[1:])
    bucket = storage_client.bucket(bucket_name)
    blob = bucket.blob(object_name)

    json_data = blob.download_as_bytes()
    params = json.loads(json_data.decode(encoding))
    return params


# --------------------------------------------------------------------------------
# Set default arguments
# --------------------------------------------------------------------------------

# If you are running Airflow in more than one time zone
# see https://airflow.apache.org/docs/apache-airflow/stable/timezone.html
# for best practices
yesterday = datetime.now() - timedelta(days=1)

default_args = {
    'owner': 'airflow',
    'start_date': yesterday,
    'depends_on_past': False,
    'email': [''],
    'email_on_failure': False,
    'email_on_retry': False,
    'retries': 0,
    'retry_delay': timedelta(minutes=5)
}

sample_serverless_spark_mainframe_ingestion = extract_job_params('sample_serverless_spark_mainframe_ingestion','dataproc-serverless-job-executor')
for key, value in sample_serverless_spark_mainframe_ingestion.items():
	default_args['sample_serverless_spark_mainframe_ingestion'+key] = value

sample_jdbc_dataflow_ingestion = extract_job_params('sample_jdbc_dataflow_ingestion','dataflow-flextemplate-job-executor')
for key, value in sample_jdbc_dataflow_ingestion.items():
	default_args['sample_jdbc_dataflow_ingestion'+key] = value

run_dataform_tag = extract_job_params('run_dataform_tag','dataform-tag-executor')
for key, value in run_dataform_tag.items():
	default_args['run_dataform_tag'+key] = value


start_date_str = yesterday.strftime('%Y-%m-%d')
end_date_str = datetime.today().strftime('%Y-%m-%d')
# --------------------------------------------------------------------------------
# Main DAG
# --------------------------------------------------------------------------------

with models.DAG(
        'demo_pipeline_composer',
        default_args=default_args,
        params={
            "start_date_str": start_date_str,
            "end_date_str": end_date_str
        },
        catchup=False,
        schedule_interval=None) as dag:

    start = empty.EmptyOperator(
        task_id='start',
        trigger_rule='all_success'
    )

        # Start level group definition
    with TaskGroup(group_id="Level_1") as tg_Level_1:
                   # Start thread group definition
           with TaskGroup(group_id="Level_1_Thread_1") as tg_level_1_Thread_1:
                    with TaskGroup(group_id="sample_serverless_spark_mainframe_ingestion") as sample_serverless_spark_mainframe_ingestion:

                           def push_batch_id_to_xcom(**context):
                               context['task_instance'].xcom_push(key='batch-id', value=batch_id)

                           push_batch_id_to_xcom_sample_serverless_spark_mainframe_ingestion = PythonOperator(
                               task_id='batch_id',
                               python_callable=push_batch_id_to_xcom,
                               provide_context=True
                           )

                           # create batch
                           create_batch_for_sample_serverless_spark_mainframe_ingestion = DataprocCreateBatchOperator(
                               task_id="create_batch_for_sample_serverless_spark_mainframe_ingestion",
                               batch={
                                   "spark_batch": {
                                       "jar_file_uris": [default_args['sample_serverless_spark_mainframe_ingestion' + 'jar_file_location']],
                                       "main_class": default_args['sample_serverless_spark_mainframe_ingestion' + 'spark_app_main_class'],
                                       "args": default_args['sample_serverless_spark_mainframe_ingestion' + 'spark_args'],
                                   },
                                   "runtime_config": {
                                       "version": default_args['sample_serverless_spark_mainframe_ingestion' + 'dataproc_serverless_runtime_version'],
                                       "properties": default_args['sample_serverless_spark_mainframe_ingestion' + 'spark_app_properties'],
                                   },
                                   "environment_config": {
                                       "execution_config": {
                                           "service_account": default_args['sample_serverless_spark_mainframe_ingestion' + 'dataproc_service_account'],
                                           "subnetwork_uri": f"projects/{default_args['sample_serverless_spark_mainframe_ingestion' + 'dataproc_serverless_project_id']}/{default_args['sample_serverless_spark_mainframe_ingestion' + 'subnetwork']}"
                                       }
                                   }
                               },
                               batch_id="{{ task_instance.xcom_pull(task_ids='Level_1.Level_1_Thread_1.sample_serverless_spark_mainframe_ingestion.batch_id', key='batch-id') }}",
                               project_id=default_args['sample_serverless_spark_mainframe_ingestion' + 'dataproc_serverless_project_id'],
                               region=default_args['sample_serverless_spark_mainframe_ingestion' + 'dataproc_serverless_region'],
                               deferrable=True
                           )

                           wait_for_batch_completion_for_sample_serverless_spark_mainframe_ingestion = DataprocBatchSensor(
                               task_id='wait_for_batch_completion_for_sample_serverless_spark_mainframe_ingestion',
                               batch_id="{{ task_instance.xcom_pull(task_ids='Level_1.Level_1_Thread_1.sample_serverless_spark_mainframe_ingestion.batch_id', key='batch-id') }}",
                               region=default_args['sample_serverless_spark_mainframe_ingestion' + 'dataproc_serverless_region'],
                               poke_interval=400,
                               timeout=3600,
                               soft_fail=True
                           )

                           get_batch_for_sample_serverless_spark_mainframe_ingestion = DataprocGetBatchOperator(
                               task_id="get_batch_for_sample_serverless_spark_mainframe_ingestion",
                               batch_id="{{ task_instance.xcom_pull(task_ids='Level_1.Level_1_Thread_1.sample_serverless_spark_mainframe_ingestion.batch_id', key='batch-id') }}",
                               region=default_args['sample_serverless_spark_mainframe_ingestion' + 'dataproc_serverless_region']
                           )

                           push_batch_id_to_xcom_sample_serverless_spark_mainframe_ingestion >> create_batch_for_sample_serverless_spark_mainframe_ingestion >> wait_for_batch_completion_for_sample_serverless_spark_mainframe_ingestion >> get_batch_for_sample_serverless_spark_mainframe_ingestion


                    sample_serverless_spark_mainframe_ingestion
           # End thread group definition           # Start thread group definition
           with TaskGroup(group_id="Level_1_Thread_2") as tg_level_1_Thread_2:
                    with TaskGroup(group_id="sample_jdbc_dataflow_ingestion") as sample_jdbc_dataflow_ingestion:
                        dataflow_job_name = re.sub(r"^\d+", "",re.sub(r"[^a-z0-9+]", "", "sample_jdbc_dataflow_ingestion"))
                        dataflow_job_name = re.sub(r"^\d+", "", dataflow_job_name)
                        gcs_path = "gs://dataflow-templates-{region}/{version}/flex/{template}".format(region=default_args['sample_jdbc_dataflow_ingestion'+'dataflow_location'],
                                                                                                       version=default_args['sample_jdbc_dataflow_ingestion'+'dataflow_template_version'],
                                                                                                       template=default_args['sample_jdbc_dataflow_ingestion'+'dataflow_template_name'])
                        body = {
                            "launchParameter": {
                                "jobName": dataflow_job_name,
                                "parameters": default_args['sample_jdbc_dataflow_ingestion'+'dataflow_job_params'],
                                "containerSpecGcsPath": gcs_path,
                                "environment": {
                                    "tempLocation": "gs://{bucket}/dataflow/temp".format(bucket=default_args['sample_jdbc_dataflow_ingestion' + 'dataflow_temp_bucket']),
                                    "maxWorkers": str(default_args['sample_jdbc_dataflow_ingestion' + 'dataflow_max_workers']),
                                    "network": str(default_args['sample_jdbc_dataflow_ingestion' + 'network']),
                                    "subnetwork": str(default_args['sample_jdbc_dataflow_ingestion' + 'subnetwork'])}
                            }
                        }
                        dataflow_job_sample_jdbc_dataflow_ingestion = DataflowStartFlexTemplateOperator(
                            task_id="dataflow_flex_template_sample_jdbc_dataflow_ingestion",
                            location=default_args['sample_jdbc_dataflow_ingestion'+'dataflow_location'],
                            body=body
                        )

                        dataflow_job_sample_jdbc_dataflow_ingestion


                    sample_jdbc_dataflow_ingestion
           # End thread group definition

           tg_level_1_Thread_1
           tg_level_1_Thread_2
    # End level group definition
    # Start level group definition
    with TaskGroup(group_id="Level_2") as tg_Level_2:
                   # Start thread group definition
           with TaskGroup(group_id="Level_2_Thread_3") as tg_level_2_Thread_3:
                    with TaskGroup(group_id="run_dataform_tag") as run_dataform_tag:
                           # compilation
                           create_compilation_result_for_run_dataform_tag = DataformCreateCompilationResultOperator(
                               project_id=default_args['run_dataform_tag'+'dataform_project_id'],
                               region=default_args['run_dataform_tag'+'dataform_location'],
                               repository_id=default_args['run_dataform_tag'+'repository_name'],
                               task_id="compilation_task_run_dataform_tag",
                               compilation_result={
                                   "git_commitish": default_args['run_dataform_tag'+'branch'],
                                   "code_compilation_config": {
                                       "vars": {
                                            "start_date": "{{ params.start_date_str }}",
                                            "end_date": "{{ params.end_date_str }}"
                                       }
                                   }
                               }
                           )

                           # workflow invocation in dataform
                           create_workflow_run_dataform_tag_invocation = DataformCreateWorkflowInvocationOperator(
                               project_id=default_args['run_dataform_tag' + 'dataform_project_id'],
                               region=default_args['run_dataform_tag' + 'dataform_location'],
                               repository_id=default_args['run_dataform_tag' + 'repository_name'],
                               task_id='workflow_inv_run_dataform_tag',
                               asynchronous=True,
                               workflow_invocation={
                                   "compilation_result": "{{ task_instance.xcom_pull('Level_2.Level_2_Thread_3.run_dataform_tag.compilation_task_run_dataform_tag')['name'] }}",
                                   "invocation_config": { "included_tags": default_args['run_dataform_tag'+'tags'],
                                                          "transitive_dependencies_included": True
                                                        }
                               },
                               trigger_rule='all_success'
                           )

                           is_workflow_run_dataform_tag_invocation_done = DataformWorkflowInvocationStateSensor(
                               project_id=default_args['run_dataform_tag' + 'dataform_project_id'],
                               region=default_args['run_dataform_tag' + 'dataform_location'],
                               repository_id=default_args['run_dataform_tag' + 'repository_name'],
                               task_id="is_workflow_run_dataform_tag_invocation_done",
                               workflow_invocation_id=("{{ task_instance.xcom_pull('Level_2.Level_2_Thread_3.run_dataform_tag.workflow_inv_run_dataform_tag')['name'].split('/')[-1] }}"),
                               expected_statuses={WorkflowInvocation.State.SUCCEEDED},
                               failure_statuses={WorkflowInvocation.State.FAILED, WorkflowInvocation.State.CANCELLED},
                           )

                           create_compilation_result_for_run_dataform_tag >> create_workflow_run_dataform_tag_invocation >> is_workflow_run_dataform_tag_invocation_done


                    run_dataform_tag
           # End thread group definition

           tg_level_2_Thread_3
    # End level group definition


    end = empty.EmptyOperator(
        task_id='end',
        trigger_rule='all_success'
    )

    start >> tg_Level_1 >> tg_Level_2 >> end