Python Code in Airflow Dag

Question

I have a python code in Airflow Dag. This Dag performs 3 tasks:

1. Authenticate the user and get access token
2. Create a Databricks cluster using rest API and Submit a notebook job on a cluster using rest API. Pass access token created in the first step as input.
3. Check the status of notebook job

Please help me with code review for this Airflow Dag.

from airflow.operators.python_operator import PythonOperator
from airflow.models import DAG
from datetime import datetime,timedelta
import json
import os
import requests
import time
import logging


auth_file = "my-token.txt"
idds_cluster_id_file = "my-cluster-id.txt"
path = "/airflow/tokens/"
cluster_id = "0408-000631"
path= "/airflow/tokens/"
temp_run_id_file = "my_save_run_id.txt"
databeicks_token_file= "databricks-token.txt"
auth_url = "https://snapdeal.com/authentication"
cluster_url = "https://snapdeal.com/clustermanager"
exec_url = "https://snapdeal.com/databricksexecutor"

args = {
    'owner': "airflow",
    'start_date': datetime.now()
}
dag = DAG(
    dag_id="my-notebook-job-trigger-dev",
    default_args=args,
    max_active_runs=1,
    concurrency=1
)
status_dict = {
    "action": "get_job_status"
    }
auth_dict = {
    "refresh_token": ""
}

cluster_dict = {
    "payload": {
        "cluster_id": cluster_id
    },
    "databricks_token": ""
}

cluster_status_dict = {
    "payload": {
        "cluster_id": cluster_id
    },
    "databricks_token": ""
}

search_dict= {
  "action": "run_job",
  "auth_token": "",
  "payload": {
    "job_id": "956"
  }
 }

def write_json(response):
    with open(path + auth_file, 'wb') as file:  
        file.write(json.dumps(response["result"]))


def read_tokens():
    with open(path + auth_file) as f:
        data = json.load(f)
    return data


def get_authentication():
    data = read_tokens()
    auth_dict["refresh_token"] = data["refresh_token"]
    auth_response_raw = requests.post(url=auth_url, json=auth_dict, headers={"Content-Type": "application/json"})
    if auth_response_raw.status_code == 200:
        auth_response = auth_response_raw.json()
        if auth_response["status"] == "FAILED":
            logging.info(json.dumps(auth_response_raw.json()))
            status_message = auth_response["error"]["message"]
            raise Exception(status_message)
        else:
            write_json(auth_response)
            logging.info("Successfully got authentication.")
    else:
        logging.info(json.dumps(auth_response_raw.json()))
        raise Exception("Got a bad response.")


def read_databricks_token():
    with open(path + databeicks_token_file, 'r') as infile:
        token = infile.read()
    return token.rstrip()

def check_cluster_status(status):
    data = read_tokens()
    access_token = data["token_type"] + " " + data["access_token"]
    cluster_status_dict["operation"] = "getinfo"
    cluster_status_dict["databricks_token"] = read_databricks_token().rstrip()

    cluster_response_raw = requests.post(url=cluster_url, json=cluster_status_dict,
                                         headers={"Content-Type": "application/json", "Authorization": access_token})
    if cluster_response_raw.status_code == 200:
        cluster_response = cluster_response_raw.json()
        if cluster_response["status"] == "FAILED":
            logging.info(json.dumps(cluster_response_raw.json()))
            status_message = cluster_response["error"]["message"]
            raise Exception(status_message)
        else:
            while cluster_response["result"]["info"]["state"] != status:
                logging.info("Waiting for Cluster to be in "+ status + " status")
                time.sleep(30)
                cluster_response_raw = requests.post(url=cluster_url, json=cluster_status_dict,
                                                     headers={"Content-Type": "application/json",
                                                              "Authorization": access_token})
                if cluster_response_raw.status_code != 200:
                    logging.info(json.dumps(cluster_response_raw.json()))
                    raise Exception("Got a bad response.")
                else:
                    cluster_response = cluster_response_raw.json()

            logging.info("Cluster is " + status + " now")
    else:
        logging.info(json.dumps(cluster_response_raw.json()))
        raise Exception("Got a bad response.")


def create_cluster():
    data = read_tokens()
    access_token = data["token_type"] + " " + data["access_token"]
    cluster_dict["databricks_token"] = read_databricks_token()
    cluster_response_raw = requests.post(url=cluster_url, json=cluster_dict,
                                         headers={"Content-Type": "application/json", "Authorization": access_token})
    if cluster_response_raw.status_code == 200:
        cluster_response = cluster_response_raw.json()
        if cluster_response["status"] == "FAILED":
            logging.info(json.dumps(cluster_response_raw.json()))
            status_message = cluster_response["error"]["message"]
            raise Exception(status_message)
        else:
            logging.info("Cluster Creation Started")
    else:
        logging.info(json.dumps(cluster_response_raw.json()))
        raise Exception("Got a bad response.")
    cluster_status_dict["payload"]["cluster_id"] = cluster_response["result"]["cluster_id"]
    with open(path + idds_cluster_id_file, 'wb') as file:
        file.write(cluster_response["result"]["cluster_id"])
    check_cluster_status("RUNNING")


def submit_job():
    auth = read_tokens()
    access_token = auth["token_type"] + " " + auth["access_token"]
    create_cluster()
    search_dict["auth_token"] = read_databricks_token()
    job_response_raw = requests.post(url=exec_url, json=search_dict,
                                     headers={"Content-Type": "application/json", "Authorization": access_token})
    if job_response_raw.status_code == 200:
        job_response = job_response_raw.json()
        if job_response["status"] == "FAILED":
            logging.error(json.dumps(job_response_raw.json()))
            status_message = job_response["error"]["message"]
            raise Exception(status_message)
        else:
            logging.info("Starting to write run id in file")
            run_id = str(job_response["result"]["run_id"])
            with open(path + temp_run_id_file, 'w') as file:
                file.write(run_id)
            return run_id
    else:
        logging.error("Failed to execute Job")
        logging.error(json.dumps(job_response_raw.json()))
        raise Exception("Got a bad response.")


def delete_cluster():
    data = read_tokens()
    access_token = data["token_type"] + " " + data["access_token"]
    cluster_status_dict["operation"] = "delete"
    cluster_status_dict["databricks_token"] = read_databricks_token().rstrip()
    with open(path + idds_cluster_id_file, 'r') as file:
        cluster_id = file.read()
    cluster_status_dict["payload"]["cluster_id"] = cluster_id
    cluster_response_raw = requests.post(url=cluster_url, json=cluster_status_dict,
                                         headers={"Content-Type": "application/json", "Authorization": access_token})
    if cluster_response_raw.status_code == 200:
        cluster_response = cluster_response_raw.json()
        if cluster_response["status"] == "FAILED":
            logging.info(json.dumps(cluster_response_raw.json()))
            status_message = cluster_response["error"]["message"]
            raise Exception(status_message + " " + cluster_status_dict["payload"]["cluster_id"])
        else:
            logging.info("Deleting Cluster")
    else:
        logging.info(json.dumps(cluster_response_raw.json()))
        raise Exception("Got a bad response.")

def check_status():
    with open(path + temp_run_id_file, 'r') as file:
        RUN_ID = file.read()
    os.remove(path + temp_run_id_file)
    logging.info(RUN_ID)
    status_dict["run_id"] = RUN_ID
    status_dict["auth_token"] = read_databricks_token()
    logging.info(status_dict)
    auth = read_tokens()
    access_token = auth["token_type"] + " " + auth["access_token"]
    job_status_response_raw = requests.post(url=exec_url, json=status_dict,
                                            headers={"Content-Type": "application/json", "Authorization": access_token})
    if job_status_response_raw.status_code == 200:
        job_status_response = job_status_response_raw.json()
        logging.info(job_status_response)
    else:
        logging.info(json.dumps(job_status_response_raw.json()))
        raise Exception("Got a bad response.")

    while job_status_response["result"]["state"]["life_cycle_state"] == "PENDING" or job_status_response["result"]["state"]["life_cycle_state"] == "RUNNING":
        logging.info("Waiting for job to finish!")
        time.sleep(30)
        job_status_response_raw = requests.post(url=exec_url, json=status_dict,
                                                headers={"Content-Type": "application/json",
                                                         "Authorization": access_token})
        if job_status_response_raw.status_code == 200:
            job_status_response = job_status_response_raw.json()
            logging.info(job_status_response)
        else:
            raise Exception("Got a bad response.")
    if "result_state" in job_status_response["result"]["state"]:
        if job_status_response["result"]["state"]["result_state"] == "FAILED":
            logging.info(json.dumps(job_status_response_raw.json()))
            status_message = job_status_response["result"]["state"]["state_message"]
            logging.info("Failed with error: " + status_message)
            logging.info(job_status_response)
            delete_cluster()
            check_cluster_status("TERMINATED")
            raise Exception(status_message)
        elif job_status_response["result"]["state"]["result_state"] == "SUCCESS":
            logging.info("Job ran successfully")
    logging.info(job_status_response)
    delete_cluster()
    check_cluster_status("TERMINATED")

t1 = PythonOperator(
    task_id='get_authentication',
    python_callable=get_authentication,
    dag=dag
)


t2 = PythonOperator(
    task_id='run_notebook_job',
    python_callable=submit_job,
    dag=dag
)
t3 = PythonOperator(
    task_id='check_notebook_status',
    python_callable=check_status,
    dag=dag
)

t1.set_downstream(t2)
t2.set_downstream(t3)

Answer 1

Typo

databeicks_token_file -> databricks_token_file

Exceptions

        raise Exception(status_message)

is going to make it problematic for callers to meaningfully catch this separate from other exceptions, if they need to. Instead, raise a more specific exception, possibly a custom one - they're easy to define in Python.

Global code

t1 = PythonOperator(
    task_id='get_authentication',
    python_callable=get_authentication,
    dag=dag
)


t2 = PythonOperator(
    task_id='run_notebook_job',
    python_callable=submit_job,
    dag=dag
)
t3 = PythonOperator(
    task_id='check_notebook_status',
    python_callable=check_status,
    dag=dag
)

t1.set_downstream(t2)
t2.set_downstream(t3)

should be pulled into a main function.

Global variables

These:

auth_file = "my-token.txt"
idds_cluster_id_file = "my-cluster-id.txt"
path = "/airflow/tokens/"
cluster_id = "0408-000631"
path= "/airflow/tokens/"
temp_run_id_file = "my_save_run_id.txt"
databeicks_token_file= "databricks-token.txt"
auth_url = "https://snapdeal.com/authentication"
cluster_url = "https://snapdeal.com/clustermanager"
exec_url = "https://snapdeal.com/databricksexecutor"

are fine-ish where they are, though the names should be capitalized. These might also benefit from being pulled out into a configuration file or environment variables.

However, the rest of it (args through search_dict) probably doesn't belong here. args should just be moved to a nested literal initializer inside the DAG constructor; these other dictionaries should be passed - or parts of them passed, where appropriate - between functions. Having various functions mutate various keys in these dictionaries before they're passed to requests is not very maintainable. Something like check_cluster_status is better off constructing cluster_status_dict in its own scope.