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

recurvedata/operators/transfer_operator/dump_task_cass.py

@@ -0,0 +1,155 @@
+import copy
+
+try:
+    import arrow
+
+    from recurvedata.pigeon.dumper.cass import CassandraDumper
+    from recurvedata.pigeon.row_factory import ordered_dict_factory
+except ImportError:
+    pass
+
+from recurvedata.operators.transfer_operator import utils
+from recurvedata.operators.transfer_operator.task import DumpTask
+from recurvedata.utils import extract_dict
+
+
+class CassandraDumpTask(DumpTask):
+    enabled = False
+    worker_install_require = ["pigeon[cassandra]"]
+    ds_name_fields = ("data_source_name",)
+
+    def determine_partitions(self):
+        if self.config.partitions:
+            return self.config.partitions
+
+        if not self.config.incremental_by_time:
+            return None
+
+        if self.dag.is_once:
+            return None
+
+        start_date, end_date = self.get_schedule_time_range()
+        partitions = arrow.Arrow.range(self.config.time_granularity, start_date, end_date)
+        partitions = [x.datetime for x in partitions[:-1]]
+        return partitions
+
+    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=["table", "columns", "where", "partition_column", "concurrency"]
+        )
+        dump_options.update({"connector": ds.connector, "handler_factories": [hf]})
+        partitions = self.determine_partitions()
+        if partitions:
+            dump_options.update({"partitions": partitions})
+        dumper = CassandraDumper(**dump_options)
+        # if self.has_custom_transformer():
+        dumper.row_factory = ordered_dict_factory
+        return dumper.execute()
+
+    @classmethod
+    def config_schema(cls):
+        # dss = cls.get_connection_names_by_type('cassandra')
+        return {
+            "type": "object",
+            "properties": {
+                "data_source_name": {
+                    "type": "string",
+                    "title": "Data Source",
+                    "ui:field": "ProjectConnectionSelectorField",
+                    "ui:options": {
+                        "supportTypes": [
+                            "cassandra",
+                        ],
+                    },
+                    # 'default': cls.first_or_default(dss, ''),
+                },
+                "table": {
+                    "type": "string",
+                    "title": "Table Name",
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "columns": {
+                    "type": "string",
+                    "title": "Columns",
+                    "description": "要导出的列，用 `,` 分隔；默认导出所有列（*）",
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "where": {
+                    "type": "string",
+                    "title": "Where Clause",
+                    "description": "Where 条件",
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "partition_column": {
+                    "type": "string",
+                    "title": "Partition Column",
+                    "description": "分区键，通常名为 date",
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "partitions": {
+                    "type": "string",
+                    "title": "Partitions",
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "concurrency": {
+                    "type": "number",
+                    "ui:options": {"controls": False},
+                    "title": "Concurrency",
+                    "default": 1,
+                    "description": "并发数，1~20",
+                    "minimum": 1,
+                    "maximum": 20,
+                },
+                "transform": copy.deepcopy(utils.TRANSFORM),
+                "incremental_by_time": {
+                    "type": "boolean",
+                    "title": "Incremental By Time",
+                    "default": False,
+                    "description": "是否按时间进行增量同步，这个时间必须是分区键",
+                    "ui:widget": "BaseCheckbox",
+                    "ui:options": {
+                        "label": "Incremental By Time",
+                    },
+                },
+                "time_granularity": {
+                    "ui:hidden": "{{!parentFormData.incremental_by_time}}",
+                    "type": "string",
+                    "title": "Time Granularity",
+                    "default": "day",
+                    "description": "分区键的时间粒度，用于生成分区值",
+                    "enum": ["day", "hour"],
+                    "enumNames": ["day", "hour"],
+                },
+                "time_auto_round": {
+                    "ui:hidden": "{{!parentFormData.incremental_by_time}}",
+                    "type": "boolean",
+                    "title": "Round Time Resolution",
+                    "default": True,
+                    "description": "是否把数据时间范围 round 到合适的粒度。比如每天 01:23 同步上一个自然日的数据，"
+                    "则运行时间是 01:23，数据范围是 [T-1 00:00, T 00:00)；否则数据范围是 [T-1 01:23, T 01:23)。"
+                    "开启后，每天运行的任务，数据范围会 round 到 0 点，即自然日；"
+                    "每周运行的任务，会 round 到周一 0 点；"
+                    "每月运行的任务，会 round 到每月 1 日 0 点",
+                },
+            },
+            # NOTE：前端用的 vue-json-schema 有 bug，enum 字段必须被 required...
+            "required": ["data_source_name", "table", "time_granularity"],
+            # 处理表单联动，只有 incremental_by_time 为 True 时，才需要显示其他两个输入框
+        }

recurvedata/operators/transfer_operator/dump_task_dbapi.py

@@ -0,0 +1,209 @@
+import copy
+
+import jsonschema
+
+from recurvedata.connectors.service import list_sql_operator_types
+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.dbapi import DBAPIDumper
+from recurvedata.pigeon.row_factory import ordered_dict_factory
+from recurvedata.pigeon.utils.sql import apply_where_safely
+from recurvedata.utils import date_time, extract_dict
+
+
+class DBAPIDumpTask(DumpTask):
+    ds_name_fields = ("data_source_name",)
+    worker_install_require = ["pigeon"]
+
+    @property
+    def time_column_tz(self):
+        return self.config.get("time_column_tz", "UTC")
+
+    def determine_time_range(self):
+        start_date, end_date = self.get_schedule_time_range()
+        if self.config.time_column_type == "date":
+            return start_date.date(), end_date.date()
+
+        # convert timezone
+        start_date = date_time.astimezone(start_date, tz=self.time_column_tz)
+        end_date = date_time.astimezone(end_date, tz=self.time_column_tz)
+
+        return start_date.replace(tzinfo=None), end_date.replace(tzinfo=None)
+
+    def derive_sql_query(self, connector, base_query: str):
+        base_query = base_query.strip().rstrip(";")
+        comment = self.get_query_comment_conf()
+        if not self.config.incremental_by_time or self.dag.is_once:
+            annotated_query = connector.add_leading_comment(base_query, comment)
+            return annotated_query
+
+        if not base_query:
+            base_query = f"SELECT * FROM {connector.quote_identifier(self.config.table)}"
+        annotated_query = connector.add_leading_comment(base_query, comment)
+
+        start, end = self.determine_time_range()
+        col = connector.quote_identifier(self.config.time_column)
+
+        if connector.is_phoenix():
+            where = f"{col} >= TIMESTAMP '{start}' AND {col} < TIMESTAMP '{end}'"
+        else:
+            where = f"{col} >= '{start}' AND {col} < '{end}'"
+        return apply_where_safely(annotated_query, where)
+
+    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=["table", "splitby", "splits", "concurrency"])
+        dump_options.update(
+            {
+                "connector": ds.connector,
+                "sql": self.derive_sql_query(ds.connector, self.rendered_config.get("sql", "")),
+                "handler_factories": [hf],
+            }
+        )
+        if not dump_options.get("splitby"):
+            dump_options["splits"] = dump_options["concurrency"] = 1
+        dumper = DBAPIDumper(**dump_options)
+        dumper.row_factory = ordered_dict_factory
+        return dumper.execute()
+
+    @classmethod
+    def validate(cls, configuration):
+        conf = super().validate(configuration)
+        if not (conf.get("table") or conf.get("sql")):
+            raise jsonschema.ValidationError(message="either table or sql is required", path=("table", "sql"))
+        return conf
+
+    @classmethod
+    def config_schema(cls):
+        return {
+            "type": "object",
+            "properties": {
+                "data_source_name": {
+                    "type": "string",
+                    "title": _l("Data Source"),
+                    "description": _l("Database connection to extract data from"),
+                    "ui:field": "ProjectConnectionSelectorField",
+                    "ui:options": {
+                        "supportTypes": list_sql_operator_types(),
+                    },
+                },
+                "table": {
+                    "type": "string",
+                    "title": _l("Source Table"),
+                    "description": _l(
+                        "Table name including schema (if required). Either specify a table name or SQL query."
+                    ),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "sql": {
+                    "type": "string",
+                    "title": _l("Custom Query"),
+                    "description": _l(
+                        "Custom SELECT query with Jinja template support. Takes precedence over table name if both are specified."
+                    ),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "code",
+                        "lang": "sql",
+                        "sqlLang": "sql",
+                    },
+                },
+                "splitby": {
+                    "type": "string",
+                    "title": _l("Split Column"),
+                    "description": _l(
+                        "Column to partition data by for parallel processing. Must be indexed, sortable and non-null."
+                    ),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "splits": {
+                    "ui:hidden": "{{ !parentFormData.splitby }}",
+                    "type": "number",
+                    "ui:options": {"controls": False},
+                    "title": _l("Number of Splits"),
+                    "default": 1,
+                    "minimum": 1,
+                    "maximum": 2000,
+                },
+                "concurrency": {
+                    "ui:hidden": "{{ !parentFormData.splitby }}",
+                    "type": "number",
+                    "ui:options": {"controls": False},
+                    "title": _l("Parallel Threads"),
+                    "default": 1,
+                    "description": _l("Number of concurrent extraction threads (1-20)"),
+                    "minimum": 1,
+                    "maximum": 20,
+                },
+                "transform": copy.deepcopy(utils.TRANSFORM),
+                "incremental_by_time": {
+                    "type": "boolean",
+                    "title": _l("Enable Time-based Incremental Sync"),
+                    "default": False,
+                    "description": _l("Sync data incrementally based on a time column"),
+                    "ui:widget": "BaseCheckbox",
+                    "ui:options": {
+                        "label": _l("Enable Time-based Incremental Sync"),
+                    },
+                },
+                "time_column": {
+                    "ui:hidden": "{{!parentFormData.incremental_by_time}}",
+                    "type": "string",
+                    "title": _l("Time Column Name"),
+                    "default": "snapshot_time",
+                    "description": _l(
+                        "Name of the time column used for incremental sync. Column should be indexed for better performance."
+                    ),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "time_column_tz": {
+                    "ui:hidden": "{{!parentFormData.incremental_by_time}}",
+                    "type": "string",
+                    "title": _l("Time Column Timezone"),
+                    "default": "UTC",
+                    "enum": [
+                        "UTC",
+                        "Asia/Shanghai",
+                    ],
+                    "enumNames": [
+                        "UTC",
+                        "Asia/Shanghai",
+                    ],
+                },
+                "time_column_type": {
+                    "ui:hidden": "{{!parentFormData.incremental_by_time}}",
+                    "type": "string",
+                    "title": _l("Timestamp Format"),
+                    "default": "datetime",
+                    "enum": ["datetime", "date"],
+                    "enumNames": ["datetime", "date"],
+                },
+                "time_auto_round": {
+                    "ui:hidden": "{{!parentFormData.incremental_by_time}}",
+                    "type": "boolean",
+                    "title": "Auto Round Time Range",
+                    "default": True,
+                    "description": _l(
+                        "Automatically round time ranges to appropriate intervals. For example:\n"
+                        "- Daily tasks running at 01:23 will sync previous day's data from 00:00 to 00:00\n"
+                        "- Weekly tasks will round to Monday 00:00\n"
+                        "- Monthly tasks will round to 1st day 00:00\n"
+                        "If disabled, exact execution times will be used (e.g. 01:23 to 01:23)"
+                    ),
+                },
+            },
+            "required": [
+                "data_source_name",
+            ],
+        }

recurvedata/operators/transfer_operator/dump_task_es.py

@@ -0,0 +1,113 @@
+import copy
+import json
+
+try:
+    from recurvedata.pigeon.dumper.es import ElasticSearchDumper
+except ImportError:
+    pass
+
+from recurvedata.core.translation import _l
+from recurvedata.operators.transfer_operator import utils
+from recurvedata.operators.transfer_operator.task import DumpTask
+from recurvedata.utils import extract_dict
+
+
+class ElasticSearchDumpTask(DumpTask):
+    ds_name_fields = ("data_source_name",)
+    worker_install_require = ["pigeon[elasticsearch]"]
+
+    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=["index", "doc_type", "query", "fields", "meta_fields"])
+        if self.rendered_config.get("search_kwargs"):
+            search_kwargs = json.loads(self.rendered_config.get("search_kwargs"))
+            dump_options["search_kwargs"] = search_kwargs
+        dump_options.update({"connector": ds.connector, "handler_factories": [hf]})
+        dumper = ElasticSearchDumper(**dump_options)
+        return dumper.execute()
+
+    @classmethod
+    def config_schema(cls):
+        # get_choices_by_type = cls.get_connection_names_by_type
+        # dss = get_choices_by_type('elasticsearch')
+        return {
+            "type": "object",
+            "properties": {
+                "data_source_name": {
+                    "type": "string",
+                    "title": _l("Elasticsearch Data Source"),
+                    "ui:field": "ProjectConnectionSelectorField",
+                    "ui:options": {
+                        "supportTypes": [
+                            "elasticsearch",
+                        ],
+                    },
+                    # 'default': cls.first_or_default(dss, ''),
+                },
+                "index": {
+                    "type": "string",
+                    "title": _l("Elasticsearch Index"),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "doc_type": {
+                    "type": "string",
+                    "title": _l("Document Type"),
+                    "default": "_doc",
+                    "description": _l("The type of documents to query"),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "query": {
+                    "type": "string",
+                    "title": _l("Search Query"),
+                    "default": "*",
+                    "description": _l("Elasticsearch query string to filter documents. Supports Jinja templating."),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "fields": {
+                    "type": "string",
+                    "title": _l("Document Fields"),
+                    "description": _l(
+                        "Comma-separated list of document fields to retrieve. Leave empty to get all fields."
+                    ),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "meta_fields": {
+                    "type": "array",
+                    "uniqueItems": True,
+                    "items": {
+                        "type": "string",
+                        "enum": ["_index", "_type", "_id"],
+                        "enumNames": ["_index", "_type", "_id"],
+                    },
+                    "title": _l("Document Metadata Fields"),
+                    "description": _l("Additional metadata fields to include with each document."),
+                    "ui:widget": "SelectWidget",
+                },
+                "search_kwargs": {
+                    "type": "string",
+                    "title": _l("Advanced Search Options"),
+                    "description": _l(
+                        "Additional options for Elasticsearch scan operation in JSON format (e.g. size, scroll)."
+                    ),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "transform": copy.deepcopy(utils.TRANSFORM),
+            },
+            "required": ["data_source_name", "index"],
+        }

recurvedata/operators/transfer_operator/dump_task_feishu_sheet.py

@@ -0,0 +1,114 @@
+import json
+import logging
+import urllib.parse
+
+try:
+    import pandas as pd
+
+    from recurvedata.pigeon.connector.feishu import FeishuBot
+except ImportError:
+    pass
+
+from recurvedata.core.translation import _l
+from recurvedata.operators.transfer_operator.dump_sheet_task_base import SheetDumpTaskBase
+
+logger = logging.getLogger(__name__)
+
+
+class FeishuSheetDumpTask(SheetDumpTaskBase):
+    _AUTO_REGISTER = True
+    ds_name_fields = ("feishu_bot",)
+    worker_install_require = ["pandas", "numpy", "pigeon[feishu]"]
+
+    custom_config_schema_properties = {
+        "feishu_bot": {
+            "type": "string",
+            "title": _l("Feishu Bot Connection"),
+            "description": _l(
+                "Select the Feishu bot connection that has permissions to access and read the target sheet"
+            ),
+            "ui:field": "ProjectConnectionSelectorField",
+            "ui:options": {
+                "supportTypes": [
+                    "feishu_bot",
+                ],
+            },
+            # 'default': cls.first_or_default(dss, ''),
+        },
+        "file_url": {
+            "type": "string",
+            "title": _l("Feishu Document URL"),
+            "description": _l(
+                "URL of the Feishu spreadsheet or document to read data from. For multi-sheet documents, the first sheet will be used by default unless a specific sheet ID is included in the URL"
+            ),
+            "ui:field": "CodeEditorWithReferencesField",
+            "ui:options": {
+                "type": "plain",
+            },
+        },
+        "cell_range": {
+            "type": "string",
+            "title": _l("Data Range"),
+            "description": _l(
+                "Optional range of cells to read in A1 notation (e.g. A1:D100). Leave empty to read all data from the sheet"
+            ),
+            "ui:field": "CodeEditorWithReferencesField",
+            "ui:options": {
+                "type": "plain",
+            },
+        },
+    }
+    custom_config_schema_required = ["feishu_bot", "file_url"]
+
+    def read_origin_df(self) -> "pd.DataFrame":
+        conf = self.rendered_config
+        if conf.extra_read_kwargs:
+            extra_read_kwargs = json.loads(conf.extra_read_kwargs)
+        else:
+            extra_read_kwargs = {}
+
+        ds = self.must_get_connection_by_name(conf.feishu_bot)
+        bot = FeishuBot(**ds.extra)
+        file_type, file_token, sheet = self.parse_feishu_sheets_url(conf.file_url)
+        logger.info(f"reading {conf.file_url}")
+        if file_type == "file":
+            df = bot.read_feishuexcel(file_token, **extra_read_kwargs)
+        else:  # sheets
+            if not sheet:
+                logger.info("sheet_id not found in url, use the first sheet as default")
+                sheet_ids, _ = bot.get_sheet_ids(file_token)
+                sheet = sheet_ids[0]
+            if conf.cell_range:
+                sheet = f"{sheet}!{conf.cell_range}"
+            df = bot.read_feishusheet(file_token, sheet, **extra_read_kwargs)
+
+        logger.info(f"original DataFrame shape {df.shape}, dtypes:\n{df.dtypes}")
+        logger.info(df.head())
+        return df
+
+    @staticmethod
+    def parse_feishu_sheets_url(url: str) -> tuple[str, str, str]:
+        rv = urllib.parse.urlparse(url)
+        if "/file/" in url:
+            file_type = "file"
+        elif "/sheets/" in url:
+            file_type = "sheets"
+        elif "/wiki/" in url:
+            file_type = "wiki"
+        else:
+            raise ValueError(f"unsupported url {url}")
+
+        file_token = rv.path.rsplit("/", 1)[-1]
+        sheet = urllib.parse.parse_qs(rv.query).get("sheet") or None
+        if sheet:
+            sheet = sheet[0]
+
+        bot = FeishuBot()
+        if file_type == "wiki":
+            obj_type, obj_token = bot.get_wiki_type_token(wiki_token=file_token)
+            if obj_type == "sheet":
+                file_type, file_token = obj_type, obj_token
+            else:
+                raise ValueError(f"unsupported url {url}")
+
+        return file_type, file_token, sheet