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

recurvedata/schedulers/cli.py

@@ -0,0 +1,228 @@
+import json
+import logging
+from typing import Any
+
+import dateutil
+import typer
+from airflow.models import DAG
+
+from recurvedata.schedulers.airflow_db_process import AirflowDbService
+from recurvedata.schedulers.schemas import WorkflowNodeDebugDetail
+from recurvedata.schedulers.service import get_job_dag
+from recurvedata.schedulers.task_status import TaskStatusScanner
+from recurvedata.schedulers.utils import clear_task_instance, format_dag_id, init_client
+from recurvedata.utils import init_logging
+from recurvedata.utils._typer import RecurveTyper
+from recurvedata.utils.date_time import to_local_datetime
+
+logger = logging.getLogger(__name__)
+cli = RecurveTyper()
+
+
+def _ensure_dag_exists(job_id: int, raise_error=True) -> DAG | None:
+    dag = get_job_dag(job_id)
+    if not dag:  # job deleted
+        if not raise_error:
+            logger.info(f"dag missing for job {job_id}")
+            return
+        raise ValueError(f"dag not exists: {job_id}")
+    return dag
+
+
+@cli.callback()
+def init():
+    init_logging()
+
+
+@cli.command()
+def update_dag(job_id: int = typer.Option(..., "--job_id")):
+    dag = _ensure_dag_exists(job_id, raise_error=False)
+    AirflowDbService.update_dag(dag)
+
+
+@cli.command()
+def activate_dag(job_id: int = typer.Option(..., "--job_id")):
+    dag = _ensure_dag_exists(job_id, raise_error=False)
+    AirflowDbService.activate_dag(dag)
+
+
+@cli.command()
+def deactivate_dag(job_id: int = typer.Option(..., "--job_id")):
+    dag = _ensure_dag_exists(job_id, raise_error=False)
+    AirflowDbService.deactivate_dag(dag)
+
+
+@cli.command()
+def delete_dag(job_id: int = typer.Option(..., "--job_id"), job_name: str = typer.Option(..., "--job_name")):
+    dag_id = format_dag_id(job_id)
+    AirflowDbService.delete_dag(dag_id, job_name)
+
+
+@cli.command()
+def clear(
+    job_id: int = typer.Option(..., "--job_id"),
+    node_key: str = typer.Option(..., "--node_key"),
+    execution_date: str = typer.Option(..., "--execution_date", callback=to_local_datetime),
+    only_failed: bool = typer.Option(False, "--only_failed"),
+    including_downstream: bool = typer.Option(False, "--including_downstream"),
+):
+    dag = get_job_dag(job_id)
+    clear_task_instance(dag, node_key, execution_date, only_failed, including_downstream)
+
+
+@cli.command()
+def start_workflow_node_debug(
+    workflow_id: int = typer.Option(..., "--workflow_id"),
+    node_key: str = typer.Option(..., "--node_key"),
+    schedule_type: str = typer.Option(..., "--schedule_type"),
+    schedule_interval: str = typer.Option(..., "--schedule_interval"),
+    execution_date: str = typer.Option(..., "--execution_date"),
+    timezone: str = typer.Option(..., "--timezone"),
+):
+    from celery.result import AsyncResult
+
+    from recurvedata.schedulers.debug_celery import debug_node
+
+    celery_kwargs = {
+        "workflow_id": workflow_id,
+        "node_key": node_key,
+        "schedule_type": schedule_type,
+        "schedule_interval": schedule_interval,
+        "execution_date": execution_date,
+        "timezone": timezone,
+    }
+    result: AsyncResult = debug_node.apply_async(kwargs=celery_kwargs)
+    logger.info(f"sent debug_node {celery_kwargs}, celery_id: {result.task_id}")
+    return {
+        "celery_task_id": result.task_id,
+    }
+
+
+@cli.command()
+def abort_workflow_node_debug(
+    workflow_id: int = typer.Option(..., "--workflow_id"),
+    node_key: str = typer.Option(..., "--node_key"),
+    celery_task_id: str = typer.Option(None, "--celery_task_id"),
+):
+    import recurvedata.schedulers.debug_celery
+
+    if not celery_task_id:
+        client = init_client()
+        detail: WorkflowNodeDebugDetail = client.get_workflow_node_debug_detail(
+            workflow_id=workflow_id, node_key=node_key
+        )
+        celery_task_id = detail.celery_task_id
+
+    if not celery_task_id:
+        logger.info("skip revoke_debug, no celery_task_id found")
+        return
+    logger.info(f"start revoke debug: {workflow_id} {node_key} {celery_task_id}")
+    recurvedata.schedulers.debug_celery.revoke_task(celery_task_id)
+    logger.info(f"finish revoke debug: {workflow_id} {node_key} {celery_task_id}")
+
+
+@cli.command()
+def sync_task_status(interval: int = typer.Option(5, "--interval")):
+    scanner = TaskStatusScanner()
+    scanner.run(interval)
+
+
+@cli.command()
+def trigger_job_run(
+    job_id: int = typer.Option(..., "--job_id"),
+    execution_date: str = typer.Option(..., "--execution_date", callback=to_local_datetime),
+    include_past: bool = typer.Option(False, "--include_past"),
+    include_future: bool = typer.Option(False, "--include_future"),
+    run_type: str = typer.Option(None, "--run_type"),
+    conf: str = typer.Option(None, "--conf"),
+):
+    dag = _ensure_dag_exists(job_id)
+    if conf:
+        conf: dict[str, Any] = json.loads(conf)
+    AirflowDbService.trigger_job_run(dag, execution_date, include_past, include_future, run_type, conf)
+
+
+@cli.command()
+def rerun_job_run(
+    job_id: int = typer.Option(..., "--job_id"),
+    run_id: str = typer.Option(None, "--run_id"),
+    min_execution_date: str = typer.Option(None, "--min_execution_date"),
+    max_execution_date: str = typer.Option(None, "--max_execution_date"),
+    failed_only: bool = typer.Option(False, "--failed_only"),
+):
+    dag = _ensure_dag_exists(job_id)
+    if min_execution_date:
+        min_execution_date = dateutil.parser.parse(min_execution_date)
+    if max_execution_date:
+        max_execution_date = dateutil.parser.parse(max_execution_date)
+    AirflowDbService.rerun_job_run(dag, run_id, min_execution_date, max_execution_date, failed_only)
+
+
+@cli.command()
+def rerun_task_run(
+    job_id: int = typer.Option(..., "--job_id"),
+    run_id: str = typer.Option(None, "--run_id"),
+    node_key: str = typer.Option(..., "--node_key"),
+    include_upstream: bool = typer.Option(False, "--include_upstream"),
+    include_downstream: bool = typer.Option(False, "--include_downstream"),
+    min_execution_date: str = typer.Option(None, "--min_execution_date"),
+    max_execution_date: str = typer.Option(None, "--max_execution_date"),
+    failed_only: bool = typer.Option(False, "--failed_only"),
+):
+    dag = _ensure_dag_exists(job_id)
+    if min_execution_date:
+        min_execution_date = dateutil.parser.parse(min_execution_date)
+    if max_execution_date:
+        max_execution_date = dateutil.parser.parse(max_execution_date)
+    AirflowDbService.rerun_task_run(
+        dag=dag,
+        run_id=run_id,
+        node_key=node_key,
+        min_execution_date=min_execution_date,
+        max_execution_date=max_execution_date,
+        include_upstream=include_upstream,
+        include_downstream=include_downstream,
+        failed_only=failed_only,
+    )
+
+
+@cli.command()
+def init_airflow_tables():
+    AirflowDbService.init_airflow_tables()
+
+
+@cli.command()
+def stop_dev_run(job_id: int = typer.Option(..., "--job_id")):
+    logger.info(f"start stop dev run job_id: {job_id}")
+    dag = _ensure_dag_exists(job_id)
+
+    AirflowDbService.mark_dag_run_failed(dag, whole_dag=True)
+
+
+@cli.command()
+def start_dev_run(
+    job_id: int = typer.Option(..., "--job_id"),
+    execution_date: str = typer.Option(..., "--execution_date", callback=to_local_datetime),
+):
+    logger.info(f"start stop dev run job_id: {job_id}")
+    dag = _ensure_dag_exists(job_id)
+
+    AirflowDbService.update_dag(dag)
+    # AirflowDbService.mark_dag_run_failed(dag, whole_dag=True)
+    AirflowDbService.delete_whole_dag_dr_ti(dag)
+    AirflowDbService.activate_dag(dag)
+    AirflowDbService.trigger_job_run(dag, execution_date, False, False, "manual")
+
+
+@cli.command()
+def terminate_task_run(
+    job_id: int = typer.Option(..., "--job_id"),
+    run_id: str = typer.Option(..., "--run_id"),
+    node_key: str = typer.Option(..., "--node_key"),
+):
+    dag = _ensure_dag_exists(job_id)
+    AirflowDbService.terminate_task_run(dag, run_id, node_key)
+
+
+if __name__ == "__main__":
+    cli()

recurvedata/schedulers/client.py

@@ -0,0 +1,56 @@
+from recurvedata.client import Client
+from recurvedata.schedulers.schemas import JobListResponse, TaskStatusCursor, WorkflowNodeDebugDetail
+from recurvedata.utils import get_env_id
+
+
+class SchedulerClient(Client):
+    def list_jobs(self, sharding_size: int = None, sharding_key: int = None) -> JobListResponse:
+        if not sharding_size:
+            sharding_size = 1
+            sharding_key = 0
+
+        params = {
+            "env_id": get_env_id(),
+            "sharding_key": sharding_key,
+            "sharding_size": sharding_size,
+        }
+
+        return self.request("GET", path="/api/scheduler/jobs", response_model_class=JobListResponse, params=params)
+
+    def get_task_status_cursor(self) -> TaskStatusCursor:
+        params = {"env_id": get_env_id()}
+        return self.request(
+            "GET", path="/api/scheduler/task-status-cursor", response_model_class=TaskStatusCursor, params=params
+        )
+
+    def sync_task_status(self, job_runs: list[dict] | None = None, task_runs: list[dict] | None = None):
+        params = {"env_id": get_env_id()}
+        payload = {
+            "job_runs": job_runs,
+            "task_runs": task_runs,
+        }
+        return self.request("POST", path="/api/scheduler/sync-task-status", params=params, json=payload)
+
+    def get_workflow_node_debug_detail(self, workflow_id: int, node_key: str) -> WorkflowNodeDebugDetail:
+        params = {
+            "env_id": get_env_id(),
+            "workflow_id": workflow_id,
+            "node_key": node_key,
+        }
+        return self.request(
+            "GET",
+            path="/api/scheduler/workflow-node-debug-detail",
+            response_model_class=WorkflowNodeDebugDetail,
+            params=params,
+        )
+
+    def on_job_run_finished(self, job_run_result: dict):
+        params = {"env_id": get_env_id()}
+        payload = {
+            "job_id": job_run_result["job_id"],
+            "run_id": job_run_result["run_id"],
+            "task_info_map": job_run_result["task_info_map"],
+            "state": job_run_result["state"],
+            "data_interval_end": job_run_result["data_interval_end"],
+        }
+        return self.request("POST", path="/api/scheduler/on-job-run-finished", params=params, json=payload)

recurvedata/schedulers/consts.py

@@ -0,0 +1,52 @@
+import re
+from enum import Enum
+
+
+class OperatorEnum(str, Enum):
+    SQLOperator = "SQLOperator"
+    TransferOperator = "TransferOperator"
+    PythonOperator = "PythonOperator"
+    SparkOperator = "SparkOperator"
+    NotifyOperator = "NotifyOperator"
+    LinkOperator = "LinkOperator"
+
+
+WORK_DIR = "/opt/airflow"  # todo: use /opt/recurve
+
+
+class Operator(str, Enum):  # todo
+    SQLOperator = "SQLOperator"
+    TransferOperator = "TransferOperator"
+    PythonOperator = "PythonOperator"
+    SparkOperator = "SparkOperator"
+    NotifyOperator = "NotifyOperator"
+    LinkOperator = "LinkOperator"
+
+
+class ScheduleType(str, Enum):
+    crontab = "crontab"
+    customization = "customization"  # 快捷设置
+    manual = "manual"  # 手动触发
+
+
+SYSTEM_SYNC_STATUS_DAG_ID = "system_sync_status"
+
+
+def format_recurve_env_key(key: str) -> str:
+    return f"RECURVE__{key.upper()}"
+
+
+def get_dag_file_loc(job_id: int) -> str:
+    # todo: configuration
+    idx = job_id % 7
+    return f"/opt/airflow/dags/autogen_sharding_{idx}.py"
+
+
+DEFAULT_RETRY_NUMBER = 2
+DEFAULT_RETRY_DELAY = 60 * 5  # 5 minutes
+
+
+def is_dev_run_job(job_name: str) -> bool:
+    pattern = r"dev_run_.*_\d+_\d+"
+    match = re.match(pattern, job_name)
+    return match is not None

recurvedata/schedulers/debug_celery.py

@@ -0,0 +1,62 @@
+import logging
+import os
+
+from airflow.providers.celery.executors.celery_executor import app as celery_app
+from celery import Task
+
+from recurvedata.executors.client import ExecutorClient
+from recurvedata.operators.config import CONF
+from recurvedata.utils.mp import run_subprocess
+
+logger = logging.getLogger(__name__)
+
+
+@celery_app.task(bind=True)
+def debug_node(
+    self: Task,
+    workflow_id: int,
+    node_key: str,
+    schedule_type: str,
+    schedule_interval: str,
+    execution_date: str,
+    timezone: str,
+):
+    task_id = self.request.id
+    logging.info(
+        f"start {task_id} {workflow_id} {node_key} {schedule_type} {schedule_interval} {execution_date} {timezone}"
+    )
+
+    executor_client = ExecutorClient()
+
+    executor_client.debug_start(workflow_id, node_key, task_id)
+    try:
+        run_subprocess(
+            [
+                os.path.join(CONF.RECURVE_EXECUTOR_PYENV_BIN_PATH, "recurve_executor"),
+                "debug",
+                "--workflow_id",
+                f"{workflow_id}",
+                "--node_key",
+                f"{node_key}",
+                "--schedule_type",
+                schedule_type,
+                "--schedule_interval",
+                schedule_interval,
+                "--execution_date",
+                execution_date,
+                "--timezone",
+                timezone,
+                "--celery_task_id",
+                task_id,
+            ],
+            env=os.environ.copy(),
+        )
+        is_success = True
+    except Exception as e:
+        logger.exception(f"{workflow_id} {node_key} {execution_date} debug failed, err: {e}")
+        is_success = False
+    executor_client.debug_end(workflow_id, node_key, task_id, is_success)
+
+
+def revoke_task(task_id: str = None, terminate=True):
+    return celery_app.control.revoke(task_id, terminate=terminate)

recurvedata/schedulers/model.py

@@ -0,0 +1,63 @@
+from dataclasses import dataclass, field
+from typing import Optional
+
+from recurvedata.operators.models import DagBase, NodeBase
+
+
+@dataclass
+class SchedulerDag(DagBase):
+    workflow_version: str | None = None
+
+
+@dataclass
+class SchedulerNode(NodeBase):
+    """
+    调度器的 Node 对象
+    """
+
+    operator: str
+
+    scheduler_settings: Optional[dict] = None
+    skip_self: Optional[bool] = None
+    skip_downstream: Optional[bool] = None
+    latest_only: Optional[bool] = None
+
+
+@dataclass
+class LinkNodeItem:
+    """
+    the node linked by LinkOperator
+    """
+
+    link_wf_id: int
+    link_wf_version: str
+    link_node_id: int
+    link_node_name: str
+    link_node_key: str
+    link_latest_only: bool
+    link_operator: str
+    link_skip_downstream: bool
+    link_skip_self: bool
+    link_scheduler_settings: dict = None
+    link_config: dict = None  # used in CustomAirflowOperator
+    node_id: int = None
+    plan_id: int = None
+
+    @property
+    def config(self):
+        # for CustomAirflowOperator
+        return self.link_config
+
+
+@dataclass
+class LinkWorkflowItem:
+    """
+    LinkOperator - link workflow
+    """
+
+    node_id: int
+    link_wf_id: int
+    link_wf_name: str
+    link_wf_version: str
+    link_graph: list[tuple[str, str]] = field(default_factory=list)  # [(upstream_node_key, downstream_node_key),]
+    link_nodes: list[LinkNodeItem] = None

recurvedata/schedulers/schemas.py

@@ -0,0 +1,97 @@
+import datetime
+
+from pydantic import BaseModel
+
+from recurvedata.consts import Operator, ScheduleType
+
+
+class SchedulerLinkWorkflowNodeItem(BaseModel):
+    link_wf_id: int
+    link_wf_version: str
+    link_node_id: int
+    link_node_key: str
+    link_node_name: str
+    link_operator: Operator
+    link_skip_self: bool
+    link_skip_downstream: bool
+    link_latest_only: bool
+    link_scheduler_settings: dict | None
+    plan_id: int | None = None
+
+
+class SchedulerLinkNodeItem(BaseModel):
+    node_id: int
+    link_wf_id: int
+    link_wf_version: str
+    link_node_id: int
+    link_node_key: str
+    link_node_name: str
+    link_operator: Operator
+    link_skip_self: bool
+    link_skip_downstream: bool
+    link_latest_only: bool
+    link_scheduler_settings: dict | None
+
+
+class SchedulerNodeItem(BaseModel):
+    id: int
+    node_key: str
+    name: str
+    operator: Operator
+    scheduler_settings: dict | None
+    skip_self: bool
+    skip_downstream: bool
+    latest_only: bool
+
+
+class JobItem(BaseModel):
+    id: int
+    name: str
+    schedule_type: ScheduleType
+    schedule_interval: str
+    timezone: str | None
+    workflow_version: str
+    start_date: datetime.datetime | None
+    end_date: datetime.datetime | None
+    scheduler_settings: dict | None
+    owner_username: str
+
+    nodes: list[SchedulerNodeItem]
+
+    graph: list
+    project_id: int = None
+    project_name: str = None
+    workflow_id: int = None
+    workflow_name: str = None
+
+    skip_data_tests: bool = False
+    retries: int | None = None
+    retry_delay: int | None = None
+
+
+class SchedulerLinkWorkflowItem(BaseModel):
+    node_id: int
+    link_wf_id: int
+    link_wf_name: str
+    link_wf_version: str
+    link_nodes: list[SchedulerLinkWorkflowNodeItem]
+    link_graph: list
+
+
+class JobListResponse(BaseModel):
+    jobs: list[JobItem]
+    link_nodes: list[SchedulerLinkNodeItem]
+    link_workflows: list[SchedulerLinkWorkflowItem]
+
+
+class TaskStatusCursor(BaseModel):
+    job_run: datetime.datetime | None = None
+    task_run: datetime.datetime | None = None
+    limit: int = 30
+    sliding_time: int = 1
+    unfinished: dict | None = None
+
+
+class WorkflowNodeDebugDetail(BaseModel):
+    celery_task_id: str | None = None
+    state: str | None = None

recurvedata/schedulers/service.py

@@ -0,0 +1,20 @@
+import logging
+from typing import Optional
+
+from airflow.models import DAG
+
+from recurvedata.schedulers.airflow import AirflowScheduler
+from recurvedata.schedulers.consts import get_dag_file_loc
+
+logger = logging.getLogger(__name__)
+
+
+def get_job_dag(job_id: int) -> Optional["DAG"]:
+    scheduler = AirflowScheduler(sharding_size=job_id, sharding_key=0)
+    dag_dct = scheduler.execute()
+    dag_ids = list(dag_dct.keys())
+    if not dag_ids:
+        return
+    dag = dag_dct[dag_ids[0]]
+    dag.fileloc = get_dag_file_loc(job_id)
+    return dag

recurvedata/schedulers/system_dags.py

@@ -0,0 +1,59 @@
+import os
+from datetime import datetime, timedelta
+
+from airflow import DAG
+from airflow.operators.bash import BashOperator
+
+from recurvedata.schedulers.consts import SYSTEM_SYNC_STATUS_DAG_ID
+
+
+def create_system_dags():
+    return [
+        create_sync_status_dag(),
+    ]
+
+
+def _prepare_bash_env():
+    dct = {}
+    for key, val in os.environ.items():
+        if key.startswith("RECURVE__"):
+            dct[key] = val
+        elif key.startswith("AIRFLOW"):
+            dct[key] = val
+        elif key in (
+            "PATH",
+            "PYENV_ROOT",
+        ):
+            dct[key] = os.environ[key]
+    return dct
+
+
+def create_sync_status_dag():
+    start_date = datetime(2024, 8, 5)
+    default_args = {
+        "depends_on_past": False,
+        "retries": 150,
+        "retry_delay": timedelta(seconds=10),
+        "priority_weight": 100,
+        "retry_exponential_backoff": True,
+        "max_retry_delay": timedelta(seconds=30),
+    }
+    dag = DAG(
+        SYSTEM_SYNC_STATUS_DAG_ID,
+        default_args=default_args,
+        description="A DAG to sync db status",
+        schedule_interval="0 */6 * * *",  # Run every 6 hours
+        start_date=start_date,
+        catchup=False,
+        dagrun_timeout=timedelta(minutes=60 * 6),
+        max_active_runs=1,  # todo: retry may delay the future dag_run
+        is_paused_upon_creation=False,
+    )
+
+    BashOperator(
+        task_id="sync_status",
+        bash_command="recurve_scheduler sync-task-status --interval=5",
+        dag=dag,
+        env=_prepare_bash_env(),
+    )
+    return dag