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

recurvedata/operators/context.py

@@ -0,0 +1,255 @@
+import contextvars
+import functools
+import inspect
+import os
+import warnings
+from typing import TYPE_CHECKING, Any, Callable
+
+if TYPE_CHECKING:
+    from recurvedata.operators.task import BaseTask
+
+try:
+    from fsspec.asyn import get_running_loop, sync  # todo
+except ImportError:
+    pass
+
+from recurvedata.operators.operator import BaseOperator, _registry, get_operator_class
+
+
+class Context(object):
+    """
+    Web/Worker 调用的对象，提供：
+    1. 注册函数功能
+        调用之前需要注册一些函数，用于 config_schema 以及 validate，包括：
+            get_connection_names_by_type: 根据连接源 type 返回连接源 names。支持同步/异步写法
+            get_connection_by_name: 根据连接源 name 返回连接源对象。支持同步/异步写法 todo
+    2. list_config_schemas
+    3. get_supported_operators
+    4. get_config_schema
+    5. Worker 执行时，需要的一些函数
+
+    sync/async 调用
+    Context 支持同步/异步的调用方式，
+    1. Web 调用是异步，Worker 调用是同步；Web 端注册的 get_connection_names_by_type 等函数是异步的，Worker 端注册的函数是同步的
+    2. 为了避免在 operator 里引入 async/await 语法，保持 operator 代码的简洁，
+        各个 Operator 里，统一使用同步的写法，
+        Operator 里只提供同步的 config_schema, validate, execute 写法，
+        context 也提供同步的 get_connection_names_by_type 等方法，供 operator 调用
+    3. Web 注册的 get_connection_names_by_type 是异步的，Context 提供的 get_connection_names_by_type 是同步的。
+        Context 为了提供同步的 get_connection_names_by_type，把 Web 注册的异步的 get_connection_names_by_type 转成了同步
+    4. Operator 的 config_schema 是同步的，而 Web 端调用需要异步的方法，
+        所以 Context.get_config_schema 方法通过异步的方式调用 Operator.config_schema
+    """
+
+    def __init__(self):
+        # 根据连接源 type 返回连接源 names。支持同步/异步写法
+        self._get_connection_names_by_type: Callable = None
+
+        # 根据连接源 name 返回连接源对象。支持同步/异步写法
+        self._get_connection_by_name: Callable = None
+
+        self.current_project_id = contextvars.ContextVar("Recurve Project ID")
+
+        self._pid = os.getpid()
+        self._loop = None
+        self.async_mode = False
+        self._functions = {}
+
+    def init_context(self, get_connection_names_by_type: Callable = None, get_connection_by_name: Callable = None):
+        """
+        :param get_connection_names_by_type: 根据连接源 type 返回连接源 names。支持同步/异步写法
+                get_connection_names_by_type 函数定义: get_connection_names_by_type(project_id, connection_type)
+        :param get_connection_by_name: 根据连接源 name 返回连接源对象。支持同步/异步写法
+                get_connection_by_name 函数定义：get_connection_by_name(project_id, connection_name)
+        """
+        self._get_connection_names_by_type = get_connection_names_by_type
+        self._get_connection_by_name = get_connection_by_name
+        if inspect.iscoroutinefunction(self._get_connection_names_by_type):
+            self.async_mode = True
+        else:
+            self.async_mode = False
+
+    @property
+    def loop(self):
+        if self._pid != os.getpid():
+            raise RuntimeError("This class is not fork-safe")
+        if self._loop:
+            return self._loop
+        # self._loop = asyncio.get_event_loop() # todo: get_running_loop?
+        self._loop = get_running_loop()
+        # self._loop = get_loop() # todo: maybe have problem
+        return self._loop
+
+    def get_connection_names_by_type(self, connection_type: str) -> list[str]:
+        """
+        根据连接源类型，返回连接源名称
+        Web 端调用的时候，self._get_connection_names_by_type 是异步方法
+        Worker 端调用，self._get_connection_names_by_type 是同步方法
+        :param connection_type:
+        :return:
+        """
+        project_id = self.current_project_id.get()
+        if inspect.iscoroutinefunction(self._get_connection_names_by_type):
+            return sync(self.loop, self._get_connection_names_by_type, project_id, connection_type)
+        return self._get_connection_names_by_type(project_id, connection_type)
+
+    def get_connection_by_name(self, connection_name: str):
+        project_id = self.current_project_id.get()
+        if inspect.iscoroutinefunction(self._get_connection_by_name):
+            return sync(self.loop, self._get_connection_by_name, project_id, connection_name)
+        return self._get_connection_by_name(project_id, connection_name)
+
+    def get_connection_choices_by_type(self, connection_type):
+        warnings.warn(
+            "This function is deprecated. Please use `get_connection_names_by_type`",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        return self.get_connection_names_by_type(connection_type)
+
+    async def async_call_synchronous_func(self, func, *args):
+        project_id = self.current_project_id.get()
+        loop = self.loop
+        res = await loop.run_in_executor(None, self.contextvars_wrapper(project_id, func), *args)
+        return res
+
+    def must_get_connection_by_name(self, connection_name: str):
+        connection = self.get_connection_by_name(connection_name)
+        if not connection:
+            raise ValueError(f"connection {connection_name} not exists")
+        return connection
+
+    async def validate_operator_configuration(self, operator_name: str, configuration: dict, project_id: str):
+        self.current_project_id.set(project_id)
+        operator_cls: BaseOperator = get_operator_class(operator_name)
+        if not operator_cls:
+            raise ValueError(f"no operator {operator_name}")
+        return await self.async_call_synchronous_func(operator_cls.ui_validate, configuration)
+
+    def validate_operator_configuration_synchronously(self, operator_name: str, configuration: dict, project_id: str):
+        self.current_project_id.set(project_id)
+        operator_cls: BaseOperator = get_operator_class(operator_name)
+        if not operator_cls:
+            raise ValueError(f"no operator {operator_name}")
+        if not self.async_mode:
+            return operator_cls.ui_validate(configuration)
+        else:
+            return sync(self.loop, operator_cls.ui_validate, configuration)
+
+    # validate_operator_configuration_synchronously = sync_wrapper(validate_operator_configuration)
+
+    @staticmethod
+    def get_ds_name_field_values(operator_name: str, rendered_config: dict) -> list[str]:
+        operator_cls: BaseOperator = get_operator_class(operator_name)
+        if not operator_cls:
+            raise ValueError(f"no operator {operator_name}")
+        return operator_cls.get_ds_name_field_values(rendered_config)
+
+    def contextvars_wrapper(self, project_id, func):
+        """
+        init contextvars in asyncio
+        """
+
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            token = self.current_project_id.set(project_id)
+            res = func(*args, **kwargs)
+            self.current_project_id.reset(token)
+            return res
+
+        return wrapper
+
+    async def get_config_schema(self, operator_name: str, project_id: str):
+        """
+        默认返回的是 get_ui_config_schema
+        :param operator_name:
+        :param project_id:
+        :return:
+        """
+        self.current_project_id.set(project_id)
+        operator_cls: BaseOperator = get_operator_class(operator_name)
+        if operator_cls:
+            return await self.async_call_synchronous_func(operator_cls.ui_config_schema)
+
+    def get_config_schema_synchronously(self, operator_name: str, project_id: str):
+        self.current_project_id.set(project_id)
+        operator_cls: BaseOperator = get_operator_class(operator_name)
+        if not operator_cls:
+            raise ValueError(f"no operator {operator_name}")
+        if not self.async_mode:
+            return operator_cls.ui_config_schema()
+        else:
+            return sync(self.loop, operator_cls.ui_config_schema)
+
+    # get_config_schema_synchronously = sync_wrapper(get_config_schema)
+
+    @staticmethod
+    def get_supported_operators() -> list[str]:
+        res_lst = []
+        for op_name, op_cls in _registry.items():
+            if not op_cls.enabled:
+                continue
+            res_lst.append(op_name)
+        return res_lst
+
+    async def list_config_schemas(self, project_id: str):
+        self.current_project_id.set(project_id)
+        res_lst = []
+
+        for operator_name, operator_cls in _registry.items():
+            res_lst.append(
+                {
+                    "name": operator_name,
+                    "config_schema": await self.async_call_synchronous_func(operator_cls.config_schema),
+                }
+            )
+        return res_lst
+
+    def list_config_schemas_synchronously(self, project_id: str):
+        self.current_project_id.set(project_id)
+        res_lst = []
+
+        for operator_name, operator_cls in _registry.items():
+            res_lst.append(
+                {
+                    "name": operator_name,
+                    "config_schema": self.get_config_schema_synchronously(operator_name, project_id),
+                }
+            )
+        return res_lst
+
+    # list_config_schemas_synchronously = sync_wrapper(list_config_schemas)
+
+    @property
+    def client(self):
+        return self._client
+
+    @client.setter
+    def client(self, client):
+        self._client = client
+
+    def register_function(self, name: str, function: Callable):
+        self._functions[name] = function
+
+    def init_task_instance_on_task_start(self, task: "BaseTask", *args, **kwargs) -> int:
+        func = self._functions.get("init_task_instance_on_task_start")
+        if func:
+            return func(task, *args, **kwargs)
+
+    def update_task_instance_on_task_finish(
+        self,
+        task: "BaseTask",
+        ti_id: int,
+        task_status: str,
+        meta: Any,
+        error: Exception,
+        error_stack: str,
+        *args,
+        **kwargs,
+    ):
+        func = self._functions.get("update_task_instance_on_task_finish")
+        if func:
+            return func(task, ti_id, task_status, meta, error, error_stack, *args, **kwargs)
+
+
+context = Context()

recurvedata/operators/dbt_operator/__init__.py

@@ -0,0 +1,2 @@
+from recurvedata.operators.dbt_operator.model_pipeline_link_operator import LinkModelPipelineOperator
+from recurvedata.operators.dbt_operator.operator import DBTOperator

recurvedata/operators/dbt_operator/model_pipeline_link_operator.py

@@ -0,0 +1,55 @@
+from recurvedata.core.translation import _l
+from recurvedata.operators.link_operator import LinkOperator
+
+
+class LinkModelPipelineOperator(LinkOperator):
+    @classmethod
+    def config_schema(cls) -> dict:  # front-end does not use this config schema to show
+        return {
+            "type": "object",
+            "properties": {
+                "model_pipeline_id": {
+                    "type": "string",
+                    "title": _l("Model Pipeline ID"),
+                    "description": _l("Model Pipeline ID"),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "workflow_id": {
+                    "type": "string",
+                    "title": _l("Workflow ID"),
+                    "description": _l("Workflow ID"),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "workflow_version": {
+                    "type": "string",
+                    "title": _l("Workflow Version"),
+                    "description": _l("Workflow Version"),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "variables": {
+                    "type": "string",
+                    "title": _l("Variables"),
+                    "default": "{}",
+                    "description": _l("Variables in JSON format"),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "code",
+                        "lang": "json",
+                    },
+                },
+            },
+            "required": [
+                "model_pipeline_id",
+                "workflow_id",
+                "workflow_version",
+            ],
+        }

recurvedata/operators/dbt_operator/operator.py

@@ -0,0 +1,353 @@
+import datetime
+import logging
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Union
+
+from recurvedata.core.translation import _l
+from recurvedata.dbt.utils import parse_run_model_log
+from recurvedata.exceptions import MaxRetriesExceededException
+from recurvedata.operators.operator import BaseOperator
+from recurvedata.operators.task import BaseTask
+from recurvedata.utils.date_time import utcnow_aware
+from recurvedata.utils.helpers import get_environment_variable
+
+if TYPE_CHECKING:
+    from recurvedata.dbt.consts import DbtMaterialization
+    from recurvedata.dbt.schemas import PreviewResult
+    from recurvedata.dbt.service import DbtService
+
+    try:
+        from dbt.cli.main import dbtRunnerResult
+    except ImportError:
+        dbtRunnerResult = None
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class TaskRuntimeException:
+    exception: Exception
+
+    def to_dict(self):
+        return {
+            "success": False,
+            "exception": {
+                "type": f"TaskRuntimeException-{type(self.exception).__name__}",
+                "message": str(self.exception),
+            },
+        }
+
+
+@dataclass
+class DbtResultConstructor:
+    project_id: int
+    model_name: str
+    materialization: "DbtMaterialization"
+    compiled_code: str = None
+
+    @staticmethod
+    def _construct_timing(action_name: str, start_time: datetime.datetime, end_time: datetime.datetime) -> list[dict]:
+        def _format_time(dt: datetime.datetime) -> str:
+            dt_utc = dt.astimezone(datetime.timezone.utc)
+            ds = dt_utc.strftime("%Y-%m-%dT%H:%M:%S.%fZ")
+            return ds
+
+        return [
+            {"name": action_name, "started_at": _format_time(start_time), "completed_at": _format_time(end_time)},
+        ]
+
+    def construct_ephemeral_materialized_result(
+        self, materialized_result_dct: dict, start_time: datetime.datetime, end_time: datetime.datetime
+    ) -> dict:
+        """
+        For ephemeral materialization, the materialized_result_dct $result.results is empty.
+        CP relies on $result.results to show message and start/end time
+        """
+        if not materialized_result_dct["success"]:
+            return materialized_result_dct
+        result_dct = materialized_result_dct["result"]
+        if not result_dct:
+            return materialized_result_dct
+        sub_results: list[dict] = result_dct["results"]
+        if not sub_results:
+            result_dct["results"] = [
+                {
+                    "unique_id": self.format_model_unique_id(),
+                    "status": "success",
+                    "timing": self._construct_timing("execute", start_time=start_time, end_time=end_time),
+                    "message": "Ephemeral model compiled successfully",
+                    "compiled_code": self.compiled_code,
+                }
+            ]
+        return materialized_result_dct
+
+    def format_model_unique_id(self) -> str:
+        return f"model.project_{self.project_id}.{self.model_name}"
+
+
+class DbtTask(BaseTask):
+    def execute_impl(self):
+        from recurvedata.dbt.schemas import PreviewResult
+        from recurvedata.dbt.service import DbtService
+        from recurvedata.dbt.utils import format_var
+        from recurvedata.utils.redis_lock import RedisLock
+
+        model_name = self.rendered_config.get("model_name") or self.rendered_config.get("entity_name")
+
+        lock = RedisLock(
+            f"dbt_task_{self.dag.project_id}_{model_name}", auto_extend=True, expire=60, timeout=60 * 60 * 1
+        )
+        lock.acquire()
+
+        try:
+            service = DbtService(self.dag.project_id)
+            service.prepare()
+            model_id: int = int(self.rendered_config["entity_id"])
+            var_str = format_var(service, self.get_template_context())
+
+            materialize_start_time = utcnow_aware()
+            full_refresh = self.dag.full_refresh_models
+
+            model_run_result = service.run_model(model_name, var_str, full_refresh=full_refresh, include_run_log=True)
+            compiled_code = model_run_result.compiled_sql
+            materialized_result = model_run_result.result
+            run_sql = model_run_result.run_sql
+            run_log = model_run_result.run_log
+
+            materialize_end_time = utcnow_aware()
+
+            if not compiled_code:
+                logger.info("compiled_code empty, use un-compiled sql")
+                compiled_code = service.read_model_sql(model_name)
+
+            if not compiled_code:
+                logger.info("compiled_code still empty, set materialized_result to failed")
+                materialized_result.success = False
+                materialized_result.exception = RuntimeError("Materialization failed due to empty compiled_code")
+                materialized_result.result = None
+
+            if not materialized_result.success:
+                self.send_dbt_model_result(
+                    service,
+                    compiled_sql=compiled_code,
+                    run_sql=run_sql,
+                    run_log=run_log,
+                    try_number=get_environment_variable("AIRFLOW_RETRY_NUMBER", int),
+                    materialized_result=materialized_result,
+                    materialize_start_time=materialize_start_time,
+                    materialize_end_time=materialize_end_time,
+                )
+                raise Exception(f"run model {model_name} materialized failed")
+
+            test_case_skipped = False
+            if self.dag.skip_data_tests:
+                logger.info("skip data tests")
+                test_result = None
+                test_case_sample_result = None
+                test_case_skipped = True
+            else:
+                logger.info("run data tests")
+                test_result = service.run_test(model_id, var_str)
+                test_case_sample_result: dict[str, PreviewResult] = service.run_test_sample_data(test_result)
+            self.send_dbt_model_result(
+                service,
+                compiled_sql=compiled_code,
+                run_sql=run_sql,
+                run_log=run_log,
+                try_number=get_environment_variable("AIRFLOW_RETRY_NUMBER", int),
+                materialized_result=materialized_result,
+                test_case_result=test_result,
+                test_case_sample_result=test_case_sample_result,
+                materialize_start_time=materialize_start_time,
+                materialize_end_time=materialize_end_time,
+                test_case_skipped=test_case_skipped,
+            )
+
+            test_case_result_dct = self.format_test_case_result(test_result)
+            if test_case_result_dct and not test_case_result_dct["success"]:
+                raise Exception("Task Run failed due to Error / Failed test cases")
+        except Exception as e:
+            raise e from None
+        finally:
+            lock.release()
+
+    @staticmethod
+    def format_materialized_result(
+        project_id: int,
+        model_name: str,
+        materialization: Union["DbtMaterialization", str],
+        compiled_code: str,
+        materialized_result: Union["dbtRunnerResult", "TaskRuntimeException"],
+        materialize_start_time: datetime.datetime = None,
+        materialize_end_time: datetime.datetime = None,
+    ) -> dict | None:
+        """
+        materialized: model, ephemeral, view, incremental
+        """
+        from recurvedata.dbt.consts import DbtMaterialization
+        from recurvedata.dbt.utils import dbt_runner_result_to_dict
+
+        if not materialized_result:
+            return
+        if isinstance(materialized_result, TaskRuntimeException):
+            materialized_result_dct = materialized_result.to_dict()
+        else:
+            materialized_result_dct = dbt_runner_result_to_dict(materialized_result)
+
+        if materialization == DbtMaterialization.EPHEMERAL:
+            constructor = DbtResultConstructor(
+                project_id=project_id,
+                model_name=model_name,
+                materialization=materialization,
+                compiled_code=compiled_code,
+            )
+            materialized_result_dct = constructor.construct_ephemeral_materialized_result(
+                materialized_result_dct, materialize_start_time, materialize_end_time
+            )
+
+        if materialized_result_dct["success"]:
+            results = materialized_result_dct.get("result", {}).get("results")
+            if not results:
+                # The selection criterion '' does not match any nodes
+                materialized_result_dct["success"] = False
+
+        return materialized_result_dct
+
+    @staticmethod
+    def format_test_case_result(test_case_result: Union["dbtRunnerResult", "TaskRuntimeException"]) -> dict | None:
+        from recurvedata.dbt.utils import dbt_runner_result_to_dict
+
+        if not test_case_result:
+            return
+        if isinstance(test_case_result, TaskRuntimeException):
+            test_case_result_dct = test_case_result.to_dict()
+        else:
+            test_case_result_dct = dbt_runner_result_to_dict(test_case_result)
+        return test_case_result_dct
+
+    @property
+    def model_name(self) -> str:
+        return self.rendered_config.get("model_name") or self.rendered_config.get("entity_name")
+
+    @property
+    def materialization(self) -> str:
+        return self.rendered_config.get("materialized")
+
+    def send_dbt_model_result(
+        self,
+        service: "DbtService",
+        compiled_sql: str | None,
+        try_number: int,
+        run_sql: str | None = None,
+        run_log: str | None = None,
+        materialized_result: Union["dbtRunnerResult", "TaskRuntimeException"] = None,
+        test_case_result: Union["dbtRunnerResult", "TaskRuntimeException"] = None,
+        test_case_sample_result: dict[str, "PreviewResult"] = None,
+        materialize_start_time: datetime.datetime = None,
+        materialize_end_time: datetime.datetime = None,
+        test_case_skipped: bool = False,
+    ):
+        materialized_result_dct = self.format_materialized_result(
+            self.dag.project_id,
+            self.model_name,
+            self.materialization,
+            compiled_sql,
+            materialized_result,
+            materialize_start_time,
+            materialize_end_time,
+        )
+        test_case_result_dct = self.format_test_case_result(test_case_result)
+
+        if not compiled_sql:
+            logger.info(f"compiled_sql empty, materialized_result_dct: {materialized_result_dct}")
+
+        if test_case_sample_result:
+            test_case_sample_result_dct = {
+                unique_id: preview_obj.model_dump() for unique_id, preview_obj in test_case_sample_result.items()
+            }
+        else:
+            test_case_sample_result_dct = None
+
+        run_sql_log = parse_run_model_log(run_log)
+
+        logger.info(f"debug: compiled sql: {compiled_sql}")
+        logger.info(f"debug: run sql: {run_sql}")
+        logger.info(f"debug: run_log: {run_log}")
+        logger.info(f"debug: run_sql_log: {run_sql_log}")
+        logger.info(f"debug: materialized_result_dct: {materialized_result_dct}")
+        logger.info(f"debug: test_case_result_dct: {test_case_result_dct}")
+        logger.info(f"debug: test_case_sample_result_dct: {test_case_sample_result_dct}")
+
+        try:
+            service.client.send_dbt_model_result(
+                self.dag.id,
+                self.node.node_key,
+                compiled_sql,
+                run_sql,
+                run_sql_log=run_sql_log,
+                raw_materialized_result=materialized_result_dct,
+                raw_test_result=test_case_result_dct,
+                test_case_sample_data=test_case_sample_result_dct,
+                materialization=self.materialization,
+                try_number=try_number,
+                test_case_skipped=test_case_skipped,
+            )
+        except MaxRetriesExceededException as e:
+            logger.exception(f"send_dbt_model_result failed, error: {e}")
+        self.sent_dbt_model_result = True
+
+    def on_execute_impl_error(self, err: Exception):
+        from recurvedata.dbt.service import DbtService
+
+        if getattr(self, "sent_dbt_model_result", False):
+            return
+        service = DbtService(self.dag.project_id)
+        self.send_dbt_model_result(
+            service,
+            compiled_sql=None,
+            try_number=get_environment_variable("AIRFLOW_RETRY_NUMBER", int),
+            materialized_result=TaskRuntimeException(err),
+            test_case_result=None,
+        )
+
+
+class DBTOperator(BaseOperator):
+    task_cls = DbtTask
+
+    @classmethod
+    def config_schema(cls) -> dict:
+        return {
+            "type": "object",
+            "properties": {
+                "entity_name": {
+                    "type": "string",
+                    "title": _l("Entity Name"),
+                    "description": _l("Entity Name"),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "materialized": {  # for front-end display
+                    "type": "string",
+                    "title": _l("Materialized"),
+                    "default": "view",
+                    "enum": ["table", "view", "incremental", "ephemeral"],
+                    "enumNames": ["table", "view", "incremental", "ephemeral"],
+                },
+            },
+            "required": ["entity_name", "materialized"],
+        }
+
+    @classmethod
+    def validate(cls, configuration) -> dict:
+        return configuration
+
+    @classmethod
+    def ui_config_to_config(cls, configuration: dict) -> dict:
+        source = configuration["source"]
+        return source
+
+    @classmethod
+    def get_ds_name_field_values(cls, rendered_config: dict) -> list[str]:
+        return []

recurvedata/operators/link_operator/__init__.py

@@ -0,0 +1 @@
+from recurvedata.operators.link_operator.operator import LinkOperator