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

recurvedata/operators/link_operator/operator.py

@@ -0,0 +1,120 @@
+import copy
+from typing import Any, Optional
+
+from recurvedata.core.translation import _l
+from recurvedata.operators.operator import BaseOperator
+from recurvedata.operators.task import BaseTask
+
+
+class ConfigTask(BaseTask):
+    pass
+
+
+class LinkOperator(BaseOperator):
+    @classmethod
+    def get_link_setting(cls, node_config: dict) -> tuple[int, str, Optional[int]]:
+        source = node_config["source"]
+        return source["workflow_id"], source["workflow_version_tag"], source.get("link_node_key")
+
+    @classmethod
+    def config_schema(cls) -> dict:  # front-end does not use this config schema to show
+        return {
+            "type": "object",
+            "properties": {
+                "workflow_id": {
+                    "type": "string",
+                    "title": _l("Target Workflow"),
+                    "description": _l("ID of the workflow to link to"),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "workflow_version": {
+                    "type": "string",
+                    "title": _l("Workflow Version Tag"),
+                    "description": _l("Version tag of the target workflow (e.g. latest, v1.0)"),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "node_id": {
+                    "type": "string",
+                    "title": _l("Target Node"),
+                    "description": _l("ID of the node to link to in the target workflow"),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "variables": {
+                    "type": "string",
+                    "title": _l("Custom Variables"),
+                    "default": "{}",
+                    "description": _l("Custom variables to pass to the linked workflow in JSON format"),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "code",
+                        "lang": "json",
+                    },
+                },
+            },
+            "required": [
+                "workflow_id",
+                "workflow_version",
+            ],
+        }
+
+    def execute(self):
+        from recurvedata.executors import LinkExecutor
+
+        link_custom_variables = self.get_link_custom_variables(
+            self.dag, self.node, self.execution_date, self.variables, self.node.job_variable
+        )
+        link_settings = self.node.link_settings
+        link_executor = LinkExecutor(
+            self.dag,
+            self.node,
+            execution_date=self.execution_date,
+            link_workflow_id=link_settings["workflow_id"],
+            link_node_id=link_settings["node_id"],
+            custom_variables=link_custom_variables,
+            is_link_workflow=link_settings["is_link_workflow"],
+        )
+        # TODO(chenjingmeng): temporary solution to distinguish link operator
+        link_executor.node.is_link_op = True
+        link_executor.node.origin_node = self.node
+        link_executor.run()
+
+    @classmethod
+    def validate(cls, configuration) -> dict:
+        return configuration  # variables is dict type which will fail the json validation
+
+    def run_stage(self, stage):
+        return self.execute()
+
+    @classmethod
+    def ui_config_to_config(cls, configuration: dict[str, Any]) -> dict[str, Any]:
+        source = configuration["source"]
+        source["variables"] = configuration.get("variables", {})
+        return source
+
+    @classmethod
+    def get_link_custom_variables(cls, dag, node, execution_date, variables, job_variables: dict):
+        # if linkOp has not configured the variables, use job_variables https://project.feishu.cn/recurvedata/issue/detail/5342288226
+        link_variables = copy.deepcopy(job_variables) if job_variables else {}
+        task_obj = ConfigTask(dag, node, execution_date, variables)
+        link_custom_variables = task_obj.rendered_config["variables"]
+
+        if "execution_date" in variables and "execution_date" not in link_custom_variables:
+            # user may update `execution_date` in variable, which may not appear in link custom variable,
+            # so we need to pass the updated `execution_date` to link
+            link_custom_variables["execution_date"] = variables["execution_date"]
+        link_variables.update(link_custom_variables)
+        return link_variables
+
+    @classmethod
+    def get_ds_name_field_values(cls, rendered_config: dict) -> list[str]:
+        # todo: check linked node
+        return []

recurvedata/operators/models.py

@@ -0,0 +1,55 @@
+from dataclasses import dataclass
+
+from recurvedata.utils.crontab import next_schedule, previous_schedule
+from recurvedata.utils.date_time import _DATELIKE
+
+
+@dataclass
+class DagBase:
+    """
+    Recurve Web Job 对象
+    """
+
+    id: int
+    project_id: int
+    name: str
+    schedule_interval: str
+    scheduler_type: str
+    timezone: str
+    owner: str
+    full_refresh_models: bool = False
+    retries: int = None
+    retry_delay: int = None
+    # for dbt operator
+    skip_data_tests: bool = False
+
+    @classmethod
+    def normalize_schedule_interval(cls, interval: str) -> str:
+        mapping = {
+            "@once": None,
+            "@hourly": "0 * * * *",
+            "@daily": "0 0 * * *",
+            "@weekly": "0 0 * * 0",
+            "@monthly": "0 0 1 * *",
+            "@yearly": "0 0 1 1 *",
+        }
+        if interval in mapping:
+            return mapping[interval]
+        return interval
+
+    def next_schedule(self, dttm: _DATELIKE) -> _DATELIKE:
+        return next_schedule(self.schedule_interval, dttm)
+
+    def previous_schedule(self, dttm: _DATELIKE) -> _DATELIKE:
+        return previous_schedule(self.schedule_interval, dttm)
+
+    @property
+    def is_once(self):
+        return self.schedule_interval == "@once" or self.schedule_interval is None or self.schedule_interval == ""
+
+
+@dataclass
+class NodeBase:
+    id: int
+    node_key: str
+    name: str

recurvedata/operators/notify_operator/__init__.py

@@ -0,0 +1 @@
+from recurvedata.operators.notify_operator.operator import NotifyOperator

recurvedata/operators/notify_operator/operator.py

@@ -0,0 +1,180 @@
+import json
+import logging
+
+import jsonschema
+
+try:
+    from recurvedata.pigeon.connector.feishu import FeishuBot
+except ImportError:
+    pass
+
+from recurvedata.core.translation import _l
+from recurvedata.operators.operator import BaseOperator
+from recurvedata.operators.task import BaseTask
+
+logger = logging.getLogger(__name__)
+
+
+class NotifyTask(BaseTask):
+    @staticmethod
+    def split_str_lst(s: str, default_value: list[str] = None) -> list[str]:
+        if default_value is None:
+            default_value = []
+        if not s:
+            return default_value
+        return list(set(item.strip() for item in s.split(",")))
+
+    def execute_impl(self, *args, **kwargs):
+        config = self.rendered_config
+        bot = FeishuBot(**self.get_connection_by_name(config["feishu_bot"]).extra)
+
+        send_method = bot.send_message
+        send_conf = {"msg_type": config.msg_type}
+        email_lst = self.split_str_lst(config.at_user_email, default_value=None)
+        if config.msg_type == "post":
+            send_conf.update({"content": config.text_content})
+            if email_lst is None and not config.subject:
+                send_conf.update({"msg_type": "text"})
+                send_conf.update({"content": json.dumps({"text": config.text_content})})
+            else:
+                send_method = bot.send_text
+                send_conf.pop("msg_type")
+                send_conf.update({"email_lst": email_lst, "subject": config.subject})
+        else:
+            send_conf.update({"content": config.card_content})
+
+        user_lst = self.split_str_lst(config.email)
+        chat_name_lst = self.split_str_lst(config.chat_name)
+        chat_id_lst = self.split_str_lst(config.chat_id)
+        for email in user_lst:
+            send_method(receiver_type="user", user_email=email, **send_conf)
+        for chat_name in chat_name_lst:
+            send_method(receiver_type="group", chat_name=chat_name, **send_conf)
+        for chat_id in chat_id_lst:
+            send_method(receiver_type="group", chat_id=chat_id, **send_conf)
+        logger.info("Message was successfully sent to the receivers.")
+
+    @classmethod
+    def config_schema(cls):
+        # get_choices_by_type = cls.get_connection_names_by_type
+        return {
+            "type": "object",
+            "properties": {
+                "feishu_bot": {
+                    "type": "string",
+                    "title": _l("Feishu Bot Connection"),
+                    "description": _l("Select the Feishu bot connection to use for sending messages"),
+                    "ui:field": "ProjectConnectionSelectorField",
+                    "ui:options": {
+                        "supportTypes": [
+                            "feishu_bot",
+                        ],
+                    },
+                },
+                "msg_type": {
+                    "type": "string",
+                    "title": _l("Message Format"),
+                    "description": _l("Choose between simple text/post format or interactive card format"),
+                    "default": "post",
+                    "enum": ["post", "interactive"],
+                    "enumNames": ["post", "interactive"],
+                },
+                "subject": {
+                    "ui:hidden": '{{parentFormData.msg_type === "interactive"}}',
+                    "type": "string",
+                    "title": _l("Message Subject"),
+                    "description": _l("Subject line for the message. Supports template variables."),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "at_user_email": {
+                    "ui:hidden": '{{parentFormData.msg_type === "interactive"}}',
+                    "type": "string",
+                    "title": _l("Mention Users"),
+                    "description": _l(
+                        "Email addresses of users to @mention in the message. Separate multiple emails with commas. Use 'all' to @mention everyone."
+                    ),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "text_content": {
+                    "ui:hidden": '{{parentFormData.msg_type === "interactive"}}',
+                    "type": "string",
+                    "title": _l("Message Text"),
+                    "default": "",
+                    "description": _l("Main text content of the message. Supports template variables."),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "card_content": {
+                    "ui:hidden": '{{parentFormData.msg_type === "post"}}',
+                    "type": "string",
+                    "title": _l("Interactive Card JSON"),
+                    "description": _l("JSON definition for the interactive message card. Supports template variables."),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "code",
+                        "lang": "json",
+                    },
+                },
+                "email": {
+                    "type": "string",
+                    "title": _l("Individual Recipients"),
+                    "description": _l(
+                        "Email addresses of individual users to receive the message. Separate multiple emails with commas."
+                    ),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "chat_name": {
+                    "type": "string",
+                    "title": _l("Group Recipients (by Name)"),
+                    "description": _l(
+                        "Names of Feishu chat groups to receive the message. Separate multiple names with commas. Bot must be a member of the groups."
+                    ),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "chat_id": {
+                    "type": "string",
+                    "title": _l("Group Recipients (by ID)"),
+                    "description": _l(
+                        "IDs of Feishu chat groups to receive the message. Separate multiple IDs with commas. Bot must be a member of the groups."
+                    ),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+            },
+            "required": ["feishu_bot", "msg_type"],
+        }
+
+    @classmethod
+    def validate(cls, configuration) -> dict:
+        config = super().validate(configuration)
+        if not any([config["user_email"], config["chat_name"], config["chat_id"]]):
+            raise jsonschema.ValidationError(
+                message="at least one of (User Email, Chat Group Name, Chat Group ID) must be entered",
+                path=("user_email", "chat_name", "chat_id"),
+            )
+        if config["card_content"]:
+            try:
+                json.loads(config["card_content"])
+            except Exception:
+                raise jsonschema.ValidationError(message="Card Content should be valid JSON", path=("card_content",))
+        return config
+
+
+class NotifyOperator(BaseOperator):
+    task_cls = NotifyTask

recurvedata/operators/operator.py

@@ -0,0 +1,119 @@
+import datetime
+from typing import TYPE_CHECKING, Callable
+
+from recurvedata.operators.base import Configurable
+from recurvedata.operators.ui import format_config_schema
+from recurvedata.utils.registry import Registry
+
+if TYPE_CHECKING:
+    from recurvedata.executors.models import ExecutorDag, ExecutorNode
+    from recurvedata.operators.task import BaseTask
+
+_registry = Registry(key_callback=lambda x: x.name())
+
+
+def get_operator_class(name):
+    return _registry.get(name)
+
+
+class BaseOperator(Configurable):
+    stages = ()
+    task_cls = None
+    web_install_require = []  # the python modules needed by Recurve Web
+    worker_install_require = []  # the python modules needed by Recurve Worker
+
+    def __init__(
+        self, dag: "ExecutorDag", node: "ExecutorNode", execution_date: datetime.datetime, variables: dict = None
+    ):
+        self.dag: "ExecutorDag" = dag
+        self.node: "ExecutorNode" = node
+        self.execution_date: datetime.datetime = execution_date
+        self.task: BaseTask = None
+        self.variables: dict = variables or {}
+        self.init_task()
+
+    def __init_subclass__(cls, **kwargs):
+        _registry.add(cls)
+
+    def execute(self):
+        task_obj = self.task
+        if task_obj:
+            task_obj.execute()
+
+    def init_task(self):
+        if self.task_cls:
+            self.task = self.task_cls(self.dag, self.node, self.execution_date, self.variables)
+
+    def set_execution_date(self, execution_date):
+        if execution_date == self.execution_date:
+            return
+        # 有些 T0 任务，需要修改 operator 的 execute_date
+        # task 也需要设置 execute_date
+        self.execution_date = execution_date
+        if self.task:
+            self.task.set_execution_date(execution_date)
+
+    @classmethod
+    def config_schema(cls) -> dict:
+        if cls.task_cls:
+            return cls.task_cls.config_schema()
+
+    @classmethod
+    def to_dict(cls) -> dict:
+        return {
+            "name": cls.name(),
+            "config_schema": {"source": cls.config_schema()},
+        }
+
+    @classmethod
+    def get_setup_install_require(cls) -> dict:
+        return {
+            "web": cls.web_install_require,
+            "worker": cls.worker_install_require,
+        }
+
+    @classmethod
+    def ui_config_schema(cls) -> dict:
+        res = {
+            "source": {
+                "name": "Source",
+                "config_schema": format_config_schema(cls.config_schema(), "source"),
+            }
+        }
+        return res
+
+    @staticmethod
+    def _add_schema_name_to_json_schema_error(schema_name: str, validate_func: Callable, *args, **kwargs):
+        try:
+            return validate_func(*args, **kwargs)
+        except Exception as e:
+            e.schema_name = schema_name
+            raise e
+
+    @classmethod
+    def ui_validate(cls, configuration: dict) -> dict:
+        return {
+            "source": cls._add_schema_name_to_json_schema_error(
+                "source", cls.validate, cls.ui_config_to_config(configuration)
+            ),
+        }
+
+    @classmethod
+    def ui_config_to_config(cls, configuration: dict) -> dict:
+        """
+        ui_config: 前端保存时传来的配置，和 ui_config_schema 里一一对应，通常是
+            {'source': source_dct, 'meta': meta_dct} 格式
+        config: Operator 具体的配置，通常指 ui_config 里的 source。
+
+        区分 ui config 和 config 的原因:
+        前端页面，根据 ui config_schema 里配置的，基本分为 Source, Meta 两大块。
+        其中 Meta 是调度器相关配置，和具体的 Operator 关联不大。
+        为了把调度器相关的校验逻辑、schema 和具体 Operator 区分开，
+        设置了 ui_config
+        """
+        return configuration["source"]
+
+    @classmethod
+    def get_ds_name_field_values(cls, rendered_config: dict) -> list[str]:
+        config = cls.ui_config_to_config(rendered_config)
+        return cls.task_cls.get_ds_name_field_values(config)

recurvedata/operators/python_operator/__init__.py

@@ -0,0 +1 @@
+from recurvedata.operators.python_operator.operator import PythonOperator

recurvedata/operators/python_operator/operator.py

@@ -0,0 +1,132 @@
+import ast
+import logging
+import os
+from tempfile import NamedTemporaryFile
+
+from recurvedata.config import PY_PACKAGES_PATH, RECURVE_EXECUTOR_PYENV_NAME
+from recurvedata.core.translation import _l
+from recurvedata.operators.config import CONF
+from recurvedata.operators.operator import BaseOperator
+from recurvedata.operators.task import BaseTask
+from recurvedata.utils.mp import robust_run_subprocess, run_subprocess
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_PY_VERSION = os.environ.get("RECURVE_OPERATOR_PYTHON_DEFAULT_VERSION", "3.11.9")
+
+
+class PythonRequirementsMixin:
+    @staticmethod
+    def _install_requirements(requirements: str, pyenv_name: str):
+        if pyenv_name != RECURVE_EXECUTOR_PYENV_NAME:
+            requirements += "\nrecurvedata-lib[slim]"
+        if not requirements:
+            return
+        logger.info("installing requirements")
+        # Install recurvedata-lib from local package if it's a new virtualenv
+        if pyenv_name != RECURVE_EXECUTOR_PYENV_NAME:
+            python = CONF.PYENV_PYTHON_PATH.format(pyenv=pyenv_name)
+            run_subprocess(
+                f"{python} -m pip install -v --no-index --find-links={PY_PACKAGES_PATH} recurvedata-lib[slim]".split()
+            )
+        with NamedTemporaryFile(mode="w+t", prefix="recurve_python_requirements_", suffix=".txt") as requirements_path:
+            requirements_path.write(requirements)
+            requirements_path.flush()
+            python = CONF.PYENV_PYTHON_PATH.format(pyenv=pyenv_name)
+            run_subprocess(f"{python} -m pip install -r {requirements_path.name}".split())
+
+
+class PythonTask(BaseTask, PythonRequirementsMixin):
+    ds_name_fields = ("python_env",)
+
+    @classmethod
+    def config_schema(cls) -> dict:
+        return {
+            "type": "object",
+            "properties": {
+                "python_env": {
+                    "type": "string",
+                    "title": _l("Python Env"),
+                    "description": _l(
+                        "Python virtual environment name that will be created and can be shared between tasks"
+                    ),
+                    "ui:field": "ProjectConnectionSelectorField",
+                    "ui:options": {
+                        "supportTypes": [
+                            "python",
+                        ]
+                    },
+                },
+                "code": {
+                    "type": "string",
+                    "title": _l("Code"),
+                    "description": _l(
+                        "Python code that will be executed. Supports Jinja templating for dynamic code generation and variable substitution."
+                    ),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "code",
+                        "lang": "python",
+                    },
+                },
+            },
+            "required": ["python_env", "code"],
+        }
+
+    def __custom_os_env(self) -> dict:
+        custom_env = os.environ.copy()
+        return custom_env
+
+    def __run_python(self, filename: str, pyenv: str, os_env: dict):
+        script_path = os.path.abspath(filename)
+        python = CONF.PYENV_PYTHON_PATH.format(pyenv=pyenv)
+        output, ret_code = robust_run_subprocess([python, script_path], env=os_env, _logger=logger)
+        if ret_code:
+            raise RuntimeError(f"Python Operator Error:\n{output}")
+
+    def _prepare_env(self, python_config: dict):
+        pyenv_name: str = python_config.get("pyenv")
+        py_version: str = python_config.get("python_version", DEFAULT_PY_VERSION)
+        self._install_virtualenv(py_version, pyenv_name)
+        self._install_requirements(python_config.get("requirements", ""), pyenv_name)
+
+    @staticmethod
+    def _install_virtualenv(py_version: str, pyenv_name: str):
+        python_path: str = CONF.PYENV_PYTHON_PATH.format(pyenv=pyenv_name)
+        if os.path.exists(python_path):
+            return
+        run_subprocess(["pyenv", "virtualenv", py_version, pyenv_name])
+
+    def execute_impl(self, *args, **kwargs):
+        config = self.rendered_config
+        code = config.code
+        os_env = self.__custom_os_env()
+
+        conn_config: dict = self.get_connection_by_name(config.python_env).extra
+        self._prepare_env(conn_config)
+        pyenv = conn_config["pyenv"]
+        prefix = f"recurve_python_{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()
+            logger.info(code)
+            self.__run_python(tmp_file.name, pyenv, os_env)
+
+
+class PythonOperator(BaseOperator):
+    task_cls = PythonTask
+
+    @classmethod
+    def validate(cls, configuration) -> dict:
+        res = super().validate(configuration)
+        # syntax_error = cls._get_python_code_syntax_error(res['code'])
+        # if syntax_error:
+        #     raise jsonschema.ValidationError(f'Python Syntax Error {syntax_error}')
+        return res
+
+    @staticmethod
+    def _get_python_code_syntax_error(code):
+        try:
+            ast.parse(code)
+        except SyntaxError as e:
+            return e

recurvedata/operators/sensor_operator/__init__.py

@@ -0,0 +1 @@
+from recurvedata.operators.sensor_operator.operator import SensorOperator

recurvedata/operators/sensor_operator/airflow_utils.py

@@ -0,0 +1,63 @@
+import datetime
+import os
+
+from airflow.models import DAG
+from airflow.timetables.interval import CronDataIntervalTimetable
+
+from recurvedata.schedulers.airflow import AirflowScheduler
+from recurvedata.schedulers.model import SchedulerNode
+from recurvedata.schedulers.utils import format_dag_id
+
+
+def prepare_airflow_env():
+    if "_AIRFLOW__AS_LIBRARY" in os.environ:
+        del os.environ["_AIRFLOW__AS_LIBRARY"]
+
+
+def get_dag_from_db(dag_id):
+    from airflow.exceptions import AirflowException
+    from airflow.models import DagBag
+
+    dagbag = DagBag(read_dags_from_db=True)
+    dag = dagbag.get_dag(dag_id)
+    if not dag:
+        raise AirflowException(f"Dag {dag_id!r} could not be found; either it does not exist or it failed to parse.")
+    return dag
+
+
+def build_execute_context(dag, task, run_id) -> dict:
+    dag_run = dag.get_dagrun(run_id=run_id)
+    task_instance = dag_run.get_task_instance(task.task_id)
+    context = {
+        "dag": dag,
+        "task": task,
+        "dag_run": dag_run,
+        "task_instance": task_instance,
+        "execution_date": dag_run.execution_date,
+        "logical_date": dag_run.logical_date,
+        "data_interval_end": dag_run.data_interval_end,
+    }
+
+    return context
+
+
+def format_external_dag_id(job_id: int) -> str:
+    return format_dag_id(job_id)
+
+
+def format_external_task_id(node_key: str) -> str:
+    node = SchedulerNode(node_key=node_key, operator="external_operator", id=0, name="external_task")
+    return AirflowScheduler.format_task_id(node)
+
+
+def data_interval_end_to_data_interval_start(
+    dag: DAG, data_interval_end: datetime.datetime
+) -> datetime.datetime | None:
+    if hasattr(dag, "timetable") and isinstance(dag.timetable, CronDataIntervalTimetable):
+        data_interval_start = dag.timetable._get_prev(data_interval_end)
+        data_interval_end2 = dag.timetable._get_next(data_interval_start)
+        if data_interval_end2 != data_interval_end:
+            data_interval_start = dag.timetable._get_prev(data_interval_start)
+        return data_interval_start
+    else:
+        return data_interval_end