Implement self sufficient DAG for LivyOperatorAsync (#235)

Author: sunank200

.circleci/integration-tests/master_dag.py

@@ -138,6 +138,11 @@ def prepare_dag_dependency(task_info, execution_time):
     dag_run_ids.extend(ids)
     chain(*snowflake_trigger_tasks)
 
+    livy_task_info = [{"livy_dag": "example_livy_operator"}]
+    livy_trigger_tasks, ids = prepare_dag_dependency(livy_task_info, "{{ ds }}")
+    dag_run_ids.extend(ids)
+    chain(*livy_trigger_tasks)
+
     report = PythonOperator(
         task_id="get_report",
         python_callable=get_report,
@@ -158,6 +163,7 @@ def prepare_dag_dependency(task_info, execution_time):
         databricks_trigger_tasks[0],
         http_trigger_tasks[0],
         snowflake_trigger_tasks[0],
+        livy_trigger_tasks[0],
     ]
 
     last_task = [
@@ -168,6 +174,7 @@ def prepare_dag_dependency(task_info, execution_time):
         databricks_trigger_tasks[-1],
         http_trigger_tasks[-1],
         snowflake_trigger_tasks[-1],
+        livy_trigger_tasks[-1],
     ]
 
     last_task >> end

astronomer/providers/apache/livy/example_dags/example_livy.py

@@ -3,38 +3,321 @@
 The tasks below trigger the computation of pi on the Spark instance
 using the Java and Python executables provided in the example library.
 """
+import logging
 import os
+import time
 from datetime import datetime
+from typing import Any, List
 
-from airflow import DAG
+import boto3
+import paramiko
+from airflow import DAG, settings
+from airflow.models import Connection, Variable
+from airflow.operators.python import PythonOperator
+from airflow.providers.amazon.aws.operators.emr import (
+    EmrCreateJobFlowOperator,
+    EmrTerminateJobFlowOperator,
+)
+from requests import get
 
 from astronomer.providers.apache.livy.operators.livy import LivyOperatorAsync
 
 LIVY_JAVA_FILE = os.environ.get("LIVY_JAVA_FILE", "/spark-examples.jar")
 LIVY_PYTHON_FILE = os.environ.get("LIVY_PYTHON_FILE", "/user/hadoop/pi.py")
+JOB_FLOW_ROLE = os.environ.get("EMR_JOB_FLOW_ROLE", "EMR_EC2_DefaultRole")
+SERVICE_ROLE = os.environ.get("EMR_SERVICE_ROLE", "EMR_DefaultRole")
+PEM_FILENAME = os.environ.get("PEM_FILENAME", "providers_team_keypair")
+PRIVATE_KEY = Variable.get("providers_team_keypair")
+
+AWS_S3_CREDS = {
+    "aws_access_key_id": os.environ.get("AWS_ACCESS_KEY", "sample_aws_access_key_id"),
+    "aws_secret_access_key": os.environ.get("AWS_SECRET_KEY", "sample_aws_secret_access_key"),
+    "region_name": os.environ.get("AWS_REGION_NAME", "us-east-2"),
+}
+
+COMMAND_TO_CREATE_PI_FILE: List[str] = [
+    "curl https://raw.githubusercontent.com/apache/spark/master/examples/src/main/python/pi.py >> pi.py",
+    "hadoop fs -copyFromLocal pi.py /user/hadoop",
+]
+
+JOB_FLOW_OVERRIDES = {
+    "Name": "team-provider-example-dag-livy-test",
+    "ReleaseLabel": "emr-5.35.0",
+    "Applications": [
+        {"Name": "Spark"},
+        {
+            "Name": "Livy",
+        },
+        {
+            "Name": "Hive",
+        },
+        {
+            "Name": "Hadoop",
+        },
+    ],
+    "Instances": {
+        "InstanceGroups": [
+            {
+                "Name": "Primary node",
+                "Market": "ON_DEMAND",
+                "InstanceRole": "MASTER",
+                "InstanceType": "m4.large",
+                "InstanceCount": 1,
+            },
+        ],
+        "Ec2KeyName": PEM_FILENAME,
+        "KeepJobFlowAliveWhenNoSteps": True,
+        "TerminationProtected": False,
+    },
+    "Steps": [],
+    "JobFlowRole": JOB_FLOW_ROLE,
+    "ServiceRole": SERVICE_ROLE,
+}
+
+
+def create_airflow_connection(task_instance: Any) -> None:
+    """
+    Checks if airflow connection exists, if yes then deletes it.
+    Then, create a new livy_default connection.
+    """
+    conn = Connection(
+        conn_id="livy_default",
+        conn_type="livy",
+        host=task_instance.xcom_pull(
+            key="cluster_response_master_public_dns", task_ids=["describe_created_cluster"]
+        )[0],
+        login="",
+        password="",
+        port=8998,
+    )  # create a connection object
+
+    session = settings.Session()
+    connection = session.query(Connection).filter_by(conn_id=conn.conn_id).one_or_none()
+    if connection is None:
+        logging.info("Connection %s doesn't exist.", str(conn.conn_id))
+    else:
+        session.delete(connection)
+        session.commit()
+        logging.info("Connection %s deleted.", str(conn.conn_id))
+
+    session.add(conn)
+    session.commit()  # it will insert the connection object programmatically.
+    logging.info("Connection livy_default is created")
+
+
+def add_inbound_rule_for_security_group(task_instance: Any) -> None:
+    """
+    Sets the inbound rule for the aws security group, based on
+    current ip address of the system.
+    """
+    current_docker_ip = get("https://api.ipify.org").text
+    logging.info("Current ip address is: %s", str(current_docker_ip))
+    client = boto3.client("ec2", **AWS_S3_CREDS)
+
+    response = client.describe_security_groups(
+        GroupIds=[
+            task_instance.xcom_pull(
+                key="cluster_response_master_security_group", task_ids=["describe_created_cluster"]
+            )[0]
+        ]
+    )
+    current_ip_permissions = response["SecurityGroups"][0]["IpPermissions"]
+    ip_exists = False
+    for current_ip in current_ip_permissions:
+        ip_ranges = current_ip["IpRanges"]
+        for ip in ip_ranges:
+            if ip["CidrIp"] == str(current_docker_ip) + "/32":
+                ip_exists = True
+
+    if not ip_exists:
+        # open port for port 8998
+        client.authorize_security_group_ingress(
+            GroupId=task_instance.xcom_pull(
+                key="cluster_response_master_security_group", task_ids=["describe_created_cluster"]
+            )[0],
+            IpPermissions=[
+                {
+                    "IpProtocol": "tcp",
+                    "FromPort": 8998,
+                    "ToPort": 8998,
+                    "IpRanges": [{"CidrIp": str(current_docker_ip) + "/32"}],
+                }
+            ],
+        )
+
+        # open port for port 22 for ssh and copy file for hdfs
+        client.authorize_security_group_ingress(
+            GroupId=task_instance.xcom_pull(
+                key="cluster_response_master_security_group", task_ids=["describe_created_cluster"]
+            )[0],
+            IpPermissions=[
+                {
+                    "IpProtocol": "tcp",
+                    "FromPort": 22,
+                    "ToPort": 22,
+                    "IpRanges": [{"CidrIp": str(current_docker_ip) + "/32"}],
+                }
+            ],
+        )
+
+
+def create_key_pair() -> None:
+    """
+    Load the private_key from airflow variable and creates a pem_file
+    at /tmp/.
+    """
+    # remove the file if it exists
+    if os.path.exists(f"/tmp/{PEM_FILENAME}.pem"):
+        os.remove(f"/tmp/{PEM_FILENAME}.pem")
+
+    # read the content for pem file from Variable set on Airflow UI.
+    with open(f"/tmp/{PEM_FILENAME}.pem", "w+") as fh:
+        fh.write(PRIVATE_KEY)
+
+    # write private key to file with 400 permissions
+    os.chmod(f"/tmp/{PEM_FILENAME}.pem", 0o400)
+
+
+def ssh_and_run_command(task_instance: Any, **kwargs: Any) -> None:
+    """
+    SSH into the machine and execute the bash script from the list
+    of commands.
+    """
+    key = paramiko.RSAKey.from_private_key_file(kwargs["path_to_pem_file"])
+    client = paramiko.SSHClient()
+    client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+    # Connect/ssh to an instance
+    cluster_response_master_public_dns = task_instance.xcom_pull(
+        key="cluster_response_master_public_dns", task_ids=["describe_created_cluster"]
+    )[0]
+    try:
+        client.connect(hostname=cluster_response_master_public_dns, username=kwargs["username"], pkey=key)
+
+        # Execute a command(cmd) after connecting/ssh to an instance
+        for command in kwargs["command"]:
+            stdin, stdout, stderr = client.exec_command(command)
+            stdout.read()
+
+        # close the client connection once the job is done
+        client.close()
+    except Exception as exc:
+        raise Exception("Got an exception as %s.", str(exc))
+
+
+def get_cluster_details(task_instance: Any) -> None:
+    """
+    Fetches the cluster details and stores EmrManagedMasterSecurityGroup and
+    MasterPublicDnsName in the XCOM.
+    """
+    client = boto3.client("emr", **AWS_S3_CREDS)
+    response = client.describe_cluster(
+        ClusterId=str(task_instance.xcom_pull(key="return_value", task_ids=["cluster_creator"])[0])
+    )
+    while (
+        "MasterPublicDnsName" not in response["Cluster"]
+        and response["Cluster"]["Status"]["State"] != "WAITING"
+    ):
+        logging.info("wait for 11 minutes to get the MasterPublicDnsName")
+        time.sleep(660)
+        response = client.describe_cluster(
+            ClusterId=str(task_instance.xcom_pull(key="return_value", task_ids=["cluster_creator"])[0])
+        )
+        logging.info("current response from ams emr: %s", str(response))
+    task_instance.xcom_push(
+        key="cluster_response_master_public_dns", value=response["Cluster"]["MasterPublicDnsName"]
+    )
+    task_instance.xcom_push(
+        key="cluster_response_master_security_group",
+        value=response["Cluster"]["Ec2InstanceAttributes"]["EmrManagedMasterSecurityGroup"],
+    )
+
 
 with DAG(
     dag_id="example_livy_operator",
-    default_args={"args": [10]},
     schedule_interval=None,
     start_date=datetime(2021, 1, 1),
     catchup=False,
     tags=["example", "async", "deferrable", "LivyOperatorAsync"],
 ) as dag:
+    # [START howto_create_key_pair_file]
+    create_key_pair_file = PythonOperator(
+        task_id="create_key_pair_file",
+        python_callable=create_key_pair,
+    )
+    # [END howto_create_key_pair_file]
+
+    # [START howto_operator_emr_create_job_flow]
+    cluster_creator = EmrCreateJobFlowOperator(
+        task_id="cluster_creator",
+        job_flow_overrides=JOB_FLOW_OVERRIDES,
+    )
+    # [END howto_operator_emr_create_job_flow]
+
+    # [START describe_created_cluster]
+    describe_created_cluster = PythonOperator(
+        task_id="describe_created_cluster", python_callable=get_cluster_details
+    )
+    # [END describe_created_cluster]
 
-    # [START create_livy]
+    # [START add_example_pi_file_in_hdfs]
+    ssh_and_copy_pifile_to_hdfs = PythonOperator(
+        task_id="ssh_and_copy_pifile_to_hdfs",
+        python_callable=ssh_and_run_command,
+        op_kwargs={
+            "path_to_pem_file": f"/tmp/{PEM_FILENAME}.pem",
+            "username": "hadoop",
+            "command": COMMAND_TO_CREATE_PI_FILE,
+        },
+    )
+    # [END add_example_pi_file_in_hdfs]
+
+    # [START add_ip_address_for_inbound_rules]
+    get_and_add_ip_address_for_inbound_rules = PythonOperator(
+        task_id="get_and_add_ip_address_for_inbound_rules",
+        python_callable=add_inbound_rule_for_security_group,
+    )
+    # [END add_ip_address_for_inbound_rules]
+
+    # [START create_airflow_connection_for_livy]
+    create_airflow_connection_for_livy = PythonOperator(
+        task_id="create_airflow_connection_for_livy", python_callable=create_airflow_connection
+    )
+    # [END create_airflow_connection_for_livy]
+
+    # [START run_pi_example_without_polling_interval]
     livy_java_task = LivyOperatorAsync(
-        task_id="pi_java_task",
+        task_id="livy_java_task",
         file=LIVY_JAVA_FILE,
         num_executors=1,
         conf={
             "spark.shuffle.compress": "false",
         },
         class_name="org.apache.spark.examples.SparkPi",
-        polling_interval=0,
     )
+    # [END run_pi_spark_without_polling_interval]
 
-    livy_python_task = LivyOperatorAsync(task_id="pi_python_task", file=LIVY_PYTHON_FILE, polling_interval=30)
+    # [START run_py_example_with_polling_interval]
+    livy_python_task = LivyOperatorAsync(
+        task_id="livy_python_task", file=LIVY_PYTHON_FILE, polling_interval=30
+    )
+    # [END run_py_example_with_polling_interval]
 
-    livy_java_task >> livy_python_task
-    # [END create_livy]
+    # [START howto_operator_emr_terminate_job_flow]
+    remove_cluster = EmrTerminateJobFlowOperator(
+        task_id="remove_cluster",
+        job_flow_id=cluster_creator.output,
+        trigger_rule="all_done",
+    )
+    # [END howto_operator_emr_terminate_job_flow]
+
+    (
+        create_key_pair_file
+        >> cluster_creator
+        >> describe_created_cluster
+        >> get_and_add_ip_address_for_inbound_rules
+        >> ssh_and_copy_pifile_to_hdfs
+        >> create_airflow_connection_for_livy
+        >> livy_java_task
+        >> livy_python_task
+        >> remove_cluster
+    )

setup.cfg

@@ -69,6 +69,7 @@ snowflake =
     apache-airflow-providers-snowflake
 apache.livy =
     apache-airflow-providers-apache-livy
+    paramiko
 docs =
     sphinx
     sphinx-autoapi
@@ -117,6 +118,7 @@ all =
     gcloud-aio-storage
     google-api-core>=1.25.1,<2.0.0
     kubernetes_asyncio
+    paramiko
 
 [options.packages.find]
 include =