recurvedata-lib 0.1.487__py2.py3-none-any.whl

recurvedata/operators/sensor_operator/operator.py

@@ -0,0 +1,172 @@
+import logging
+import os
+from tempfile import NamedTemporaryFile
+
+from recurvedata.core.translation import _l
+from recurvedata.operators.operator import BaseOperator
+from recurvedata.operators.task import BaseTask
+from recurvedata.utils.mp import robust_run_subprocess
+
+logger = logging.getLogger(__name__)
+
+AIRFLOW_PYTHON_PATH = "python"  # system python path
+
+
+class SensorTask(BaseTask):
+    """
+    Sensor operator is used to create a sensor that could check upstream task status,
+    will wait until the upstream task is success.
+    It uses airflow ExternalTaskSensor to check the status of the upstream task.
+    For modeling pipeline,
+        the node_key is the node_key of the selected model.
+    for advanced pipeline,
+        - if normal Operator like SQLOperator, the node_key is the node_key of the SQLOperator.
+        - if Modeling Pipeline ( which is LinkModelPipelineOperator), the node_key is  "{node_key of the LinkModelPipelineOperator}.{node_key of the model}"
+
+    """
+
+    # todo: same schedule interval dependency
+
+    @classmethod
+    def config_schema(cls):
+        return {
+            "type": "object",
+            "properties": {
+                "project_id": {
+                    "type": "string",
+                    "title": _l("Project Name"),
+                    "description": _l("Project Name containing the external task"),
+                    "ui:field": "SensorOperatorProjectSelectField",
+                },
+                "job_id": {
+                    "type": "string",
+                    "title": _l("Job Name"),
+                    "description": _l("Job Name of the external task"),
+                    "ui:field": "SensorOperatorJobSelectField",
+                },
+                "node_key": {
+                    "type": "string",
+                    "title": _l("Node Name"),
+                    "description": _l("Node Name of the external task"),
+                    "ui:field": "SensorOperatorNodeSelectField",
+                },
+                "wait_time": {
+                    "type": "integer",
+                    "title": _l("Wait Time"),
+                    "description": _l("Wait time in seconds between checks"),
+                    "ui:options": {
+                        "min": 0,
+                        "step": 1,
+                    },
+                },
+                "timeout": {
+                    "type": "integer",
+                    "title": _l("Timeout"),
+                    "description": _l("Timeout"),
+                    "default": 60,
+                    "ui:options": {
+                        "min": 0,
+                        "step": 1,
+                    },
+                },
+            },
+            "required": ["project_id", "job_id", "node_key"],
+        }
+
+    def generate_airflow_operator_code(self):
+        config = self.rendered_config
+        external_job_id = config["job_id"]
+        external_node_key = config["node_key"]
+        execution_delta = config.get("wait_time", 0)
+        timeout = config.get("timeout", 3600)
+        dag_name = self.dag.name
+
+        return f"""
+import sys
+from recurvedata.operators.sensor_operator.airflow_utils import format_external_dag_id, format_external_task_id, data_interval_end_to_data_interval_start
+from airflow.sensors.external_task import ExternalTaskSensor
+import datetime
+import logging
+from recurvedata.schedulers.consts import is_dev_run_job
+
+logger = logging.getLogger()
+external_dag_id = format_external_dag_id({external_job_id!r})
+external_task_id = format_external_task_id({external_node_key!r})
+
+external_dag = get_dag_from_db(external_dag_id)
+if not external_dag:
+    raise ValueError("External DAG not found")
+external_task = external_dag.get_task(external_task_id)
+if not external_task:
+    raise ValueError("External Task not found")
+
+data_interval_end = context["data_interval_end"]
+external_data_interval_end = data_interval_end - datetime.timedelta(seconds={execution_delta})
+external_data_interval_start = data_interval_end_to_data_interval_start(external_dag, external_data_interval_end)
+
+logger.debug("external_data_interval_start " + str(external_data_interval_start))
+
+tmp_task_id="tmp_task_id_for_external_task_sensor"
+operator = ExternalTaskSensor(
+        dag=dag,
+        task_id=tmp_task_id,
+        external_dag_id=external_dag_id,
+        external_task_id=external_task_id,
+        execution_date_fn = lambda *args, **kwargs: external_data_interval_start,
+        execution_timeout = datetime.timedelta(seconds={timeout}),
+    )
+if is_dev_run_job({dag_name!r}):
+    logger.info(f"dag_name: {dag_name!r}")
+    logger.info(f"skip: SensorOperator is not working in dev mode")
+    sys.exit(0)
+
+"""
+
+    def generate_airflow_code(self) -> str:
+        config = self.rendered_config
+        timeout = config.get("timeout", 3600)
+        operator_code = self.generate_airflow_operator_code()
+        return """
+import os
+import time
+from recurvedata.operators.sensor_operator.airflow_utils import prepare_airflow_env, get_dag_from_db, \
+build_execute_context
+from recurvedata.utils.timeout import timeout
+
+prepare_airflow_env()
+
+dag_id = os.environ.get("AIRFLOW_CTX_DAG_ID")
+task_id = os.environ.get("AIRFLOW_CTX_TASK_ID")
+run_id = os.environ.get("AIRFLOW_CTX_DAG_RUN_ID")
+
+dag = get_dag_from_db(dag_id)
+task = dag.get_task(task_id)
+context = build_execute_context(dag, task, run_id)
+
+{operator_code}
+
+with timeout({timeout}):
+    operator.execute(context)
+    """.format(
+            operator_code=operator_code,
+            timeout=timeout,
+        )
+
+    def __run_airflow_operator(self, filename: str):
+        script_path = os.path.abspath(filename)
+        env = os.environ.copy()
+        output, ret_code = robust_run_subprocess([AIRFLOW_PYTHON_PATH, script_path], _logger=logger, env=env)
+        if ret_code:
+            raise RuntimeError(f"Airflow Error:\n{output}")
+
+    def execute_impl(self, *args, **kwargs):
+        code = self.generate_airflow_code()
+        prefix = f"reorc_sensor_operator_{self.dag.id}_{self.node.id}_"
+        with NamedTemporaryFile(mode="w+t", prefix=prefix, suffix=".py") as tmp_file:
+            tmp_file.write(code)
+            tmp_file.flush()
+            self.__run_airflow_operator(tmp_file.name)
+
+
+class SensorOperator(BaseOperator):
+    task_cls = SensorTask

recurvedata/operators/spark_operator/__init__.py

@@ -0,0 +1 @@
+from recurvedata.operators.spark_operator.operator import SparkOperator

recurvedata/operators/spark_operator/operator.py

@@ -0,0 +1,200 @@
+import json
+import logging
+import os
+from functools import lru_cache
+from importlib import resources
+from subprocess import PIPE, STDOUT, Popen
+from tempfile import NamedTemporaryFile
+from textwrap import dedent
+
+from recurvedata.core.translation import _l
+from recurvedata.operators.operator import BaseOperator
+from recurvedata.operators.task import BaseTask
+
+logger = logging.getLogger(__name__)
+
+
+@lru_cache()
+def get_sample_code():
+    return resources.files("recurvedata.operators.spark_operator").joinpath("spark_sample.py").read_text()
+
+
+class SparkTask(BaseTask):
+    @classmethod
+    def config_schema(cls):
+        return {
+            "type": "object",
+            "properties": {
+                "spark_source": {
+                    "type": "string",
+                    "title": _l("Spark Environment"),
+                    "ui:field": "ProjectConnectionSelectorField",
+                    "ui:options": {
+                        "supportTypes": [
+                            "spark",
+                        ],
+                    },
+                    "description": _l("Select the Spark environment and version to use for this task"),
+                },
+                "env": {
+                    "type": "string",
+                    "title": _l("Environment Variables"),
+                    "default": "{}",
+                    "description": _l(
+                        'Additional environment variables in JSON format (e.g. {"HADOOP_CONF_DIR": "/etc/hadoop/conf"})'
+                    ),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "code",
+                        "lang": "json",
+                    },
+                },
+                "execution_config": {
+                    "type": "string",
+                    "title": _l("Spark Configuration"),
+                    "default": dedent(
+                        """\
+                        {
+                            "master": "yarn",
+                            "executor-memory": "4g",
+                            "num-executors": "10",
+                            "executor-cores": "2",
+                            "queue": "default",
+                            "conf": {
+                                "spark.dynamicAllocation.enabled": "False"
+                            }
+                        }
+                    """
+                    ),
+                    "description": _l(
+                        "Spark execution parameters and configurations. See "
+                        "<a target='_blank' href='https://spark.apache.org/docs/latest/configuration.html'>"
+                        "Spark Docs</a> for available options"
+                    ),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "code",
+                        "lang": "json",
+                    },
+                },
+                "code": {
+                    "type": "string",
+                    "title": _l("Spark Code"),
+                    "default": get_sample_code(),
+                    "description": _l(
+                        "PySpark code to execute. The default template shows how to create a SparkSession "
+                        "(Spark 2.3+). Supports Jinja templating for dynamic code generation."
+                    ),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "code",
+                        "lang": "python",
+                    },
+                },
+            },
+            "required": ["spark_source", "env", "execution_config", "code"],
+        }
+
+    def __create_env(self, source_env, extra_env):
+        env = os.environ.copy()
+        env.update(source_env)
+        env.update(extra_env)
+        return env
+
+    def __create_bash_command(self, script_path, submitter, execution_config, source_conf: dict):
+        execution_conf_list = []
+        conf_names = []
+        for k, v in execution_config.items():
+            if k == "conf":
+                for k2, v2 in v.items():
+                    execution_conf_list.append(f"--conf {k2}={v2}")
+                    conf_names.append(k2)
+            else:
+                execution_conf_list.append(f"--{k} {v}")
+        for k, v in source_conf.items():
+            if k in conf_names:
+                continue
+            execution_conf_list.append(f"--conf {k}={v}")
+        execution_conf_str = " ".join(execution_conf_list)
+        bash_command = submitter + " " + execution_conf_str + " " + script_path
+        return bash_command
+
+    def __execute_command(self, bash_command, env):
+        logger.info("Running command: %s", bash_command)
+        sub_process = Popen(["bash", "-c", bash_command], stdout=PIPE, stderr=STDOUT, env=env)
+        logger.info("Output:")
+        for raw_line in iter(sub_process.stdout.readline, b""):
+            line = raw_line.decode("utf8").rstrip()
+            logger.info(line)
+        sub_process.wait()
+        logger.info("Node exited with return code %s", sub_process.returncode)
+        if sub_process.returncode:
+            raise Exception("Spark node failed")
+
+    @classmethod
+    def __filter_empty_value_in_dict(cls, dct: dict):
+        if not dct:
+            return dct
+        return {k: v for (k, v) in dct.items() if (v is not None and v != "" and v != {})}
+
+    @classmethod
+    def _merge_dict(cls, priority_dct: dict, other_dct: dict):
+        """
+        Filter out empty values
+        """
+        if not (other_dct and priority_dct):
+            return priority_dct or other_dct
+        result_dct = {}
+        for key in set(list(priority_dct.keys()) + list(other_dct.keys())):
+            if key not in priority_dct:
+                result_dct[key] = other_dct[key]
+                continue
+            if key not in other_dct:
+                result_dct[key] = priority_dct[key]
+                continue
+            if isinstance(priority_dct[key], dict) and isinstance(other_dct, dict):
+                result_dct[key] = cls._merge_dict(priority_dct[key], other_dct[key])
+                continue
+            result_dct[key] = priority_dct[key]
+        return result_dct
+
+    def __excute_spark_code(self, config):
+        spark_source = self.must_get_connection_by_name(config.spark_source)
+        submitter = spark_source.extra.get("submitter")
+        source_conf = spark_source.extra.get("conf", {})
+        source_env = self.__filter_empty_value_in_dict(
+            spark_source.extra.get("env")
+        )  # Some empty values may be saved when saving on the page
+        execution_config = self._merge_dict(
+            self.__filter_empty_value_in_dict(json.loads(config.execution_config)),
+            self.__filter_empty_value_in_dict(spark_source.extra.get("execution_config")),
+        )
+        extra_env = json.loads(config.env)  # User input, don't filter empty values
+        code = config.code
+
+        prefix = f"recurve_pyspark_{self.dag.dag_id}_{self.node.id}_"
+        with NamedTemporaryFile(mode="w+t", prefix=prefix, suffix=".py") as tmp_file:
+            tmp_file.write(code)
+            tmp_file.flush()
+            logger.info(code)
+            script_path = os.path.abspath(tmp_file.name)
+            bash_command = self.__create_bash_command(script_path, submitter, execution_config, source_conf)
+            env = self.__create_env(source_env, extra_env)
+            self.__execute_command(bash_command, env)
+
+    def execute_impl(self, *args, **kwargs):
+        config = self.rendered_config
+        self.__excute_spark_code(config)
+        return None
+
+
+class SparkOperator(BaseOperator):
+    task_cls = SparkTask
+
+    @classmethod
+    def validate(cls, configuration):
+        config = super().validate(configuration)
+        # execution_config = json.loads(config['execution_config'])
+        # if execution_config['master'] != 'yarn':
+        #     raise jsonschema.ValidationError(message='master should be yarn')
+        return config

recurvedata/operators/spark_operator/spark_sample.py

@@ -0,0 +1,47 @@
+try:
+    from pyspark.sql import SparkSession
+    from pyspark.sql.types import IntegerType, StringType, StructField, StructType
+
+    # Initialize a Spark session
+    spark = SparkSession.builder.appName("PySpark SQL Example").getOrCreate()
+
+    # Define the schema
+    schema = StructType(
+        [
+            StructField("name", StringType(), True),
+            StructField("age", IntegerType(), True),
+            StructField("city", StringType(), True),
+        ]
+    )
+
+    # Create a DataFrame manually
+    data = [
+        ("Alice", 34, "New York"),
+        ("Bob", 45, "Los Angeles"),
+        ("Cathy", 29, "Chicago"),
+        ("David", 31, "New York"),
+        ("Emma", 42, "San Francisco"),
+    ]
+
+    df = spark.createDataFrame(data, schema)
+
+    # Show the DataFrame
+    df.show()
+
+    # Register the DataFrame as a temporary view
+    df.createOrReplaceTempView("people")
+
+    # Perform a SQL query
+    result_df = spark.sql("SELECT * FROM people WHERE age > 30")
+
+    # Show the result of the SQL query
+    result_df.show()
+
+    # Write the result to another CSV file
+    result_df.write.csv("output.csv", header=True)
+
+    # Stop the Spark session
+    spark.stop()
+
+except ImportError:
+    pass

recurvedata/operators/sql_operator/__init__.py

@@ -0,0 +1 @@
+from recurvedata.operators.sql_operator.operator import SQLOperator

recurvedata/operators/sql_operator/operator.py

@@ -0,0 +1,90 @@
+from typing import Any
+
+import jsonschema
+
+from recurvedata.connectors.service import list_sql_operator_types
+from recurvedata.core.translation import _l
+from recurvedata.operators.operator import BaseOperator
+from recurvedata.operators.task import BaseTask
+from recurvedata.operators.utils import lineage
+
+
+class SQLTask(BaseTask):
+    no_template_fields = ("autocommit", "data_source_name")
+
+    @classmethod
+    def config_schema(cls):
+        # get_names_by_type = cls.get_connection_names_by_type
+        return {
+            "type": "object",
+            "properties": {
+                "data_source_name": {
+                    "type": "string",
+                    "title": _l("Data Source"),
+                    "ui:field": "ProjectConnectionSelectorField",
+                    "ui:options": {"supportTypes": list_sql_operator_types()},
+                },
+                # "database": {
+                #     "type": "string",
+                #     "title": _l("Database"),
+                #     "ui:field": "CodeEditorWithReferencesField",
+                #     "ui:options": {
+                #         "type": "plain",
+                #     },
+                # },
+                "sql": {
+                    "type": "string",
+                    "title": _l("SQL Query"),
+                    "description": _l(
+                        "Execute single or multiple SQL statements. "
+                        "Supports Jinja templating for variables and dynamic queries."
+                    ),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "code",
+                        "lang": "sql",
+                        "sqlLang": "sql",
+                    },
+                },
+            },
+            "required": ["data_source_name", "sql"],
+        }
+
+    @classmethod
+    def validate(cls, configuration: dict[str, Any]) -> dict[str, Any]:
+        config = super().validate(configuration)
+
+        ds = cls.must_get_connection_by_name(configuration["data_source_name"])
+        if not ds.is_dbapi:
+            raise jsonschema.ValidationError(message="only DBAPI is supported", path=("data_source_name",))
+        return config
+
+    def execute_impl(self, *args, **kwargs):
+        config = self.rendered_config
+        ds = self.get_connection_by_name(config.data_source_name)
+        connector = ds.connector
+        # if config.database:
+        #     connector.database = config.database
+
+        queries = config.sql
+        if connector.is_hive():
+            # Set spark.app.name to help locate the specific Recurve task in the YARN UI
+            queries = f"SET spark.app.name=recurve.{self.dag.name}.{self.node.name};\n{queries}"
+
+        comment = self.get_query_comment_conf()
+        annotated_queries = connector.add_leading_comment(queries, comment)
+
+        connector.execute(annotated_queries, autocommit=config.get("autocommit", False))
+        return None
+
+    def parse_lineage(self):
+        config = self.rendered_config
+        ds = self.get_connection_by_name(config.data_source_name)
+        if not lineage.supported_recurve_ds_type(ds.ds_type):
+            return
+        res = lineage.parse_lineage(config.sql, config.database, ds.name, ds.ds_type)
+        return res
+
+
+class SQLOperator(BaseOperator):
+    task_cls = SQLTask