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

recurvedata/operators/task.py

@@ -0,0 +1,211 @@
+import datetime
+import json
+import logging
+import urllib.parse
+from functools import cached_property
+from typing import TYPE_CHECKING, Any, Optional, Union
+
+import jsonschema
+
+from recurvedata.consts import ETLExecutionStatus
+from recurvedata.core.templating import Renderer
+from recurvedata.operators.base import Configurable
+from recurvedata.operators.context import context
+from recurvedata.operators.models import DagBase, NodeBase
+from recurvedata.utils.attrdict import AttrDict
+
+if TYPE_CHECKING:
+    from recurvedata.executors.client import ExecutorClient
+
+
+logger = logging.getLogger(__name__)
+
+
+class LineageTaskMixin(object):
+    # todo: move to utils
+    def process_lineage(self):
+        try:
+            lineage = self.parse_lineage()
+            self.save_lineage(lineage)
+        except Exception as e:
+            # lineage_fail_notify(self)
+            logger.exception(f"failed to process lineage, error: {e}")
+
+    def parse_lineage(self):
+        pass
+
+    def save_lineage(self, lineage):
+        if not lineage:
+            return
+
+        self.save_lineage(self, lineage)
+        # todo: worker sdk
+
+
+class BaseTask(Configurable, LineageTaskMixin):
+    no_template_fields = ()  # 不使用 jinja 渲染的字段
+    ds_name_fields = ()
+
+    def __init__(self, dag: DagBase, node: NodeBase, execution_date: datetime.datetime, variables: dict = None):
+        self.dag: DagBase = dag
+        self.node: NodeBase = node
+        self.execution_date: datetime.datetime = execution_date
+        self.variables: dict = variables or {}
+
+        self.config = AttrDict(self.node.configuration)
+        self.task_instance_id: int = 0
+
+    @classmethod
+    def validate(cls, configuration: dict) -> dict:
+        config = super().validate(configuration)
+
+        # validate data sources
+        for name in cls.ds_name_fields:
+            ds = context.get_connection_by_name(connection_name=configuration[name])
+            if not ds:
+                raise jsonschema.ValidationError(
+                    message=f"Unknown data source {repr(configuration[name])}", path=(name,)
+                )
+        return config
+
+    @classmethod
+    def get_ds_name_field_values(cls, rendered_config: dict) -> list[str]:
+        res = set()
+        for field in cls.ds_name_fields:
+            if field in rendered_config:
+                ds_name = rendered_config[field]
+                res.add(ds_name)
+            elif "." in field:
+                tmp_rendered_config = rendered_config
+                for sub_field in field.split("."):
+                    if sub_field not in tmp_rendered_config:
+                        break
+                    tmp_rendered_config = tmp_rendered_config[sub_field]
+                else:
+                    if isinstance(tmp_rendered_config, str):
+                        ds_name = tmp_rendered_config
+                        res.add(ds_name)
+        return list(res)
+
+    @cached_property
+    def rendered_config(self) -> AttrDict:
+        return self.render_config()
+
+    def render_config(self) -> AttrDict:
+        result = {}
+        env = Renderer()
+        ctx = self.get_template_context()
+
+        for k, v in self.config.items():
+            if v is None or k in self.__class__.no_template_fields or not isinstance(v, (str, dict, list, tuple)):
+                result[k] = v
+            else:
+                result[k] = env.render_template(v, ctx)
+        return AttrDict(result)
+
+    def get_template_context(self) -> dict[str, Any]:
+        ctx = Renderer.init_context(self.execution_date, self.dag.schedule_interval)
+        ctx.update(self.variables)
+        return ctx
+
+    def execute(self, *args, **kwargs):
+        # TODO: create new task instance, send request to server or message queue?
+
+        self.on_task_start()
+
+        self.before_execute_hook()
+
+        error = None
+        meta = None
+        error_stack = None
+
+        logger.info("task configuration: %s", json.dumps(self.rendered_config, indent=2, ensure_ascii=False))
+        try:
+            meta = self.execute_impl(*args, **kwargs)
+        except Exception as exc:
+            error = exc
+            error_stack = exc.__repr__()
+            self.on_execute_impl_error(exc)
+
+        self.after_execute_hook()
+
+        self.on_task_finish(meta, error, error_stack)  # todo: try except?
+
+        if error is not None:
+            raise error
+
+    def on_task_start(self):
+        self.task_instance_id = context.init_task_instance_on_task_start(self)
+
+    def on_task_finish(self, meta: Any, error: Exception, error_stack: str):
+        try:
+            if meta:
+                meta = meta.to_json()
+        except Exception as e:
+            logger.debug(f"failed to get json from meta {meta}, error: {e}")
+            meta = None
+        if error_stack:
+            task_status = ETLExecutionStatus.FAILED
+        else:
+            task_status = ETLExecutionStatus.SUCCESS
+        context.update_task_instance_on_task_finish(self, self.task_instance_id, task_status, meta, error, error_stack)
+
+    def before_execute_hook(self):
+        pass
+
+    def after_execute_hook(self):
+        pass
+
+    def on_execute_impl_error(self, exc: Exception):
+        """callback function to be called if `execute_impl` throws exceptions"""
+        pass
+
+    def execute_impl(self, *args, **kwargs):
+        raise NotImplementedError
+
+    def get_query_comment_conf(self) -> str:
+        query_config = {
+            "Source": "Recurve",
+            "Owner": self.dag.owner,
+            "Node": self.node_url,
+        }
+        return ", ".join(["{}: {}".format(k, v) for k, v in query_config.items()])
+
+    def set_execution_date(self, execution_date):
+        if execution_date == self.execution_date:
+            return
+        _ = self.rendered_config
+        # rendered_config 依赖 self.execution_date
+        # 需要用旧的 execution_date 渲染后，再替换掉 self.execution_date
+        self.execution_date = execution_date
+
+    @property
+    def node_url(self) -> str:
+        # https://dev-test.recurve.test.recurvedata.com/datawork/workflow?p_id=257942399102349312&wf_id=258282502478635008&open_drawer=true&node_key=D2f0I
+        host = context.client.base_url  # todo: correct it
+        query_string = urllib.parse.urlencode(
+            {"p_id": self.dag.project_id, "job_id": self.dag.id, "node_key": self.node.node_key, "open_drawer": "true"}
+        )
+        return f"{host}/datawork/workspace/job?{query_string}"
+
+    # add proxy methods to avoid importing context everywhere
+
+    @staticmethod
+    def get_connection_by_name(name: str):
+        return context.get_connection_by_name(name)
+
+    @staticmethod
+    def must_get_connection_by_name(name: str):
+        return context.must_get_connection_by_name(name)
+
+    @staticmethod
+    def get_connection_names_by_type(connection_type: Union[str, list[str]]) -> list[str]:
+        return context.get_connection_names_by_type(connection_type)
+
+    @property
+    def stage(self) -> Optional[str]:
+        return
+
+    @property
+    def client(self) -> "ExecutorClient":
+        return context.client

recurvedata/operators/transfer_operator/__init__.py

@@ -0,0 +1,40 @@
+from recurvedata.operators.transfer_operator.dump_aliyun_sls import AliyunSLSDumpTask
+from recurvedata.operators.transfer_operator.dump_task_dbapi import DBAPIDumpTask
+from recurvedata.operators.transfer_operator.dump_task_es import ElasticSearchDumpTask
+from recurvedata.operators.transfer_operator.dump_task_feishu_sheet import FeishuSheetDumpTask
+from recurvedata.operators.transfer_operator.dump_task_ftp import FTPDumpTask
+from recurvedata.operators.transfer_operator.dump_task_google_sheet import GoogleSheetDumpTask
+
+# from recurvedata.operators.transfer_operator.dump_task_cass import CassandraDumpTask
+from recurvedata.operators.transfer_operator.dump_task_mongodb import MongoDBDumpTask
+from recurvedata.operators.transfer_operator.dump_task_oss import AliyunOSSDumpTask
+from recurvedata.operators.transfer_operator.dump_task_python import PythonDumpTask
+from recurvedata.operators.transfer_operator.dump_task_s3 import S3DumpTask
+from recurvedata.operators.transfer_operator.dump_task_sftp import SFTPDumpTask
+from recurvedata.operators.transfer_operator.load_task_aliyun_oss import AliyunOSSLoadTask
+from recurvedata.operators.transfer_operator.load_task_azure_blob import AzureBlobStorageLoadTask
+from recurvedata.operators.transfer_operator.load_task_clickhouse import ClickHouseLoadTask
+
+# from recurvedata.operators.transfer_operator.load_task_filebrowser import FileBrowserLoadTask
+# from recurvedata.operators.transfer_operator.load_task_hive import HiveLoadTask
+# from recurvedata.operators.transfer_operator.load_task_owncloud import OwnCloudLoadTask
+# from recurvedata.operators.transfer_operator.load_task_recurve_data_prep import DataPrepLoadTask
+# from recurvedata.operators.transfer_operator.load_task_yicrowds import YiCrowdsLoadTask
+# from recurvedata.operators.transfer_operator.load_task_azure_synapse import AzureSynapseLoadTask
+# from recurvedata.operators.transfer_operator.load_task_email import EmailLoadTask
+from recurvedata.operators.transfer_operator.load_task_doris import DorisLoadTask
+from recurvedata.operators.transfer_operator.load_task_es import ElasticSearchLoadTask
+from recurvedata.operators.transfer_operator.load_task_ftp import FTPLoadTask
+from recurvedata.operators.transfer_operator.load_task_google_bigquery import GoogleBigqueryLoadTask
+from recurvedata.operators.transfer_operator.load_task_google_cloud_storage import GoogleCloudStorageLoadTask
+from recurvedata.operators.transfer_operator.load_task_google_sheet import GoogleSheetLoadTask
+from recurvedata.operators.transfer_operator.load_task_microsoft_fabric import MicrosoftFabricLoadTask
+from recurvedata.operators.transfer_operator.load_task_mssql import MsSQLLoadTask
+from recurvedata.operators.transfer_operator.load_task_mysql import MySQLLoadTask
+from recurvedata.operators.transfer_operator.load_task_postgresql import PostgresqlLoadTask
+from recurvedata.operators.transfer_operator.load_task_qcloud_cos import TencentCOSLoadTask
+from recurvedata.operators.transfer_operator.load_task_redshift import RedshiftLoadTask
+from recurvedata.operators.transfer_operator.load_task_s3 import S3LoadTask
+from recurvedata.operators.transfer_operator.load_task_sftp import SFTPLoadTask
+from recurvedata.operators.transfer_operator.load_task_starrocks import StarRocksLoadTask
+from recurvedata.operators.transfer_operator.operator import TransferOperator

recurvedata/operators/transfer_operator/const.py

@@ -0,0 +1,10 @@
+LOAD_RENAME_OVERWRITE = "RENAME_OVERWRITE"
+LOAD_OVERWRITE = "OVERWRITE"
+LOAD_MERGE = "MERGE"
+LOAD_APPEND = "APPEND"
+
+
+FILE_TRANSFORM_FUNC_DEFAULT_VALUE = """\
+def transform(filename: str) -> str:
+    return filename
+"""

recurvedata/operators/transfer_operator/dump_aliyun_sls.py

@@ -0,0 +1,82 @@
+import copy
+
+from recurvedata.core.translation import _l
+from recurvedata.operators.transfer_operator import utils
+from recurvedata.operators.transfer_operator.task import DumpTask
+from recurvedata.pigeon.dumper.aliyun_sls import AliyunSLSDumper
+from recurvedata.utils import extract_dict
+
+
+class AliyunSLSDumpTask(DumpTask):
+    ds_name_fields = ("data_source_name",)
+
+    def execute_impl(self, *args, **kwargs):
+        ds = self.must_get_connection_by_name(self.config["data_source_name"])
+        hf = self.create_handler_factory()
+        dump_options = extract_dict(
+            self.rendered_config, keys=["project", "logstore", "query", "start_time", "end_time", "fields"]
+        )
+        access_key_id = ds.data.get("access_key_id")
+        access_key_secret = ds.data.get("access_key_secret")
+        endpoint = ds.data.get("endpoint")
+        dump_options.update(
+            {
+                "endpoint": endpoint,
+                "access_key_id": access_key_id,
+                "access_key_secret": access_key_secret,
+                "handler_factories": [hf],
+            }
+        )
+        dumper = AliyunSLSDumper(**dump_options)
+        return dumper.execute()
+
+    @classmethod
+    def config_schema(cls):
+        return {
+            "type": "object",
+            "properties": {
+                "data_source_name": {
+                    "type": "string",
+                    "title": _l("Aliyun Access Key"),
+                    "ui:field": "ProjectConnectionSelectorField",
+                    "ui:options": {
+                        "supportTypes": ["aliyun_access_key"],
+                    },
+                },
+                "project": {"type": "string", "title": _l("Project Name")},
+                "logstore": {"type": "string", "title": _l("Logstore Name")},
+                "query": {
+                    "type": "string",
+                    "title": _l("Query"),
+                    "description": _l("Query to retrieve logs from Aliyun SLS."),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "code",
+                        "lang": "sql",
+                    },
+                },
+                "start_time": {
+                    "type": "string",
+                    "description": _l(
+                        "Start time of the data to retrieve, supports Jinja templating for dynamic. Format: %Y-%m-%d %H:%M:%S"
+                    ),
+                    "title": _l("Start Time"),
+                    "default": "{{ data_interval_start }}",
+                },
+                "end_time": {
+                    "type": "string",
+                    "description": _l(
+                        "End time of the data to retrieve, supports Jinja templating for dynamic. Format: %Y-%m-%d %H:%M:%S"
+                    ),
+                    "title": _l("End Time"),
+                    "default": "{{ data_interval_end }}",
+                },
+                "fields": {
+                    "type": "string",
+                    "title": _l("Fields"),
+                    "description": _l("Comma-separated list of fields to retrieve. Leave empty to get all fields."),
+                },
+                "transform": copy.deepcopy(utils.TRANSFORM),
+            },
+            "required": ["data_source_name", "project", "logstore", "start_time", "end_time"],
+        }

recurvedata/operators/transfer_operator/dump_sheet_task_base.py

@@ -0,0 +1,292 @@
+import inspect
+import json
+import logging
+import os
+
+import jsonschema
+
+from recurvedata.pigeon.handler.csv_handler import CSVFileHandler
+from recurvedata.pigeon.utils import ensure_str_list, fs
+
+try:
+    import numpy as np
+    import pandas as pd
+except ImportError:
+    pass
+
+from recurvedata.core.translation import _l
+from recurvedata.operators.transfer_operator.task import DumpTask
+from recurvedata.operators.utils import infer_schema_from_dataframe, parse_to_date
+from recurvedata.utils.attrdict import AttrDict
+
+logger = logging.getLogger(__name__)
+_transform_default_value = """\
+import pandas as pd
+
+
+def transform(df: pd.DataFrame) -> pd.DataFrame:
+    return df
+"""
+
+
+class SheetDumpTaskBase(DumpTask):
+    _AUTO_REGISTER = False
+
+    common_config_schema_properties = {
+        "extra_read_kwargs": {
+            "type": "string",
+            "title": _l("Additional Read Parameters"),
+            "description": _l(
+                "Additional parameters to pass to pandas read_csv or read_excel functions in JSON format"
+            ),
+            "ui:field": "CodeEditorWithReferencesField",
+            "ui:options": {
+                "type": "code",
+                "lang": "json",
+            },
+        },
+        "type_mapping": {
+            "type": "string",
+            "title": _l("Column Type Mapping"),
+            "description": _l(
+                'Specify data types for columns using format {"column_name": "data_type"}. '
+                "This mapping is passed to DataFrame.astype() - see "
+                '<a target="_blank" href="https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.astype.html">'
+                "pandas documentation</a> for supported types."
+            ),
+            "ui:field": "CodeEditorWithReferencesField",
+            "ui:options": {
+                "type": "code",
+                "lang": "json",
+            },
+        },
+        "date_columns": {
+            "type": "string",
+            "title": _l("Date Format Columns"),
+            "description": _l("Comma-separated list of column names to parse as dates"),
+            "ui:field": "CodeEditorWithReferencesField",
+            "ui:options": {
+                "type": "plain",
+            },
+        },
+        "fillna_to_null": {
+            "type": "boolean",
+            "title": _l("Convert NaN to NULL"),
+            "default": True,
+        },
+        "order_by": {
+            "type": "string",
+            "title": _l("Sort Order"),
+            "description": _l(
+                "Comma-separated list of columns to sort rows by. Original order is preserved if not specified."
+            ),
+            "ui:field": "CodeEditorWithReferencesField",
+            "ui:options": {
+                "type": "plain",
+            },
+        },
+        "column_name_mapping": {
+            "type": "string",
+            "title": _l("Rename Columns"),
+            "description": _l('Map old column names to new names using format {"old_name": "new_name"}'),
+            "ui:field": "CodeEditorWithReferencesField",
+            "ui:options": {
+                "type": "code",
+                "lang": "json",
+            },
+        },
+        "result_columns": {
+            "type": "string",
+            "title": _l("Output Columns"),
+            "description": _l(
+                "Comma-separated list of columns to include in output and their order. All columns included if not specified."
+            ),
+            "ui:field": "CodeEditorWithReferencesField",
+            "ui:options": {
+                "type": "plain",
+            },
+        },
+        "primary_keys": {
+            "type": "string",
+            "title": _l("Unique Key Columns"),
+            "description": _l(
+                "Comma-separated list of columns that should contain unique values. "
+                "Task will fail if duplicates are found. Leave empty to skip uniqueness check."
+            ),
+            "ui:field": "CodeEditorWithReferencesField",
+            "ui:options": {
+                "type": "plain",
+            },
+        },
+        "not_nullable_columns": {
+            "type": "string",
+            "title": _l("Required Columns"),
+            "description": _l(
+                "Comma-separated list of columns that must not contain NULL values. "
+                "Task will fail if NULL values are found in these columns."
+            ),
+            "ui:field": "CodeEditorWithReferencesField",
+            "ui:options": {
+                "type": "plain",
+            },
+        },
+        "transform_func": {
+            "type": "string",
+            "title": _l("Custom Transform"),
+            "description": _l(
+                "Optional Python function to transform the DataFrame. Must accept and return a pandas DataFrame. "
+                "This transformation is applied after all other processing steps."
+            ),
+            "default": _transform_default_value,
+            "ui:field": "CodeEditorWithReferencesField",
+            "ui:options": {
+                "type": "code",
+                "lang": "python",
+            },
+        },
+    }
+
+    custom_config_schema_properties = {}
+
+    custom_config_schema_required = []
+
+    @classmethod
+    def config_schema(cls):
+        schema = {
+            "type": "object",
+            "properties": {},
+            "required": cls.custom_config_schema_required,
+        }
+        schema["properties"].update(cls.custom_config_schema_properties)
+        for k, v in cls.common_config_schema_properties.items():
+            if k not in schema["properties"]:
+                schema["properties"][k] = v
+        return schema
+
+    def execute_impl(self, *args, **kwargs):
+        conf = self.rendered_config
+        df = self.read_origin_df()
+
+        df = self.apply_builtin_transform(conf, df)
+        df = self.apply_validations(conf, df)
+        df = self.apply_custom_transform_func(conf, df)
+
+        self.df_to_csv(df)
+
+    def read_origin_df(self) -> "pd.DataFrame":
+        raise NotImplementedError
+
+    def df_to_csv(self, df):
+        logger.info(f"result DataFrame shape {df.shape}, dtypes:\n{df.dtypes}")
+        logger.info(df.head())
+
+        handler: CSVFileHandler = self.create_handler_factory().create_handler()
+        for row in df.itertuples(index=False):
+            handler.handle(row)
+        handler.close()
+        if handler.filename != self.filename and os.path.exists(handler.filename):
+            os.rename(handler.filename, self.filename)
+        logger.info(f"exported {len(df)} rows into {self.filename}")
+
+        schema = infer_schema_from_dataframe(df)
+        schema_filename = fs.schema_filename(self.filename)
+        schema.dump(schema_filename)
+        logger.info(f"saving schema to {schema_filename}")
+
+    @staticmethod
+    def apply_builtin_transform(conf: AttrDict, df: "pd.DataFrame") -> "pd.DataFrame":
+        logger.info("apply_builtin_transform...")
+        if conf.type_mapping:
+            logger.info(f"  * convert dtypes with {conf.type_mapping}")
+            df = df.astype(json.loads(conf.type_mapping))
+
+        if conf.date_columns:
+            cols = ensure_str_list(conf.date_columns)
+            logger.info(f"  * parse {cols} to date")
+            for col in cols:
+                df[col] = df[col].map(parse_to_date)
+
+        if conf.fillna_to_null:
+            logger.info("  * fillna with None")
+            df = df.fillna(np.nan).replace([np.nan], [None])
+
+        if conf.order_by:
+            cols = ensure_str_list(conf.order_by)
+            logger.info("  * sort by {cols")
+            df = df.sort_values(by=cols)
+
+        if conf.column_name_mapping:
+            logger.info(f"  * apply column name mapping {conf.column_name_mapping}")
+            df = df.rename(json.loads(conf.column_name_mapping), axis=1)
+
+        if conf.result_columns:
+            cols = ensure_str_list(conf.result_columns)
+            logger.info(f"  * change result columns with {cols}")
+            df = df[cols]
+
+        return df
+
+    @staticmethod
+    def apply_validations(conf: AttrDict, df: "pd.DataFrame") -> "pd.DataFrame":
+        logger.info("apply_validations...")
+        if conf.primary_keys:
+            logger.info("  * checking duplication...")
+            duplicate = df[df.duplicated(subset=ensure_str_list(conf.primary_keys))]
+            if not duplicate.empty:
+                logger.error(f"duplicate rows: {duplicate}")
+                raise ValueError("duplication detected")
+
+        if conf.not_nullable_columns:
+            cols = ensure_str_list(conf.not_nullable_columns)
+            logger.info(f"  * checking null to columns {cols}...")
+            null_cols = []
+            for col in cols:
+                if df[col].isnull().values.any():
+                    null_cols.append(col)
+            if null_cols:
+                logger.error(f"{null_cols} contains null values")
+                raise ValueError(f"{null_cols} contains null values")
+
+        return df
+
+    @staticmethod
+    def apply_custom_transform_func(conf: AttrDict, df: "pd.DataFrame") -> "pd.DataFrame":
+        if not conf.transform_func:
+            return df
+        func = validate_transform(conf.transform_func)
+        if not func:
+            return df
+
+        logger.info("apply transform function...")
+        df = func(df)
+        if not isinstance(df, pd.DataFrame):
+            raise ValueError(f"transform function must return an Pandas DataFrame object, got {type(df)} instead")
+        return df
+
+    @classmethod
+    def validate(cls, configuration):
+        conf = super().validate(configuration)
+
+        transform_func_code = conf.get("transform_func", "").strip()
+        if transform_func_code:
+            validate_transform(transform_func_code)
+        return conf
+
+
+def validate_transform(raw_code):
+    code = compile(raw_code, "", "exec")
+    ns = {}
+    exec(code, ns)
+    func = ns.get("transform")
+    if not func:
+        return None
+
+    if not callable(func):
+        raise jsonschema.ValidationError(message="transform should be callable", path=("transform_func",))
+
+    sig = inspect.signature(func)
+    if tuple(sig.parameters.keys()) != ("df",):
+        raise jsonschema.ValidationError(
+            message="transform must accept and only accept df as parameter", path=("transform_func",)
+        )
+    return func