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

recurvedata/schedulers/task_status.py

@@ -0,0 +1,279 @@
+import datetime
+import logging
+import signal
+import sys
+import time
+
+import pytz
+from airflow.models import DagRun, TaskInstance
+from airflow.models.taskinstancehistory import TaskInstanceHistory
+from airflow.utils.session import create_session
+from sqlalchemy import and_, func, tuple_
+from sqlalchemy.orm import joinedload
+
+from recurvedata.config import AgentConfig
+from recurvedata.schedulers.client import SchedulerClient
+from recurvedata.schedulers.consts import SYSTEM_SYNC_STATUS_DAG_ID
+
+logger = logging.getLogger(__name__)
+
+
+class TaskStatusScanner:
+    def __init__(self):
+        config = AgentConfig.load()
+        if config.request_timeout < 30:
+            config.request_timeout = 30
+        self.client = SchedulerClient(config)
+        self._running = False
+
+    def run(self, interval: int):
+        def signal_handler(_sig, _frame):
+            self._running = False
+
+        def handle_sigterm(_sig, _frame):
+            self._running = False
+            sys.exit(0)
+
+        signal.signal(signal.SIGINT, signal_handler)
+        signal.signal(signal.SIGTERM, handle_sigterm)
+
+        self._running = True
+        step = interval
+
+        while self._running:
+            if step >= interval:
+                status_cursor = self.client.get_task_status_cursor()
+
+                job_runs = {}
+                task_runs = {}
+
+                def _collect_job_runs(_job_runs):
+                    for jr in _job_runs:
+                        job_runs[(str(jr["job_id"]), jr["run_id"])] = jr
+
+                def _collect_task_runs(_task_runs):
+                    for tr in _task_runs:
+                        task_runs[(tr["job_id"], tr["run_id"], tr["node_key"])] = tr
+
+                _job_runs = self.scan_dag_runs(
+                    self._localize_time(status_cursor.job_run),
+                    status_cursor.limit,
+                    sliding_time=status_cursor.sliding_time,
+                )
+                _collect_job_runs(_job_runs)
+
+                _task_runs, _job_runs = self.scan_task_instances(
+                    self._localize_time(status_cursor.task_run),
+                    status_cursor.limit,
+                    sliding_time=status_cursor.sliding_time,
+                )
+                _collect_job_runs(_job_runs)
+                _collect_task_runs(_task_runs)
+
+                _task_runs, _job_runs = self.scan_unfinished_task_instances(status_cursor.unfinished)
+                _collect_job_runs(_job_runs)
+                _collect_task_runs(_task_runs)
+
+                logger.info(f"number of job runs: {len(job_runs)}")
+                logger.info(f"number of task runs: {len(task_runs)}")
+
+                # get actual job start time
+                filters = list(job_runs.keys())
+
+                if filters:
+                    with create_session() as session:
+                        query = (
+                            session.query(
+                                func.min(TaskInstance.start_date).label("start_time"),
+                                TaskInstance.dag_id,
+                                TaskInstance.run_id,
+                                TaskInstance.try_number,
+                            )
+                            .filter(tuple_(TaskInstance.dag_id, TaskInstance.run_id).in_(filters))
+                            .group_by(TaskInstance.dag_id, TaskInstance.run_id, TaskInstance.try_number)
+                        )
+
+                        # retried job runs
+                        history_filters = []
+                        for row in query.all():
+                            if row.try_number > 1:
+                                history_filters.append((str(row.dag_id), row.run_id))
+                                continue
+
+                            job_runs[(row.dag_id), row.run_id]["start_time"] = (
+                                row.start_time and row.start_time.isoformat()
+                            )
+
+                        if history_filters:
+                            query = (
+                                session.query(
+                                    func.min(TaskInstanceHistory.start_date).label("start_time"),
+                                    TaskInstanceHistory.dag_id,
+                                    TaskInstanceHistory.run_id,
+                                )
+                                .filter(
+                                    tuple_(TaskInstanceHistory.dag_id, TaskInstanceHistory.run_id).in_(history_filters)
+                                )
+                                .group_by(TaskInstanceHistory.dag_id, TaskInstanceHistory.run_id)
+                            )
+
+                            for row in query.all():
+                                job_runs[(row.dag_id, row.run_id)]["start_time"] = (
+                                    row.start_time and row.start_time.isoformat()
+                                )
+
+                self.client.sync_task_status(job_runs=list(job_runs.values()), task_runs=list(task_runs.values()))
+                step = 0
+            time.sleep(1.0)
+            step += 1
+
+    def _localize_time(self, time: datetime.datetime | None):
+        if time is not None and time.tzinfo is None:
+            timezone = pytz.timezone("UTC")
+            time = timezone.localize(time)
+        return time
+
+    def _parse_job_id(self, dag_id: str):
+        return int(dag_id.split(".")[-1])
+
+    def _sliding_time_query(self, session, model, last_updated_time, sliding_time):
+        query = session.query(model).where(model.dag_id != SYSTEM_SYNC_STATUS_DAG_ID)
+        query = query.where(
+            and_(
+                model.updated_at >= (last_updated_time - datetime.timedelta(seconds=sliding_time)),
+                model.updated_at < last_updated_time,
+            )
+        )
+        return query.order_by(model.updated_at.asc())
+
+    def _limit_query(self, session, model, last_updated_time, limit):
+        query = session.query(model).where(model.dag_id != SYSTEM_SYNC_STATUS_DAG_ID)
+        if last_updated_time is not None:
+            query = query.where(model.updated_at >= last_updated_time)
+        return query.order_by(model.updated_at.asc()).limit(limit)
+
+    def _format_job_run(self, dr: DagRun, workflow_version: str):
+        return dict(
+            job_id=self._parse_job_id(dr.dag_id),
+            run_id=dr.run_id,
+            state=dr.state,
+            start_time=dr.start_date and dr.start_date.isoformat(),
+            end_time=dr.end_date and dr.end_date.isoformat(),
+            execution_date=dr.execution_date and dr.execution_date.isoformat(),
+            workflow_version=workflow_version,
+            airflow_updated_at=dr.updated_at and dr.updated_at.isoformat(),
+            run_type=None,
+            data_interval_end=dr.data_interval_end and dr.data_interval_end.isoformat(),
+        )
+
+    def scan_dag_runs(self, last_updated_time: datetime.datetime | None, limit: int, sliding_time: int = 0):
+        job_runs = []
+        workflow_version_map = {}
+
+        with create_session() as session:
+            dag_runs: list[DagRun] = []
+            if last_updated_time and sliding_time > 0:
+                dag_runs.extend(self._sliding_time_query(session, DagRun, last_updated_time, sliding_time).all())
+            dag_runs.extend(self._limit_query(session, DagRun, last_updated_time, limit).all())
+
+            for dr in dag_runs:
+                workflow_version = workflow_version_map.get((dr.dag_id, dr.run_id))
+                if workflow_version is None:
+                    ti = (
+                        session.query(TaskInstance)
+                        .filter(TaskInstance.dag_id == dr.dag_id, TaskInstance.run_id == dr.run_id)
+                        .first()
+                    )
+                    workflow_version = ti and ti.executor_config.get("workflow_version")
+                    if workflow_version is not None:
+                        workflow_version_map[(dr.dag_id, dr.run_id)] = workflow_version
+                job_runs.append(self._format_job_run(dr, workflow_version))
+        return job_runs
+
+    def _format_task_run(self, ti: TaskInstance):
+        return dict(
+            job_id=self._parse_job_id(ti.dag_id),
+            run_id=ti.run_id,
+            node_key=ti.task_id,
+            state=ti.state,
+            try_number=ti._try_number,
+            start_time=ti.start_date and ti.start_date.isoformat(),
+            end_time=ti.end_date and ti.end_date.isoformat(),
+            execution_date=ti.execution_date and ti.execution_date.isoformat(),
+            workflow_version=ti.executor_config.get("workflow_version"),
+            operator=ti.executor_config.get("operator"),
+            task=ti.executor_config.get("task"),
+            link_workflow_id=ti.executor_config.get("link_workflow_id"),
+            link_workflow_version=ti.executor_config.get("link_workflow_version"),
+            airflow_updated_at=ti.updated_at and ti.updated_at.isoformat(),
+        )
+
+    def scan_task_instances(self, last_updated_time: datetime.datetime | None, limit: int, sliding_time: int = 0):
+        dag_runs = {}
+        task_runs = []
+
+        with create_session() as session:
+            tis: list[TaskInstance] = []
+            if last_updated_time and sliding_time > 0:
+                tis.extend(
+                    self._sliding_time_query(session, TaskInstance, last_updated_time, sliding_time)
+                    .options(joinedload(TaskInstance.dag_run))
+                    .all()
+                )
+            tis.extend(
+                self._limit_query(session, TaskInstance, last_updated_time, limit)
+                .options(joinedload(TaskInstance.dag_run))
+                .all()
+            )
+
+            for ti in tis:
+                dag_runs[(ti.dag_run.dag_id, ti.dag_run.run_id)] = (
+                    ti.dag_run,
+                    ti.executor_config.get("workflow_version"),
+                )
+                task_runs.append(self._format_task_run(ti))
+        return task_runs, [self._format_job_run(dr, workflow_version) for dr, workflow_version in dag_runs.values()]
+
+    def scan_unfinished_task_instances(self, data: dict | None):
+        if not data:
+            return [], []
+
+        dag_ids = set()
+        task_ids = set()
+        run_ids = set()
+        for dag_id, item in data.items():
+            dag_ids.add(dag_id)
+            for task_id, _run_ids in item.items():
+                task_ids.add(task_id)
+                for run_id in _run_ids:
+                    run_ids.add(run_id)
+
+        dag_runs = {}
+        task_runs = []
+
+        with create_session() as session:
+            criterion = []
+            if dag_ids:
+                criterion.append(TaskInstance.dag_id.in_(dag_ids))
+            if task_ids:
+                criterion.append(TaskInstance.task_id.in_(task_ids))
+            if run_ids:
+                criterion.append(TaskInstance.run_id.in_(run_ids))
+
+            tis: list[TaskInstance] = (
+                session.query(TaskInstance).where(*criterion).options(joinedload(TaskInstance.dag_run)).all()
+            )
+
+            for ti in tis:
+                if (
+                    ti.dag_id not in data
+                    or ti.task_id not in data[ti.dag_id]
+                    or ti.run_id not in data[ti.dag_id][ti.task_id]
+                ):
+                    continue
+                dag_runs[(ti.dag_run.dag_id, ti.dag_run.run_id)] = (
+                    ti.dag_run,
+                    ti.executor_config.get("workflow_version"),
+                )
+                task_runs.append(self._format_task_run(ti))
+        return task_runs, [self._format_job_run(dr, workflow_version) for dr, workflow_version in dag_runs.values()]

recurvedata/schedulers/utils.py

@@ -0,0 +1,73 @@
+import datetime
+import json
+import logging
+from typing import Generator
+
+from airflow.models import DAG, BaseOperator
+from airflow.utils.session import create_session
+from slugify import slugify
+
+from recurvedata.schedulers.client import SchedulerClient
+
+logger = logging.getLogger(__name__)
+
+
+def get_tasks(dag: DAG, recurve_node_key: str = None) -> Generator[BaseOperator, None, None]:
+    if recurve_node_key:
+        for task in dag.task_dict.values():
+            if task.doc_json and json.loads(task.doc_json).get("recurve_node_key") == recurve_node_key:
+                yield task
+
+
+def clear_task_instance(
+    dag: DAG,
+    recurve_node_key: str,
+    execution_date: datetime.datetime,
+    only_failed: bool = False,
+    including_downstream: bool = False,
+):
+    clear_task_ids: list[str] = []
+    for task in get_tasks(dag, recurve_node_key):
+        clear_task_ids.append(task.task_id)
+    airflow_execution_date = dag.previous_schedule(execution_date)  # todo: timezone
+    with create_session() as session:
+        dag = dag.partial_subset(task_ids_or_regex=clear_task_ids, include_downstream=including_downstream)
+        clear_task_ids = [tid for tid in dag.task_dict]
+        logger.info(f"prepare to clear dag: {dag.dag_id}, {clear_task_ids} execution_date: {airflow_execution_date}")
+
+        clear_cnt = dag.clear(
+            task_ids=clear_task_ids,
+            start_date=airflow_execution_date,
+            end_date=airflow_execution_date,
+            only_failed=only_failed,
+            session=session,
+        )
+        logger.info(
+            f"finish clear dag: {dag.dag_id}, {clear_task_ids} execution_date: {airflow_execution_date}, total clear: {clear_cnt} task_instances"
+        )
+
+
+def slugify_text(s: str) -> str:
+    """A simple wrapper to python-slugify, using custom regex_pattern to keep `.` and `_` as is
+
+    >>> slugify_text('我是谁')
+    'wo-shi-shui'
+    >>> slugify_text('load_fact_user_stats')
+    'load_fact_user_stats'
+    >>> slugify_text('tidb prepare category tables')
+    'tidb-prepare-category-tables'
+    >>> slugify_text('estimate daily deal 2017.10.20')
+    'estimate-daily-deal-2017.10.20'
+    """
+    return slugify(s, regex_pattern=r"[^-a-zA-Z0-9\._]+")
+
+
+def format_dag_id(job_id: int) -> str:
+    """
+    please do not adjust this function
+    """
+    return str(job_id)
+
+
+def init_client() -> SchedulerClient:
+    return SchedulerClient()

recurvedata/schema/field.py

@@ -0,0 +1,88 @@
+import dataclasses
+import datetime
+from typing import Any, Callable, Optional
+
+import dateutil.parser
+
+from recurvedata.schema.types import DataType
+from recurvedata.utils import json
+from recurvedata.utils.registry import GenericRegistry
+
+_registry = GenericRegistry[DataType, Callable[[str], Any]]()
+
+
+@_registry.add(DataType.INT8, DataType.INT16, DataType.INT32, DataType.INT64)
+def _(value: str) -> int:
+    if value == "":
+        return 0
+    return int(value)
+
+
+@_registry.add(DataType.FLOAT32, DataType.FLOAT64)
+def _(value: str) -> float:
+    if value == "":
+        return 0.0
+    return float(value)
+
+
+@_registry.add(DataType.BOOLEAN)
+def _(value: str) -> bool:
+    if value.lower() in ("", "0", "false"):
+        return False
+    return True
+
+
+@_registry.add(DataType.DATETIME)
+def _(value: str) -> Optional[datetime.datetime]:
+    if value == "":
+        return None
+    return dateutil.parser.parse(value)
+
+
+@_registry.add(DataType.DATE)
+def _(value: str) -> Optional[datetime.date]:
+    if value == "":
+        return None
+    return dateutil.parser.parse(value).date()
+
+
+@_registry.add(DataType.JSON)
+def _(value: str) -> Any:
+    if value in ("",):
+        # 正常情况下不会有 ''，很可能是从 CSV 文件读到了空字符，当作 None 处理
+        return None
+    return json.loads(value)
+
+
+@dataclasses.dataclass
+class Field:
+    name: str
+    type: DataType
+    size: int = None
+    comment: str = None
+    extra: dict = None
+
+    def __post_init__(self):
+        self._cast_func: Callable[[str], Any] = _registry.get(self.type, lambda x: x)
+
+    def cast(self, value: Optional[str]) -> Any:
+        if value is None:
+            return None
+        if value == "NULL":
+            return None
+        return self._cast_func(value)
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "name": self.name,
+            "type": self.type,
+            "size": self.size,
+            "comment": self.comment,
+            "extra": self.extra,
+        }
+
+    def __str__(self):
+        return f'<Field ("{self.name}", "{self.type}")>'
+
+    def __repr__(self):
+        return f'<Field ("{self.name}", "{self.type}")>'

recurvedata/schema/schema.py

@@ -0,0 +1,55 @@
+import dataclasses
+from typing import Any
+
+from recurvedata.schema.field import Field
+from recurvedata.schema.types import DataType
+from recurvedata.utils import json
+
+
+@dataclasses.dataclass
+class Schema:
+    fields: list[Field] = dataclasses.field(default_factory=list)
+
+    def add_field(self, field: Field):
+        if field.name in self.field_names:
+            raise ValueError(f"Field name {field.name} already exists")
+        self.fields.append(field)
+
+    def add_field_by_attrs(
+        self,
+        name: str,
+        type: DataType,
+        size: int = None,
+        comment: str = None,
+        extra: dict = None,
+    ):
+        self.add_field(Field(name, type, size, comment, extra))
+
+    def remove_field(self, name: str):
+        self.fields = [x for x in self.fields if x.name != name]
+
+    def keep_fields(self, names: list[str]):
+        self.fields = [x for x in self.fields if x.name in names]
+
+    @property
+    def field_names(self) -> list[str]:
+        return [x.name for x in self.fields]
+
+    def __iter__(self):
+        return iter(self.fields)
+
+    def to_list(self) -> list[dict[str, Any]]:
+        return [x.to_dict() for x in self.fields]
+
+    def to_json(self, **kwargs) -> str:
+        return json.dumps(self.to_list(), **kwargs)
+
+    def dump(self, filename: str):
+        with open(filename, "w") as f:
+            f.write(self.to_json(indent=2))
+
+    @classmethod
+    def load(cls, filename: str) -> "Schema":
+        with open(filename) as f:
+            data = json.loads(f.read())
+        return cls([Field(**item) for item in data])

recurvedata/schema/types.py

@@ -0,0 +1,17 @@
+import enum
+
+
+class DataType(str, enum.Enum):
+    INT8 = "INT8"  # 1-byte (8-bit) signed integers
+    INT16 = "INT16"  # 2-byte (16-bit) signed integers
+    INT32 = "INT32"  # 4-byte (32-bit) signed integers
+    INT64 = "INT64"  # 8-byte (64-bit) signed integers
+    FLOAT32 = "FLOAT32"  # 4-byte (32-bit) single-precision floating
+    FLOAT64 = "FLOAT64"  # 8-byte (64-bit) double-precision floating
+    BOOLEAN = "BOOLEAN"
+
+    DATETIME = "DATETIME"
+    DATE = "DATE"
+
+    STRING = "STRING"
+    JSON = "JSON"

recurvedata/server/app.py

@@ -0,0 +1,7 @@
+from recurvedata.executors.utils import patch_pandas_mysql_connector_cext_missing
+
+patch_pandas_mysql_connector_cext_missing()
+
+from recurvedata.server.main import create_app  # noqa: E402
+
+app = create_app()

recurvedata/server/connector/api.py

@@ -0,0 +1,79 @@
+from fastapi import APIRouter
+from loguru import logger
+
+from recurvedata.core.tracing import Tracing
+from recurvedata.executors.schemas import (
+    ColumnListPayload,
+    ConnectionRuntimePayload,
+    ResponseModel,
+    TableListPayload,
+    TestConnectionPayload,
+)
+from recurvedata.executors.service.connector import ConnectionService
+from recurvedata.executors.utils import run_with_result_handling_v2
+from recurvedata.server.connector.schemas import (
+    ListColumnsResponse,
+    ListDatabasesResponse,
+    ListFullDatabasesResponse,
+    ListTablesResponse,
+    TestConnectionResponse,
+)
+
+tracer = Tracing()
+router = APIRouter()
+
+
+@router.post("/test-connection")
+@tracer.create_span(sampling_rate=0.1)
+async def test_connection(*, payload: TestConnectionPayload) -> TestConnectionResponse:
+    logger.info(f"test_connection: {payload.connection_type}")
+
+    res: ResponseModel = await run_with_result_handling_v2(
+        ConnectionService.test_connection, payload.timeout, payload.connection_type, payload.config
+    )
+    logger.info("finish test_connection")
+    return TestConnectionResponse.model_validate(res.model_dump())
+
+
+@router.post("/list-databases")
+@tracer.create_span(sampling_rate=0.1)
+async def list_databases(*, payload: ConnectionRuntimePayload) -> ListDatabasesResponse:
+    logger.info(f"list_databases: {payload.connection_type}")
+    res: ResponseModel = await run_with_result_handling_v2(
+        ConnectionService.list_databases, None, payload.connection_type, payload.config
+    )
+    logger.info("finish list_databases")
+    return ListDatabasesResponse.model_validate(res.model_dump())
+
+
+@router.post("/list-tables")
+@tracer.create_span(sampling_rate=0.1)
+async def list_tables(*, payload: TableListPayload) -> ListTablesResponse:
+    logger.info(f"list_tables: {payload.connection_type} {payload.database}")
+    res: ResponseModel = await run_with_result_handling_v2(
+        ConnectionService.list_tables, None, payload.connection_type, payload.config, payload.database
+    )
+    logger.info("finish list_tables")
+    return ListTablesResponse.model_validate(res.model_dump())
+
+
+@router.post("/list-columns")
+@tracer.create_span(sampling_rate=0.1)
+async def list_columns(*, payload: ColumnListPayload) -> ListColumnsResponse:
+    logger.info(f"list_columns: {payload.connection_type} {payload.database} {payload.table}")
+    res: ResponseModel = await run_with_result_handling_v2(
+        ConnectionService.list_columns, None, payload.connection_type, payload.config, payload.database, payload.table
+    )
+    logger.info("finish list_columns")
+    return ListColumnsResponse.model_validate(res.model_dump())
+
+
+@router.post("/list-full-databases")
+@tracer.create_span(sampling_rate=0.1)
+async def list_full_databases(*, payload: ConnectionRuntimePayload) -> ListFullDatabasesResponse:
+    logger.info(f"list_full_databases: {payload.connection_type}")
+    res: ResponseModel = await run_with_result_handling_v2(
+        ConnectionService.list_full_databases, None, payload.connection_type, payload.config
+    )
+    logger.info("finish list_full_databases")
+    return ListFullDatabasesResponse.model_validate(res.model_dump())

recurvedata/server/connector/schemas.py

@@ -0,0 +1,28 @@
+from recurvedata.executors.schemas import (
+    ColumnItem,
+    FullDatabaseItem,
+    ListDatabases,
+    Pagination,
+    ResponseModel,
+    TableItem,
+)
+
+
+class TestConnectionResponse(ResponseModel):
+    pass
+
+
+class ListDatabasesResponse(ResponseModel):
+    data: ListDatabases | None
+
+
+class ListTablesResponse(ResponseModel):
+    data: Pagination[TableItem] | None
+
+
+class ListColumnsResponse(ResponseModel):
+    data: Pagination[ColumnItem] | None
+
+
+class ListFullDatabasesResponse(ResponseModel):
+    data: Pagination[FullDatabaseItem] | None