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

recurvedata/operators/transfer_operator/operator.py

@@ -0,0 +1,231 @@
+import datetime
+import logging
+import typing
+
+try:
+    from recurvedata.pigeon.utils.fs import new_stagefile_factory
+except ImportError:
+    pass
+
+from recurvedata.operators.config import CONF
+from recurvedata.operators.models import DagBase, NodeBase
+from recurvedata.operators.operator import BaseOperator
+from recurvedata.operators.transfer_operator.task import get_dump_classes, get_load_classes, get_task_class
+from recurvedata.operators.ui import format_config_schema
+from recurvedata.utils import md5hash
+
+if typing.TYPE_CHECKING:
+    from recurvedata.operators.transfer_operator.task import DumpTask, LoadTask
+
+logger = logging.getLogger(__name__)
+
+
+class TransferOperator(BaseOperator):
+    """
+    Operator that handles data transfer operations between dump and load stages.
+    Manages the execution of dump and load tasks with appropriate configurations.
+    """
+
+    stages = ("dump", "load")
+
+    def __init__(self, dag: DagBase, node: NodeBase, execution_date: datetime.datetime, variables: dict = None) -> None:
+        self.dump_task: "DumpTask" = None
+        self.load_task: "LoadTask" = None
+        self.filename: str = self._determine_filename(dag, node, execution_date)
+        # self.execution_date = as_local_datetime(execution_date)
+
+        super().__init__(dag, node, execution_date, variables)
+
+    def init_task(self):
+        params = {
+            "dag": self.dag,
+            "node": self.node,
+            "execution_date": self.execution_date,
+            "filename": self.filename,
+            "variables": self.variables,
+        }
+
+        load_config = self.node.configuration["load"]
+        load_cls = self.get_task_class(load_config["name"])
+        logger.debug(f"create load task with {params}")
+
+        self.load_task: LoadTask = load_cls(config=load_config["config"], **params)
+
+        # TODO: 最好能去掉这种配置，dump 和 load 都使用使用统一的 CSV 格式，由 loader 自己去处理
+        handler_options = {
+            "encoding": None,
+            "write_header": self._determine_write_header(),
+        }
+        if self.load_task.default_dumper_handler_options:
+            handler_options.update(self.load_task.default_dumper_handler_options)
+
+        dump_config = self.node.configuration["dump"]
+        dump_cls = self.get_task_class(dump_config["name"])
+        logger.debug(f"create dump task with {params}")
+        self.dump_task: DumpTask = dump_cls(config=dump_config["config"], handler_options=handler_options, **params)
+        self.load_task.dump_task_type = dump_cls.__name__
+
+    def set_execution_date(self, execution_date):
+        self.dump_task.set_execution_date(execution_date)
+        self.load_task.set_execution_date(execution_date)
+
+    def _determine_write_header(self):
+        return self.load_task.should_write_header
+
+    @staticmethod
+    def _determine_filename(dag: DagBase, node: NodeBase, execution_date: datetime.datetime) -> str:
+        """
+        Generate a unique filename for the transfer operation.
+
+        Args:
+            dag: The DAG instance
+            node: The node instance
+            execution_date: The execution datetime
+
+        Returns:
+            str: Generated filename
+        """
+        dag_id = dag.id
+        node_id = node.node_key
+        is_link_node = getattr(node, "is_link_op", False)
+        if not is_link_node:
+            hash_txt = md5hash(f"{dag_id}|{node_id}|{execution_date}")
+            prefix = f"{dag_id}_{node_id}_"
+        else:
+            origin_node = node.origin_node
+            hash_txt = md5hash(f"{dag_id}|{origin_node.node_key}|{node_id}|{execution_date}")
+            prefix = f"{dag_id}_{origin_node.node_key}_{node_id}_"
+            logger.info(f"link op _determine_filename: {prefix} {hash_txt}")
+
+        hash_len = max(8, len(hash_txt) - len(prefix))
+        return new_stagefile_factory(CONF.DATA_ROOT)(prefix + hash_txt[:hash_len])
+
+    def dump(self):
+        return self.dump_task.execute()
+
+    def load(self):
+        return self.load_task.execute()
+
+    def execute(self):
+        self.dump()
+        self.load()
+
+    @classmethod
+    def validate(cls, configuration: dict):
+        config = {
+            "dump": cls._validate_task_config(configuration["dump"]),
+            "load": cls._validate_task_config(configuration["load"]),
+        }
+        return config
+
+    @classmethod
+    def _validate_task_config(cls, config: dict):
+        task_cls = cls.get_task_class(config["name"])
+        cfg = task_cls.validate(config["config"])
+        return {"name": config["name"], "config": cfg}
+
+    @classmethod
+    def to_dict(cls) -> dict:
+        return {
+            "name": cls.name(),
+            "config_schema": {
+                "dump": [x.to_dict() for x in get_dump_classes()],
+                "load": [x.to_dict() for x in get_load_classes()],
+            },
+        }
+
+    @classmethod
+    def config_schema(cls):
+        return {
+            "dump": [{"name": x.name(), "config_schema": x.config_schema()} for x in cls.get_dump_classes()],
+            "load": [{"name": x.name(), "config_schema": x.config_schema()} for x in cls.get_load_classes()],
+        }
+
+    @classmethod
+    def ui_config_schema(cls):
+        return {
+            "dump": {
+                "name": "Dump",
+                "config_schema": [
+                    {"name": x.name(), "config_schema": format_config_schema(x.config_schema(), "dump")}
+                    for x in cls.get_dump_classes()
+                ],
+            },
+            "load": {
+                "name": "Load",
+                "config_schema": [
+                    {"name": x.name(), "config_schema": format_config_schema(x.config_schema(), "load")}
+                    for x in cls.get_load_classes()
+                ],
+            },
+        }
+
+    @classmethod
+    def ui_validate(cls, configuration: dict) -> dict:
+        res = {
+            "dump": cls._add_schema_name_to_json_schema_error("dump", cls._validate_task_config, configuration["dump"]),
+            "load": cls._add_schema_name_to_json_schema_error("load", cls._validate_task_config, configuration["load"]),
+        }
+        return res
+
+    @classmethod
+    def ui_config_to_config(cls, configuration: dict) -> dict:
+        return {
+            "dump": configuration["dump"],
+            "load": configuration["load"],
+        }
+
+    @classmethod
+    def get_ds_name_field_values(cls, rendered_config: dict) -> list[str]:
+        config = cls.ui_config_to_config(rendered_config)
+        res = []
+        dump_cls = cls.get_task_class(config["dump"]["name"])
+        if dump_cls:
+            res.extend(dump_cls.get_ds_name_field_values(config["dump"]["config"]))
+        load_cls = cls.get_task_class(config["load"]["name"])
+        if load_cls:
+            res.extend(load_cls.get_ds_name_field_values(config["load"]["config"]))
+        return res
+
+    @classmethod
+    def get_task_class(cls, name: str):
+        return get_task_class(name)
+
+    @classmethod
+    def get_dump_classes(cls, check_enabled=True):
+        res_lst = get_dump_classes()
+        if check_enabled:
+            res_lst = [dump_cls for dump_cls in res_lst if dump_cls.enabled]
+        return res_lst
+
+    @classmethod
+    def get_load_classes(cls, check_enabled=True):
+        res_lst = get_load_classes()
+        if check_enabled:
+            res_lst = [load_cls for load_cls in res_lst if load_cls.enabled]
+        return res_lst
+
+    @classmethod
+    def get_setup_install_require(cls) -> dict:
+        require_dct = {}
+        op_name = cls.name()
+        op_web_requires = cls.web_install_require[:]
+        op_worker_requires = cls.worker_install_require[:]
+        for dump_cls in cls.get_dump_classes():
+            if dump_cls.web_install_require:
+                require_dct[f"web.{op_name}.dump.{dump_cls.name()}"] = dump_cls.web_install_require
+                op_web_requires.extend(dump_cls.web_install_require)
+            if dump_cls.worker_install_require:
+                require_dct[f"worker.{op_name}.dump.{dump_cls.name()}"] = dump_cls.worker_install_require
+                op_worker_requires.extend(dump_cls.worker_install_require)
+
+        for load_cls in cls.get_load_classes():
+            if load_cls.web_install_require:
+                require_dct[f"web.{op_name}.load.{load_cls.name()}"] = load_cls.web_install_require
+                op_web_requires.extend(load_cls.web_install_require)
+            if load_cls.worker_install_require:
+                require_dct[f"worker.{op_name}.load.{load_cls.name()}"] = load_cls.worker_install_require
+                op_worker_requires.extend(load_cls.worker_install_require)
+        require_dct["web"] = sorted(list(set(op_web_requires)))
+        require_dct["worker"] = sorted(list(set(op_worker_requires)))
+        return require_dct

recurvedata/operators/transfer_operator/task.py

@@ -0,0 +1,223 @@
+import json
+import logging
+import os
+import shutil
+import tempfile
+import traceback
+
+import jsonschema
+
+from recurvedata.core.transformer import Transformer
+from recurvedata.operators.task import BaseTask
+from recurvedata.utils.attrdict import AttrDict
+from recurvedata.utils.date_time import round_time_resolution
+from recurvedata.utils.helpers import first
+from recurvedata.utils.registry import Registry
+
+try:
+    from recurvedata.pigeon.handler.csv_handler import create_csv_file_handler_factory
+    from recurvedata.pigeon.utils import fs, trim_suffix
+except ImportError:
+    pass
+
+from recurvedata.operators.transfer_operator import utils
+
+logger = logging.getLogger(__name__)
+_registry = Registry(key_callback=lambda x: x.name())
+_load_task_registry = Registry(key_callback=lambda x: x.ds_types)
+
+
+class Task(BaseTask):
+    worker_install_require = []
+    web_install_require = []
+
+    def __init__(self, dag, node, execution_date, variables, config, filename):
+        super().__init__(dag, node, execution_date, variables)
+
+        self.config = AttrDict(config)
+        self.filename = filename
+
+    @classmethod
+    def type(cls):
+        return None
+
+    @staticmethod
+    def first_or_default(dss, default=""):
+        return first(dss, default)
+
+
+class DumpTask(Task):
+    _AUTO_REGISTER = True
+    _MAX_ERROR_RATE = 0
+    no_template_fields = (
+        "data_source_name",
+        "filter_engine",
+    )
+
+    @classmethod
+    def type(cls):
+        return "dump"
+
+    @property
+    def stage(self) -> str:
+        return "dump"
+
+    def __init_subclass__(cls, **kwargs):
+        if cls._AUTO_REGISTER:
+            _registry.add(cls)
+
+    def __init__(self, handler_options=None, *args, **kwargs):
+        self.handler_options = handler_options or {}
+        super().__init__(*args, **kwargs)
+
+    @classmethod
+    def validate(cls, configuration):
+        config = super().validate(configuration)
+
+        transformer_code = configuration.get("transform", "").strip()
+        if not transformer_code:
+            return config
+        try:
+            utils.validate_transform(transformer_code)
+        except (ValueError, TypeError) as e:
+            raise jsonschema.ValidationError(message=str(e), path=("transform",))
+        except Exception:
+            tb = traceback.format_exc(limit=0)
+            msg = "\n".join(tb.splitlines()[1:])
+            raise jsonschema.ValidationError(message=msg, path=("transform",))
+
+        if "custom_handler_options" in config:
+            try:
+                value = json.loads(config["custom_handler_options"])
+            except Exception:
+                raise jsonschema.ValidationError(
+                    message="custom_handler_options should be valid JSON", path=("custom_handler_options",)
+                )
+            if not isinstance(value, dict):
+                raise jsonschema.ValidationError(
+                    message="custom_handler_options should be dict", path=("custom_handler_options",)
+                )
+        return config
+
+    def create_handler_factory(self):
+        self.remove_intermediate_files()
+        transformer = self.create_transformer()
+        kwargs = self.handler_options.copy()
+        encoding = self.rendered_config.get("middle_file_encoding")
+        kwargs.update(
+            {
+                "filename": self.filename,
+                "encoding": encoding,
+                "transformer": transformer,
+                "max_error_rate": self._MAX_ERROR_RATE,
+            }
+        )
+        # FIXME: ugly way to get more handler options from Transformer definition
+        kwargs.update(getattr(transformer, "handler_options", {}))
+
+        # allow user to override the default handler options
+        if self.rendered_config.custom_handler_options:
+            custom_handler_options = json.loads(self.rendered_config.custom_handler_options)
+            kwargs.update(custom_handler_options)
+        hf = create_csv_file_handler_factory(**kwargs)
+        return hf
+
+    def create_transformer(self) -> Transformer:
+        transformer_code = self.rendered_config.get("transform", "").strip()
+        if transformer_code:
+            transformer = utils.validate_transform(transformer_code)
+        else:
+            transformer = None
+        return transformer
+
+    def has_custom_transformer(self):
+        return self.rendered_config.get("transform")
+
+    def get_schedule_time_range(self):
+        end_date = self.execution_date
+        start_date = self.dag.previous_schedule(self.execution_date)
+        if self.config.get("time_auto_round", False):
+            start_date = round_time_resolution(start_date, self.dag.schedule_interval)
+            end_date = round_time_resolution(end_date, self.dag.schedule_interval)
+        return start_date, end_date
+
+    def remove_intermediate_files(self):
+        pattern = f"{self.filename}.*"
+        logger.info(f"remove intermediate files {pattern}")
+        fs.remove_files_by_pattern(pattern)
+
+    def on_execute_impl_error(self, exc: Exception):
+        logger.exception(f"caught error: {exc}")
+        self.remove_intermediate_files()
+
+
+class LoadTask(Task):
+    ds_types = ()
+    should_write_header = False
+    default_dumper_handler_options = {}
+    dump_task_type = None
+
+    def __init_subclass__(cls, **kwargs):
+        _registry.add(cls)
+        _load_task_registry.add(cls)
+
+    @classmethod
+    def type(cls):
+        return "load"
+
+    @property
+    def stage(self) -> str:
+        return "load"
+
+    @staticmethod
+    def compress_file(filename, target_filename=None, compress_mode="None"):
+        """compress file before loading, only support gzip and zip"""
+        if compress_mode == "None":
+            return filename, None
+        if compress_mode not in ("Gzip", "Zip", "Bzip2"):
+            raise ValueError(f"{compress_mode} is not supported")
+
+        logger.info(f"Compressing file using {compress_mode}")
+        compress_method, ext = {
+            "Gzip": (fs.gzip_compress, ".gz"),
+            "Zip": (fs.zip_compress, ".zip"),
+            "Bzip2": (fs.bzip2_compress, ".bz2"),
+        }[compress_mode]
+
+        # 如果指定了压缩包内的文件名，先把文件临时改名为目标文件名，压缩完了再改回来
+        if target_filename:
+            inner_filename = trim_suffix(os.path.basename(target_filename), ext)
+            tmp_dir = tempfile.mkdtemp(dir=os.path.dirname(filename))
+            file_to_compress = os.path.join(tmp_dir, inner_filename)
+            os.rename(filename, file_to_compress)
+        else:
+            target_filename = f"{filename}{ext}"
+            file_to_compress = filename
+
+        try:
+            compressed_file = compress_method(file_to_compress, target_filename=target_filename, using_cmd=True)
+        except BaseException as e:
+            raise e
+        finally:
+            # 如果发生异常，做回滚操作
+            if file_to_compress != filename:
+                os.rename(file_to_compress, filename)
+                shutil.rmtree(os.path.dirname(file_to_compress))
+        return compressed_file, ext
+
+
+def get_task_class(name):
+    return _registry[name]
+
+
+def get_dump_classes():
+    return sorted([x for x in _registry.values() if x.type() == "dump"], key=lambda x: x.name())
+
+
+def get_load_classes():
+    return sorted([x for x in _registry.values() if x.type() == "load"], key=lambda x: x.name())
+
+
+def get_load_by_ds_type(ds_type):
+    klass = _load_task_registry.get(ds_type)
+    return klass.name()

recurvedata/operators/transfer_operator/utils.py

@@ -0,0 +1,134 @@
+import importlib.util
+import sys
+
+from recurvedata.core.transformer import Transformer
+from recurvedata.core.translation import _l
+from recurvedata.operators.transfer_operator import const
+
+allowed_modes = (const.LOAD_OVERWRITE, const.LOAD_MERGE, const.LOAD_APPEND)
+
+_TRANSFORM_SKELETON = """\
+from recurvedata.core.transformer import Transformer
+
+
+class MyTransformer(Transformer):
+    def transform_impl(self, row, *args, **kwargs):
+        # The row is an OrderedDict. Write your custom transformation logic here.
+        return row
+
+
+# Instantiate the transformer, the name must be `transformer`
+transformer = MyTransformer()
+"""
+
+
+TRANSFORM = {
+    "type": "string",
+    "title": _l("Custom Transformation"),
+    "description": _l(
+        "Python code to transform data during transfer. Must implement a Transformer class with "
+        "transform_impl method that processes each row. See example code below."
+    ),
+    "default": _TRANSFORM_SKELETON,
+    "ui:field": "CodeEditorWithReferencesField",
+    "ui:options": {
+        "type": "code",
+        "lang": "python",
+    },
+}
+
+LOAD_COMMON = {
+    "mode": {
+        "type": "string",
+        "title": _l("Load Mode"),
+        "description": _l("How to handle existing data in the target table"),
+        "enum": list(allowed_modes),
+        "enumNames": list(allowed_modes),
+        "default": const.LOAD_OVERWRITE,
+    },
+    "primary_keys": {
+        "type": "string",
+        "title": _l("Primary Keys"),
+        "description": _l(
+            "Comma-separated list of columns used for deduplication in MERGE mode. "
+            "Should be primary or unique key columns."
+        ),
+        "ui:field": "CodeEditorWithReferencesField",
+        "ui:options": {
+            "type": "plain",
+        },
+        "ui:hidden": '{{parentFormData.mode !== "MERGE"}}',
+    },
+    "dedup": {
+        "type": "boolean",
+        "title": _l("Enable Deduplication"),
+        "default": False,
+        "description": _l("Remove duplicate rows from the data before loading"),
+        "ui:widget": "BaseCheckbox",
+        "ui:options": {
+            "label": _l("Enable Deduplication"),
+        },
+    },
+    "dedup_uniq_keys": {
+        "type": "string",
+        "title": _l("Deduplication Keys"),
+        "description": _l("Comma-separated list of columns that uniquely identify each row"),
+        "ui:field": "CodeEditorWithReferencesField",
+        "ui:options": {
+            "type": "plain",
+        },
+        "ui:hidden": "{{!parentFormData.dedup}}",
+    },
+    "dedup_orderby": {
+        "type": "string",
+        "title": _l("Sort Order"),
+        "description": _l("Comma-separated list of columns to sort by before deduplication"),
+        "ui:field": "CodeEditorWithReferencesField",
+        "ui:options": {
+            "type": "plain",
+        },
+        "ui:hidden": "{{!parentFormData.dedup}}",
+    },
+    # "pre_queries": {
+    #     "type": "string",
+    #     "title": "Queries Ran Before Loading",
+    #     "description": '新数据导入前运行的 SQL，多条 SQL 用 `;` 分隔；支持传入变量，详见 <a target="_blank" href="https://bit.ly/2JMutjn">文档</a>',
+    #     "ui:field": "CodeEditorWithReferencesField",
+    #     "ui:options": {
+    #         "type": "code",
+    #         "lang": "sql",
+    #         "sqlLang": "sql",
+    #     },
+    # },
+    # "post_queries": {
+    #     "type": "string",
+    #     "title": "Queries Ran After Loading",
+    #     "description": '新数据导入后运行的 SQL，多条 SQL 用 `;` 分隔；支持传入变量，详见 <a target="_blank" href="https://bit.ly/2JMutjn">文档</a>',
+    #     "ui:field": "CodeEditorWithReferencesField",
+    #     "ui:options": {
+    #         "type": "code",
+    #         "lang": "sql",
+    #         "sqlLang": "sql",
+    #     },
+    # },
+}
+
+__spec = importlib.util.spec_from_loader("recurve_hack", None)
+__recurve_hack = importlib.util.module_from_spec(__spec)
+sys.modules["recurve_hack"] = __recurve_hack
+
+
+def validate_transform(raw_code):
+    from recurvedata.pigeon.transformer import Transformer as PigeonTransformer
+
+    code = compile(raw_code, "", "exec")
+    exec(code, __recurve_hack.__dict__)
+    transformer = __recurve_hack.__dict__.get("transformer")
+    if not transformer:
+        raise ValueError("transformer is required")
+    if (
+        not isinstance(transformer, (Transformer, PigeonTransformer))
+        and transformer.__class__.__name__ != "MyTransformer"
+    ):
+        raise TypeError(f"transformer should be type of pigeon.transformer.Transformer, {type(transformer)}")
+    return transformer

recurvedata/operators/ui.py

@@ -0,0 +1,80 @@
+def format_config_schema(config_schema: dict, schema_name: str):
+    """
+    按前端需求，重新格式化 config_schema
+    """
+    if "properties" not in config_schema:
+        return config_schema
+    for field_name, field_dct in config_schema["properties"].items():
+        if field_dct["type"] == "object":
+            format_config_schema(field_dct, field_name)
+        else:
+            format_field_schema(field_dct, schema_name)
+    return config_schema
+
+
+def format_field_schema(field_dct: dict, schema_name: str):
+    """
+    field_dct 例子：
+        {
+            'type': 'string',
+            'title': 'Data Source',
+            "ui:field": "ProjectConnectionSelectorField",
+            "ui:options": {
+                "supportTypes": ["mysql", "postgres",],
+            },
+        },
+    """
+    _add_option_id(field_dct)
+    _format_input_with_variable(field_dct, schema_name)
+    _format_aliases_select_field(field_dct)
+
+
+def _add_option_id(field_dct: dict):
+    ui_field = field_dct.get("ui:field")
+    if ui_field == "CodeEditorWithReferencesField":
+        return
+    if "ui:options" not in field_dct:
+        field_dct["ui:options"] = {}
+    if "id" in field_dct["ui:options"]:
+        return
+    field_dct["ui:options"]["id"] = ""
+
+
+def _format_input_with_variable(field_dct: dict, schema_name: str):
+    ui_field = field_dct.get("ui:field")
+    if ui_field != "CodeEditorWithReferencesField":
+        return
+    ui_options: dict = field_dct.get("ui:options")
+    if not ui_options:
+        return
+    ui_type = ui_options.get("type")
+    if ui_type != "code":
+        return
+    # 全屏相关配置
+    if "toParent" in ui_options:
+        return
+    ui_options["toParent"] = ".expanded_code_position"
+    ui_options["parentName"] = schema_name
+    ui_options["needExpandBtn"] = True
+
+
+def _format_aliases_select_field(field_dct: dict):
+    """
+    ProjectConnectionSelectorField 的 supportTypes 前端需要展示 connection 的 ui_type
+    """
+    from recurvedata.connectors import get_connection_ui_type
+
+    ui_field = field_dct.get("ui:field")
+    if ui_field != "ProjectConnectionSelectorField":
+        return
+    ui_options: dict = field_dct.get("ui:options")
+    if not ui_options:
+        return
+
+    support_types = ui_options.get("supportTypes")
+    if not support_types:
+        return
+
+    ui_options["supportTypes"] = [
+        ui_type for ui_type in [get_connection_ui_type(backend_type) for backend_type in support_types] if ui_type
+    ]