cloud-composer/dags/sample-create-data-fusion.py

####################################################################################
# Copyright 2022 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.
####################################################################################

# Author:  Adam Paternostro
# Summary: Creates a data fusion instance and waits for the instance to come online ()

# [START dag]
from datetime import datetime, timedelta
from airflow.operators import bash_operator
from airflow.utils import trigger_rule
from airflow.operators.python_operator import PythonOperator
import requests
import sys
import os
import logging
import airflow
import google.auth
import google.auth.transport.requests
import time




default_args = {
    'owner': 'airflow',
    'depends_on_past': False,
    'email': None,
    'email_on_failure': False,
    'email_on_retry': False,
    'retries': 0,
    'retry_delay': timedelta(minutes=5),
    'dagrun_timeout' : timedelta(minutes=60),
}

project_id            = os.environ['ENV_PROJECT_ID'] 
region                = os.environ['ENV_DATAFUSION_REGION'] 
datafusion_name       = "data-fusion-prod-01"


# Creates a data fusion instance (name hardcoded)
# NOTE: There is a data fusion airflow operator, this is designed to show the REST API call for other tooling as well
def create_data_fusion(project_id, region, datafusion_name):

    # Get auth (default service account running composer worker node)
    creds, project = google.auth.default()
    auth_req = google.auth.transport.requests.Request() # required to acess access token
    creds.refresh(auth_req)
    access_token=creds.token
    auth_header = { 'Authorization'   : "Bearer " + access_token ,
                    'Accept'          : 'application/json',
                    'Content-Type'    : 'application/json'
    }

    # call rest api with bearer token
    createUri="https://datafusion.googleapis.com/v1beta1/projects/{}/locations/{}/instances?instanceId={}".format(project_id,region,datafusion_name)
    # https://cloud.google.com/data-fusion/docs/reference/rest/v1beta1/projects.locations.instances#Instance
    request_body= '{ "labels" : { "env" : "prod" } }'


    try:
        response = requests.post(createUri, headers=auth_header, data=request_body)
        response.raise_for_status()
        print("Create Data Fusion")
    except requests.exceptions.RequestException as err:
        print(err)
        raise err


# Waits for the instance to come online
def wait_for_data_fusion_provisioning(project_id, region, datafusion_name):

    # Get auth (default service account running composer worker node)
    creds, project = google.auth.default()
    auth_req = google.auth.transport.requests.Request() # required to acess access token
    creds.refresh(auth_req)
    access_token=creds.token
    auth_header = {'Authorization' : "Bearer " + access_token }


    instanceState = "CREATING"
    i = 1

    while instanceState == "CREATING":

        time.sleep(30)
            
        # Check the state
        stateDataFusion="https://datafusion.googleapis.com/v1beta1/projects/{}/locations/{}/instances/{}".format(project_id,region,datafusion_name)

        try:
            response = requests.get(stateDataFusion, headers=auth_header)
            response.raise_for_status()
            print("Checking State of Data Fusion Deployment")
        except requests.exceptions.RequestException as err:
            print(err)
            raise err

        instanceState = response.json()['state']

        print("Checking Data Fusion State ({}) value: {}".format(i, instanceState))
        print("JSON: ", response.json())
        i = i + 1
        
        """ Sample Response
        {
            "name": "projects/data-analytics-demo-z4x77nbcc3/locations/us-central1/instances/qadatafusion",
            "type": "BASIC",
            "networkConfig": {},
            "createTime": "2022-06-29T13:21:24.874523833Z",
            "updateTime": "2022-06-29T13:32:28.762387202Z",
            "state": "RUNNING",
            "serviceEndpoint": "https://qadatafusion-data-analytics-demo-z4x77nbcc3-dot-usw1.datafusion.googleusercontent.com",
            "version": "6.6.0",
            "serviceAccount": "cloud-datafusion-management-sa@wcb6888d249fbda6c-tp.iam.gserviceaccount.com",
            "displayName": "QADataFusion",
            "availableVersion": [
                {
                "versionNumber": "6.7.0"
                }
            ],
            "apiEndpoint": "https://qadatafusion-data-analytics-demo-z4x77nbcc3-dot-usw1.datafusion.googleusercontent.com/api",
            "gcsBucket": "gs://df-2432185209294608530-52ppjphxvyi6zaykaizbbqaaaa",
            "p4ServiceAccount": "service-1074332558384@gcp-sa-datafusion.iam.gserviceaccount.com",
            "tenantProjectId": "wcb6888d240fbda7c-tp",
            "dataprocServiceAccount": "1074332558384-compute@developer.gserviceaccount.com"
        }
        """


     
with airflow.DAG('sample-create-data-fusion',
                 default_args=default_args,
                 start_date=datetime(2021, 1, 1),
                 # Add the Composer "Data" directory which will hold the SQL scripts for deployment
                 template_searchpath=['/home/airflow/gcs/data'],
                 # Not scheduled, trigger only
                 schedule_interval=None) as dag:

    # NOTE: The service account of the Composer worker node must have access to run these commands

    # Show creating data fusion via the REST API
    create_data_fusion = PythonOperator(
        task_id='create_data_fusion',
        python_callable= create_data_fusion,
        op_kwargs = { "project_id" : project_id, 
                      "region" : region,
                      "datafusion_name" : datafusion_name,
                       },
        execution_timeout=timedelta(minutes=10),
        dag=dag,
        )    

    # Wait for the instance to come online
    wait_for_data_fusion_provisioning = PythonOperator(
        task_id='wait_for_data_fusion_provisioning',
        python_callable= wait_for_data_fusion_provisioning,
        op_kwargs = { "project_id" : project_id, 
                      "region" : region,
                      "datafusion_name" : datafusion_name,
                       },
        execution_timeout=timedelta(minutes=60),
        dag=dag,
        )    

    # DAG Graph
    create_data_fusion >> wait_for_data_fusion_provisioning
    

# [END dag]