The Power of a Data Enablers Team: A Bridge between Business and Platform

from airflow.exceptions import AirflowException
from airflow.models import Variable, Connection
from datetime import datetime
from airflow.providers.databricks.operators.databricks import DatabricksSubmitRunOperator
from airflow.providers.databricks.hooks.databricks import DatabricksHook
from airflow.providers.postgres.hooks.postgres import PostgresHook
from typing import Dict, List, Optional
import requests


DEFAULT_DATABRICKS_NODE_TYPE = 'm5.xlarge'
DEFAULT_POLLING_PERIOD_SECONDS = 15
DEFAULT_SPARK_MAX_EXECUTORS = 2
DEFAULT_SPARK_CONF = {"spark.pyspark.python": "python3"}
DEFAULT_DATABRICKS_CLUSTER = {
            'aws_attributes': {
                        'first_on_demand': 1
            },
            'autoscale': {
                        'min_workers': 1,
                        'max_workers': 3
            }
}


class DatabricksPysparkOperator(DatabricksSubmitRunOperator):
    """
    Runs Pyspark applications from repo:
    https://github.mpi-internal.com/scmspain/backend-di--etl-pyspark-databricks

    :param team: Name of the team in "es_di_team" nomenclature.
        Example:
        team="es_di_advertising"
    :type team: str
    :param object_path: Path of the module or class to run on Databricks
        Example:
        object_path="advertising.domain_fact_tables.DomainFactTables"
    :type object_path: str
    :param execution_params: List of parameters used in the execution of the module or class
        Example:
        execution_params=[f"date_start_iso={date}", f"date_end_iso={date}"]
    :type execution_params: List[str]
    :param name: Optional application name that will take the application on Databricks. 
        If not set, the name will be taken from object_path parameter.
    :type name: Optional[str]
    :param databricks_node_type: Node type used by Spark workers and driver.
        Default value m5.xlarge
    :type databricks_node_type: Optional[str]
    :param databricks_cluster_conf: Dict with cluster configuration, with autoscale workers for example.
        Default value DEFAULT_DATABRICKS_CLUSTER variable in this operator
    :type databricks_cluster_conf: Optional[Dict]
    :param polling_interval: Time in seconds between polling for job completion.
        Default value 15
    :type Optional[int]
    :param max_executors: Number of maximum executors that will run the application.
        Default value 2
    :type Optional[int] 
    :param spark_conf: Optional dictionary with spark configuration that will apply to Databricks execution
        Example:
        spark_conf={"spark.driver.memory":"3g",
                    "spark.executor.memory": "3g"}
    :type spark_conf: Optional[str]
    :param pre_action_conn_id: Redshift connection where to launch the pre_action.
    :type Optional[str] 
    :param pre_action: SQL string to launch on Redshift before the application is launched to Databricks.
    :type Optional[str]     
    """
     
    template_ext = (".sql",)
    template_fields = ("json", "pre_action")
    template_fields_renderers = {'pre_action': 'sql'}

    def __init__(
            self,
            team: str,
            object_path: str,
            execution_params: List[str] = None,            
            name: str = None,
            databricks_node_type: str = DEFAULT_DATABRICKS_NODE_TYPE,
            databricks_cluster_conf: Dict = None,
            databricks_connection_id_con_name: Optional[str] = 'databricks_srv_airflow',
            databricks_policy_id_var_name: Optional[str] = 'databricks_srv_airflow_policy_id',
            datarbicks_group_var_name: Optional[str] = 'databricks_group_name',
            render_params: Optional[bool] = False,
            view_access_control_list=[],
            polling_interval: int = DEFAULT_POLLING_PERIOD_SECONDS,
            max_executors: int = None,
            spark_conf: Dict = None,
            pre_action_conn_id: Optional[str] = None,
            pre_action: Optional[str] = None,
            **kwargs
    ):
        # Get S3 path variables
        bucket = Variable.get('es_di_enablers/sch_repos_bucket')
        code_folder = Variable.get(
            'es_di_enablers/etl_pyspark_databricks_code_folder')
        applications = Variable.get(
            'es_di_enablers/etl_pyspark_applications_zip')
        s3_execution_path = f"s3://{bucket}/{code_folder}"

        # Set application name
        name = name or object_path.replace(".", "_")
        name = f"{name}_{datetime.now().isoformat()}"

        # Get Databricks credentials
        databricks_connection_id = f'{team}/{databricks_connection_id_con_name}'
        databricks_policy_id = Variable.get(
            f"{team}/{databricks_policy_id_var_name}")

        access_control_list = [
            {
                'group_name': Variable.get(f"{team}/{datarbicks_group_var_name}"),
                'permission_level': 'CAN_MANAGE'
            }
        ]

        self.view_access_control_list = view_access_control_list

        extra_access_control_list = [{
                'group_name': group,
                'permission_level': 'CAN_VIEW'
            } for group in self.view_access_control_list]

        # Get execution parameters, append the object path as the last parameter
        parameters = execution_params.copy() if execution_params else []
        parameters.append(f"path={object_path}")

        # Get custom spark configuration and append it to default spark
        # parameters
        spark_conf = spark_conf or {}
        self.max_executors = max_executors
        if self.max_executors:
            spark_conf['spark.dynamicAllocation.maxExecutors'] = max_executors

        databricks_cluster_conf = self.get_databricks_cluster_conf(
            databricks_cluster_conf
        )

        self.render_params = render_params
        self.databricks_connection_id = databricks_connection_id
        self.pre_action = pre_action
        self.pre_action_conn_id = pre_action_conn_id

        # Call Airflow Databricks Operator
        super(DatabricksPysparkOperator, self).__init__(
            databricks_conn_id=databricks_connection_id,
            run_name=name,
            libraries=[
                {"whl": f"{s3_execution_path}/{applications}.wheelhouse.zip"}
            ],
            new_cluster=self.generate_new_cluster(
                databricks_node_type=databricks_node_type,
                databricks_policy_id=databricks_policy_id,
                databricks_cluster_conf=databricks_cluster_conf,
                spark_conf={**spark_conf, **DEFAULT_SPARK_CONF}
            ),
            spark_python_task=self.generate_spark_python_task(
                executable=f"{s3_execution_path}/launch.py",
                parameters=parameters
            ),
            polling_period_seconds=polling_interval,
            access_control_list=access_control_list+extra_access_control_list,
            **kwargs
        )

    # Generate spark_python_task dict from an executable path and
    # a list of parameters
    @staticmethod
    def generate_spark_python_task(executable: str, parameters: list):
        return {
            'python_file': executable,
            'parameters': parameters
        }

    # Generate new_cluster dict from a node_type, policy_id and
    # spark configuration dict
    @staticmethod
    def generate_new_cluster(databricks_node_type: str, 
                             databricks_policy_id: str, 
                             databricks_cluster_conf: Dict,
                             spark_conf: Dict):
        return {
            'node_type_id': databricks_node_type,
            'spark_conf': spark_conf,
            'policy_id': databricks_policy_id,
            **databricks_cluster_conf
            }      

    def get_databricks_creds(self):
        # Get credentials to call Databricks
        databricks_connection_str = Connection.get_connection_from_secrets(
            self.databricks_connection_id).get_uri()
        databricks_srv_airflow = (databricks_connection_str
                                  .replace("databricks://", "https://")
                                  .split("?token="))
        databricks_host = databricks_srv_airflow[0]
        databricks_token = databricks_srv_airflow[1]

        if databricks_host[-1] == "/":
            databricks_host = databricks_host[0:-1]

        # Mount header to Databricks call
        self.header = dict(Authorization=f"Bearer {databricks_token}")
        self.databricks_host = databricks_host

    def databricks_get_logs(self):

        self.log.info(f"\nRequesting execution log to Databricks {self.databricks_host}...")
        r = requests.get(f"{self.databricks_host}/api/2.1/jobs/runs/get-output",
                         headers=self.header, params={"run_id": self.run_id})

        if r.status_code == 200:
            if "logs" in r.json():
                self.log.info(f"{r.json()['logs']}")
            else:
                self.log.info(f"Can't get the logs for run_id {self.run_id} on Databricks")
        else:
            self.log.info(r.content)

    def execute(self, context):
        if self.pre_action is not None:
            pg_hook = PostgresHook(postgres_conn_id=self.pre_action_conn_id)
            pg_hook.run(self.pre_action)

        self.get_databricks_creds()

        try:
            result = super().execute(context) 
            self.databricks_get_logs()
            return result

        except AirflowException:
            if self.run_id:
                self.databricks_get_logs()
                db_hook = DatabricksHook(databricks_conn_id=self.databricks_connection_id)
                db_hook.cancel_run(self.run_id)
                self.log.info(
                    "Task: %s with run_id: %s was requested to be cancelled.", self.task_id, self.run_id
                )
            raise AirflowException

    def get_databricks_cluster_conf(self, databricks_cluster_conf):
        env = Variable.get('es_di_enablers/environment')

        # In case someone submits a proper Databricks conf
        if databricks_cluster_conf:
            if 'spark_env_vars' not in databricks_cluster_conf:
                databricks_cluster_conf["spark_env_vars"] = {}
            databricks_cluster_conf["spark_env_vars"]["ENVIRONMENT"] = env
            
            if 'custom_tags' not in databricks_cluster_conf:
                databricks_cluster_conf["custom_tags"] = {}
            databricks_cluster_conf["custom_tags"]["airflow-operator"] = "DatabricksPysparkOperator"            
            
            return databricks_cluster_conf

        # The default conf
        default_conf = {
            'aws_attributes': {'first_on_demand': 1},
            'autoscale': {
                'min_workers': 1,
                'max_workers': 3
            },
            'spark_env_vars': {
                'ENVIRONMENT': env
            },
            'custom_tags': {
                'airflow-operator': 'DatabricksPysparkOperator'
            },
            'apply_policy_default_values': 'true'
        }

        # Changing default conf with max_executors if used in
        if self.max_executors:
            default_conf["autoscale"]["max_workers"] = self.max_executors

        return default_conf

    def render_template_fields(self, context) -> None:
        # IF Boolean render_params is True, we run the render template on params and add it to the context
        if self.render_params:
            """ Add the rendered 'params' to the context dictionary before running the templating """
            # Like the original method, get the env if not provided
            jinja_env = self.get_template_env()

            if self.params:
                context['params'] = self.render_template(self.params, context, jinja_env, set())

            # Call the original method
            super().render_template_fields(context=context, jinja_env=jinja_env)

        else:
            super().render_template_fields(context)

fotocasa_android_fact_top_pains_comments_csat_push_alerts = DatabricksPysparkOperator(
    task_id="fotocasa_android_fact_top_pains_comments_csat_push_alerts",
    name=f"dag_fotocasa_android_fact_top_pains_comments_csat_push_alerts_{start_date_iso}",
    team="es_di_cross",
    object_path="common.csat_push_alerts.llm.fact_top_pains_comments_csat_push_alerts.DomainTopPainsCommentsCsatPushAlerts",
    execution_params=[
        f"start_date_iso={start_date_iso}",
        f"end_date_iso={end_date_iso}",
        f"marketplace=fotocasa",
        f"os_platform=Android"
    ],
    pre_action=f"DELETE FROM <table_schema>.<table_name> WHERE date_activity = '{start_date_iso}'and site = '<marketplace>' and sub_channel = 'Android'",
    pre_action_conn_id=dataland_conn_id,
    dag=dag
)