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

recurvedata/schedulers/airflow_db_process.py

@@ -0,0 +1,331 @@
+import datetime
+import logging
+from typing import Any, Optional
+
+from airflow.api.common.mark_tasks import set_dag_run_state_to_failed, set_dag_run_state_to_success
+from airflow.api.common.trigger_dag import trigger_dag
+from airflow.models import DAG, DagModel, DagRun, TaskInstance
+from airflow.models.serialized_dag import SerializedDagModel
+from airflow.utils.session import create_session, provide_session
+from airflow.utils.state import TaskInstanceState
+from sqlalchemy import Index, Table
+from sqlalchemy.orm import Session
+from sqlalchemy.schema import CreateIndex
+
+from recurvedata.utils.date_time import to_local_datetime, utcnow
+
+logger = logging.getLogger(__name__)
+
+
+class AirflowDbService:
+    @classmethod
+    def update_dag(cls, dag: DAG):
+        logger.info(f"start sync dag {dag.dag_id} to serialized_dag")
+        DAG.sync_to_db(dag)
+        SerializedDagModel.write_dag(dag)
+        logger.info(f"finish sync {dag.dag_id} to serialized_dag")
+
+    @classmethod
+    def activate_dag(cls, dag: DAG):
+        with create_session() as session:
+            d = session.query(DagModel).filter(DagModel.dag_id == dag.dag_id).one_or_none()
+            if not d:
+                raise ValueError(f"dag not exists: {dag.dag_id}")
+
+            if d.is_paused is False:
+                logger.info(f"{dag.dag_id} is active, no need to activate")
+                return
+
+            logger.info(f"start activate_dag dag {dag.dag_id}")
+
+            d.is_paused = False
+            session.merge(d)
+            session.commit()
+
+            logger.info(f"finish activate_dag dag {dag.dag_id}")
+
+    @classmethod
+    def deactivate_dag(cls, dag: DAG):
+        with create_session() as session:
+            d = session.query(DagModel).filter(DagModel.dag_id == dag.dag_id).one_or_none()
+            if not d:
+                raise ValueError(f"dag not exists: {dag.dag_id}")
+
+            if d.is_paused is True:
+                logger.info(f"{dag.dag_id} is deactive, no need to deactivate")
+                return
+
+            logger.info(f"start deactivate_dag dag {dag.dag_id}")
+
+            d.is_paused = True
+            session.merge(d)
+            session.commit()
+
+            logger.info(f"finish deactivate_dag dag {dag.dag_id}")
+
+    @classmethod
+    def delete_dag(cls, dag_id: str, job_name: str):
+        with create_session() as session:
+            d: DagModel = session.query(DagModel).filter(DagModel.dag_id == dag_id).one_or_none()
+            if not d:
+                logger.warning(f"dag not exists: {dag_id}")
+                return
+
+            logger.info(f"start delete_dag dag {job_name} {dag_id}")
+            d.is_paused = True
+            d.is_active = False
+            session.merge(d)
+            session.commit()
+
+            logger.info(f"finish delete_dag dag {job_name} {dag_id}")
+            # todo(chenjingmeng): delete dag
+
+    @classmethod
+    def trigger_job_run(
+        cls,
+        dag: DAG,
+        execution_date: datetime.datetime,
+        include_past: bool,
+        include_future: bool,
+        run_type: str,
+        conf: dict[str, Any] | None = None,
+    ):
+        execution_date_ds = execution_date.isoformat()
+        run_id = f"{run_type}__{execution_date_ds}"
+        reference_date = to_local_datetime(execution_date_ds)
+        current_date = utcnow()
+        airflow_current_date = dag.previous_schedule(current_date)
+
+        if include_past:
+            airflow_start_date = dag.start_date or dag.default_args.get("start_date")
+            if airflow_start_date:
+                tmp_date = dag.previous_schedule(reference_date)
+                while tmp_date >= airflow_start_date:
+                    cls._trigger_run_if_not_exists(
+                        dag.dag_id, run_id=f"{run_type}__{tmp_date.isoformat()}", execution_date=tmp_date, conf=conf
+                    )
+                    tmp_date = dag.previous_schedule(tmp_date)
+
+        if include_future:
+            tmp_date = dag.following_schedule(reference_date)
+            while tmp_date <= airflow_current_date:
+                cls._trigger_run_if_not_exists(
+                    dag.dag_id, run_id=f"{run_type}__{tmp_date.isoformat()}", execution_date=tmp_date, conf=conf
+                )
+                tmp_date = dag.following_schedule(tmp_date)
+
+        cls._trigger_run_if_not_exists(dag.dag_id, run_id, execution_date=execution_date, conf=conf)
+
+    @staticmethod
+    def _trigger_run_if_not_exists(
+        dag_id: str, run_id: str, execution_date: datetime.datetime, conf: dict[str, Any] | None = None
+    ):
+        import recurvedata.schedulers.airflow_trigger_dag_patch  # noqa
+
+        with create_session() as session:
+            existing_run = session.query(DagRun).filter(DagRun.dag_id == dag_id, DagRun.run_id == run_id).first()
+            if existing_run:
+                logger.info(f"Skipping existing run for {dag_id} at {run_id}")
+                return
+        logger.info(f"start trigger dag_run for {dag_id} at {run_id}")
+        trigger_dag(dag_id, run_id=run_id, execution_date=execution_date, conf=conf, replace_microseconds=False)
+        logger.info(f"finished trigger dag_run for {dag_id} at {run_id} execution_date: {execution_date}")
+
+    @staticmethod
+    @provide_session
+    def _get_rerun_earliest_execution_date(dag: DAG, session: Session = None) -> Optional[datetime.datetime]:
+        earliest_dag_run = (
+            session.query(DagRun).filter(DagRun.dag_id == dag.dag_id).order_by(DagRun.execution_date).first()
+        )
+        return earliest_dag_run and earliest_dag_run.execution_date
+
+    @classmethod
+    def rerun_job_run(
+        cls,
+        dag: DAG,
+        run_id: str | None,
+        min_execution_date: datetime.datetime | None,
+        max_execution_date: datetime.datetime | None,
+        failed_only: bool,
+    ):
+        drs: list[DagRun] = DagRun.find(
+            dag_id=dag.dag_id,
+            run_id=run_id,
+            execution_start_date=min_execution_date,
+            execution_end_date=max_execution_date,
+        )
+        if not drs:
+            logger.info(f"skip rerun, no dag_run found for {dag.dag_id} at {run_id}")
+            return
+        clear_start_date = min([dr.execution_date for dr in drs])
+        clear_end_date = max([dr.execution_date for dr in drs])
+
+        logger.info(
+            f"prepare to clear dag_run for {dag.dag_id}, start_date: {clear_start_date}, end_date: {clear_end_date}, failed_only: {failed_only}"
+        )
+
+        dag.clear(
+            start_date=clear_start_date,
+            end_date=clear_end_date,
+            only_failed=failed_only,
+        )
+
+    @classmethod
+    def rerun_task_run(
+        cls,
+        dag: DAG,
+        run_id: str,
+        node_key: str,
+        min_execution_date: datetime.datetime | None,
+        max_execution_date: datetime.datetime | None,
+        include_upstream: bool,
+        include_downstream: bool,
+        failed_only: bool,
+    ):
+        drs: list[DagRun] = DagRun.find(
+            dag_id=dag.dag_id,
+            run_id=run_id,
+            execution_start_date=min_execution_date,
+            execution_end_date=max_execution_date,
+        )
+        if not drs:
+            logger.info(f"skip rerun, no dag_run found for {dag.dag_id} at {run_id}")
+            return
+        clear_start_date = min([dr.execution_date for dr in drs])
+        clear_end_date = max([dr.execution_date for dr in drs])
+
+        clear_task_ids: list[str] = []
+        for task_id in dag.task_dict.keys():
+            if task_id.startswith(node_key):
+                clear_task_ids.append(task_id)
+
+        expanded_task_ids = set(clear_task_ids)
+        if include_upstream or include_downstream:
+            if include_upstream:
+                for task_id in clear_task_ids:
+                    if task_id in dag.task_dict:
+                        task = dag.task_dict[task_id]
+                        upstream_task_ids = [t.task_id for t in task.upstream_list]
+                        expanded_task_ids.update(upstream_task_ids)
+
+            if include_downstream:
+                for task_id in clear_task_ids:
+                    if task_id in dag.task_dict:
+                        task = dag.task_dict[task_id]
+                        downstream_task_ids = [t.task_id for t in task.downstream_list]
+                        expanded_task_ids.update(downstream_task_ids)
+
+        clear_task_ids = list(expanded_task_ids)
+
+        logger.info(
+            f"prepare to clear task: {dag.dag_id}, {clear_task_ids} start_date: {clear_start_date}, end_date: {clear_end_date}, failed_only: {failed_only}"
+        )
+
+        clear_cnt = dag.clear(
+            task_ids=clear_task_ids,
+            start_date=clear_start_date,
+            end_date=clear_end_date,
+            only_failed=failed_only,
+        )
+        logger.info(f"finish clear task: {dag.dag_id}, {clear_task_ids}, total clear: {clear_cnt} task_instances")
+
+    @classmethod
+    def init_airflow_tables(cls):
+        from airflow.settings import engine
+        from airflow.utils.db import reflect_tables
+
+        def _is_index_exists(session: Session, table_name: str, index_name: str) -> bool:
+            query = f"""
+            SELECT EXISTS (
+                SELECT 1
+                FROM pg_indexes
+                WHERE tablename = {table_name!r}
+                  AND indexname = {index_name!r}
+            )
+            """
+            result = session.execute(query)
+            return result.scalar()
+
+        with create_session() as session:
+            metadata = reflect_tables(tables=["dag_run", "task_instance"], session=session)
+            dag_run = Table("dag_run", metadata, autoload_with=engine)
+            task_instance = Table("task_instance", metadata, autoload_with=engine)
+
+            dag_run_updated_at_idx = Index("ix_dag_run_updated_at", dag_run.c.updated_at)
+            task_instance_updated_at_idx = Index("ix_task_instance_updated_at", task_instance.c.updated_at)
+
+            with engine.connect():
+                if not _is_index_exists(session, "dag_run", "ix_dag_run_updated_at"):
+                    logger.info("start creating index on dag_run.updated_at")
+                    session.execute(CreateIndex(dag_run_updated_at_idx))
+                    logger.info("Created index on dag_run.updated_at")
+                else:
+                    logger.info("Skipped creating index on dag_run.updated_at")
+
+                if not _is_index_exists(session, "task_instance", "ix_task_instance_updated_at"):
+                    logger.info("start creating index on task_instance.updated_at")
+                    session.execute(CreateIndex(task_instance_updated_at_idx))
+                    logger.info("Created index on task_instance.updated_at")
+                else:
+                    logger.info("Skipped creating index on task_instance.updated_at")
+
+    @classmethod
+    def mark_dag_run_success(cls, dag: DAG, run_id: str = None, whole_dag: bool = False):
+        if not run_id:
+            if not whole_dag:
+                logger.info("mark_dag_run need a run_id, skip mark_dag_run")
+                return
+            run_ids = cls._get_dag_run_ids(dag)
+            for run_id in run_ids:
+                cls.mark_dag_run_success(dag, run_id)
+            return
+        logger.info(f"start mark dag run {dag.dag_id} {run_id} to success")
+        set_dag_run_state_to_success(dag=dag, run_id=run_id, commit=True)
+
+    @classmethod
+    def mark_dag_run_failed(cls, dag: DAG, run_id: str = None, whole_dag: bool = False):
+        """
+        will mark un-running tasks to skipped;
+        mark running tasks to failed;
+        keep finished tasks the same.
+        """
+        if not run_id:
+            if not whole_dag:
+                logger.info("mark_dag_run need a run_id, skip mark_dag_run")
+                return
+            run_ids = cls._get_dag_run_ids(dag)
+            for run_id in run_ids:
+                cls.mark_dag_run_failed(dag, run_id)
+            return
+        logger.info(f"start mark dag run {dag.dag_id} {run_id} to success")
+        set_dag_run_state_to_failed(dag=dag, run_id=run_id, commit=True)
+
+    @staticmethod
+    @provide_session
+    def _get_dag_run_ids(dag: DAG, session: Session = None) -> list[str]:
+        query = session.query(DagRun.run_id).filter(DagRun.dag_id == dag.dag_id)
+        return [res[0] for res in query.all()]
+
+    @classmethod
+    @provide_session
+    def delete_whole_dag_dr_ti(cls, dag: DAG, session: Session = None):
+        logger.info(f"start delete whole dag_run and task_instance for {dag.dag_id}")
+        for model in (TaskInstance, DagRun):
+            session.query(model).filter(model.dag_id == dag.dag_id).delete(synchronize_session="fetch")
+        logger.info(f"finish deleted whole dag_run and task_instance for {dag.dag_id}")
+
+    @staticmethod
+    @provide_session
+    def _set_task_run_state(dag: DAG, run_id: str, node_key: str, state: TaskInstanceState, session: Session = None):
+        logger.info(f"start set task_run {dag.dag_id} {run_id} {node_key} to {state}")
+        dag.set_task_instance_state(
+            task_id=node_key,
+            run_id=run_id,
+            state=state,
+            session=session,
+        )
+        logger.info(f"finish set task_run {dag.dag_id} {run_id} {node_key} to {state}")
+
+    @staticmethod
+    def terminate_task_run(dag: DAG, run_id: str, node_key: str):
+        AirflowDbService._set_task_run_state(dag, run_id, node_key, TaskInstanceState.FAILED)

recurvedata/schedulers/airflow_operators.py

@@ -0,0 +1,61 @@
+import logging
+
+from airflow.exceptions import AirflowSkipException
+from airflow.models import TaskInstance
+from airflow.operators.bash import BashOperator
+from airflow.utils.context import Context
+from airflow.utils.task_instance_session import get_current_task_instance_session
+from sqlalchemy.orm.attributes import flag_modified
+
+from recurvedata.executors.utils import read_meta_file
+
+logger = logging.getLogger(__name__)
+
+
+class RecurveBashOperator(BashOperator):
+    def execute(self, context: Context):
+        try:
+            res = super().execute(context)
+            self.update_meta_to_task_instance_executor_config(context)
+            return res
+        except Exception:
+            self.update_meta_to_task_instance_executor_config(context)
+            raise
+
+    @staticmethod
+    def read_meta_file(context: Context) -> dict:
+        return read_meta_file(
+            context["dag"].dag_id, context["ti"].task_id, context["next_execution_date"] or context["execution_date"]
+        )
+
+    def update_meta_to_task_instance_executor_config(self, context: Context):
+        meta = self.read_meta_file(context)
+        if not meta:
+            return
+        logger.debug(f"update_meta_to_task_instance_executor_config: {str(meta)}")
+        session = get_current_task_instance_session()
+        task_instance = TaskInstance.get_task_instance(
+            dag_id=context["dag"].dag_id,
+            task_id=context["ti"].task_id,
+            run_id=context["dag_run"].run_id,
+            map_index=-1,
+            session=session,
+        )
+        if task_instance:
+            task_instance.executor_config.update(meta)
+            flag_modified(task_instance, "executor_config")
+
+
+class SkipSelfBashOperator(BashOperator):
+    ui_color = "#e8f7e4"
+
+    def execute(self, context):
+        raise AirflowSkipException("This task is skipped")
+
+
+class LinkNodeBashOperator(RecurveBashOperator):
+    ui_color = "#8DEEEE"
+
+
+class LinkErrorBashOperator(BashOperator):
+    ui_color = "red"  # not used

recurvedata/schedulers/airflow_plugin.py

@@ -0,0 +1,9 @@
+from airflow.plugins_manager import AirflowPlugin
+
+
+class RecurveAirflowPlugin(AirflowPlugin):
+    name = "recurvedata"
+
+    @classmethod
+    def on_load(cls, *args, **kwargs):
+        import recurvedata.schedulers.debug_celery  # noqa: F401

recurvedata/schedulers/airflow_trigger_dag_patch.py

@@ -0,0 +1,117 @@
+"""
+monkey patch airflow airflow/api/common/trigger_dag.py,
+airflow native _trigger_dag will create data_interval_end = execution_date dag run,
+which will cause plan run error (one data_interval_end may have multiple run_id).
+"""
+
+import json
+import logging
+from datetime import datetime
+
+import airflow.api.common.trigger_dag
+from airflow.exceptions import DagNotFound, DagRunAlreadyExists
+from airflow.models import DagBag, DagRun
+from airflow.models.dag import DAG
+from airflow.timetables.base import DataInterval
+from airflow.timetables.interval import CronDataIntervalTimetable
+from airflow.utils import timezone
+from airflow.utils.state import DagRunState
+from airflow.utils.types import DagRunType
+
+logger = logging.getLogger(__name__)
+
+
+def _recurve_get_next_data_interval(mannual_data_interval: DataInterval, dag: DAG):
+    if not isinstance(dag.timetable, CronDataIntervalTimetable):
+        return mannual_data_interval
+    next_data_interval_end = dag.timetable._get_next(mannual_data_interval.end)
+    return DataInterval(start=mannual_data_interval.end, end=next_data_interval_end)
+
+
+def _recurve_trigger_dag(
+    dag_id: str,
+    dag_bag: DagBag,
+    run_id: str | None = None,
+    conf: dict | str | None = None,
+    execution_date: datetime | None = None,
+    replace_microseconds: bool = True,
+) -> list[DagRun | None]:
+    """
+    Triggers DAG run.
+
+    :param dag_id: DAG ID
+    :param dag_bag: DAG Bag model
+    :param run_id: ID of the dag_run
+    :param conf: configuration
+    :param execution_date: date of execution
+    :param replace_microseconds: whether microseconds should be zeroed
+    :return: list of triggered dags
+    """
+    logger.info("start call _recurve_trigger_dag")
+    dag = dag_bag.get_dag(dag_id)  # prefetch dag if it is stored serialized
+
+    if dag is None or dag_id not in dag_bag.dags:
+        raise DagNotFound(f"Dag id {dag_id} not found")
+
+    execution_date = execution_date or timezone.utcnow()
+
+    if not timezone.is_localized(execution_date):
+        raise ValueError("The execution_date should be localized")
+
+    if replace_microseconds:
+        execution_date = execution_date.replace(microsecond=0)
+
+    if dag.default_args and "start_date" in dag.default_args:
+        min_dag_start_date = dag.default_args["start_date"]
+        if min_dag_start_date and execution_date < min_dag_start_date:
+            raise ValueError(
+                f"The execution_date [{execution_date.isoformat()}] should be >= start_date "
+                f"[{min_dag_start_date.isoformat()}] from DAG's default_args"
+            )
+    logical_date = timezone.coerce_datetime(execution_date)
+
+    data_interval = dag.timetable.infer_manual_data_interval(run_after=logical_date)
+
+    # recurve update start #
+    recurve_external_trigger = True
+    inferred_run_type = DagRunType.from_run_id(run_id)
+    if inferred_run_type == DagRunType.SCHEDULED:
+        new_data_interval = _recurve_get_next_data_interval(data_interval, dag)
+        logger.info(f"adjust data interval: {data_interval} -> {new_data_interval}")
+        data_interval = new_data_interval
+        recurve_external_trigger = False
+    # recurve update end #
+
+    run_id = run_id or dag.timetable.generate_run_id(
+        run_type=DagRunType.MANUAL, logical_date=logical_date, data_interval=data_interval
+    )
+    dag_run = DagRun.find_duplicate(dag_id=dag_id, execution_date=execution_date, run_id=run_id)
+
+    if dag_run:
+        raise DagRunAlreadyExists(dag_run=dag_run, execution_date=execution_date, run_id=run_id)
+
+    run_conf = None
+    if conf:
+        run_conf = conf if isinstance(conf, dict) else json.loads(conf)
+
+    # recurve update start #
+    dag_runs = []
+    dags_to_run = [dag, *dag.subdags]
+    for _dag in dags_to_run:
+        dag_run = _dag.create_dagrun(
+            run_id=run_id,
+            execution_date=execution_date,
+            state=DagRunState.QUEUED,
+            conf=run_conf,
+            external_trigger=recurve_external_trigger,
+            dag_hash=dag_bag.dags_hash.get(dag_id),
+            data_interval=data_interval,
+        )
+        dag_runs.append(dag_run)
+    # recurve update end #
+
+    return dag_runs
+
+
+logger.info("monkey patch airflow.api.common.trigger_dag._trigger_dag")
+airflow.api.common.trigger_dag._trigger_dag = _recurve_trigger_dag

recurvedata/schedulers/base.py

@@ -0,0 +1,99 @@
+import datetime
+import logging
+from dataclasses import dataclass
+from functools import cached_property
+from typing import ClassVar, Optional
+
+import pendulum
+
+from recurvedata.schedulers.client import SchedulerClient
+from recurvedata.schedulers.schemas import JobListResponse
+from recurvedata.utils.dataclass import init_dataclass_from_dict
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class DagSchema:
+    id: int  # recurve job_id
+    name: str  # recurve job_name
+    project_id: int
+    project_name: str
+    workflow_id: int
+    workflow_name: str
+    workflow_version: str
+    graph: list[tuple[str, str]]  # [(upstream_node_key, downstream_node_key),]
+    nodes: list
+    schedule_type: str
+    schedule_interval: str
+    timezone: str
+
+    owner_username: str
+    # scheduler_args: dict
+    start_date: Optional[datetime.datetime] = None
+    end_date: Optional[datetime.datetime] = None
+    scheduler_settings: Optional[dict] = None
+    retries: Optional[int] = None
+    retry_delay: Optional[int] = None
+
+    # attr for modeling pipeline
+    skip_data_tests: bool = False
+
+    @property
+    def job_id(self):
+        return self.id
+
+
+class SchedulerBase(object):
+    DEFAULT_DAG_OWNER: ClassVar[str] = "recurve"
+
+    def __init__(self, sharding_size: int = 1, sharding_key: int = 0):
+        self.sharding_size = sharding_size
+        self.sharding_key = sharding_key
+        self.client: SchedulerClient = self.init_client()
+
+    @cached_property
+    def localtz(self):  # todo: move to CONF
+        return pendulum.timezone("Asia/Shanghai")
+
+    @classmethod
+    def init_client(cls) -> SchedulerClient:
+        return SchedulerClient()
+
+    def list_scheduler_dag(self):
+        """
+        从 sdk 获取符合条件的所有 dag 信息
+        :return:
+        """
+
+        jobs: JobListResponse = self.client.list_jobs(sharding_size=self.sharding_size, sharding_key=self.sharding_key)
+
+        for job in jobs.jobs:
+            dag = init_dataclass_from_dict(DagSchema, job.model_dump())
+            yield dag
+
+    def create_dag(self, row: DagSchema):
+        """
+        生成对应调度器（airflow/...) 的对象
+        :param args:
+        :param kwargs:
+        :return:
+        """
+        try:
+            return self.create_dag_impl(row)
+        except Exception as e:
+            logger.exception(f"failed to generate dag {row.id}, %s", e)
+        return  # todo: add new client api to notify
+
+    def create_dag_impl(self, row: DagSchema):
+        pass
+
+    def execute(self, *args, **kwargs):
+        """
+        入口
+        :param args:
+        :param kwargs:
+        :return:
+        """
+        for row in self.list_scheduler_dag():
+            pass