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

recurvedata/operators/transfer_operator/dump_task_ftp.py

@@ -0,0 +1,234 @@
+import csv
+import inspect
+import logging
+import os
+
+import jsonschema
+
+from recurvedata.core.translation import _l
+from recurvedata.operators.transfer_operator.const import FILE_TRANSFORM_FUNC_DEFAULT_VALUE
+from recurvedata.operators.transfer_operator.mixin import HiveTextfileConverterMixin
+from recurvedata.operators.transfer_operator.task import DumpTask
+from recurvedata.operators.utils import file_factory as ff
+from recurvedata.utils import unescape_backslash
+
+logger = logging.getLogger(__name__)
+
+
+class FTPDumpTask(DumpTask, HiveTextfileConverterMixin):
+    ds_name_fields = ("data_source_name",)
+    worker_install_require = ["pigeon"]
+
+    def execute_impl(self, *args, **kwargs):
+        from recurvedata.pigeon.dumper.ftp import FtpDumper
+
+        conf = self.rendered_config
+
+        ds = self.must_get_connection_by_name(conf["data_source_name"])
+        dumper = FtpDumper(ds.connector, src=conf["filepath"], dst=self.filename)
+        meta = dumper.execute()
+
+        filename = self.process_file(conf)
+        filename = self.transform_file(conf, filename)
+
+        if filename != self.filename:
+            logger.info("renaming %s to %s", filename, self.filename)
+            os.rename(filename, self.filename)
+
+        # TODO: pigeon loader 要支持不同的文件格式
+        self.convert_csv_to_hive_text_if_needed()
+        return meta
+
+    def process_file(self, conf):
+        filename = self.filename
+        if conf["decompress"] == "Gzip":
+            logger.info("decompressing %s using gzip", self.filename)
+            filename = ff.gzip_decompress(self.filename, inplace=True)
+
+        skip_head_lines = conf.get("skip_head_lines", 0)
+        if conf["file_format"] == "Excel":
+            logger.info("converting Excel to CSV...")
+            filename = ff.convert_excel_to_csv(filename, skiprows=skip_head_lines, inplace=True)
+        if conf["file_format"] == "JSONLines":
+            logger.info("converting JSON lines to CSV...")
+            filename = ff.convert_jsonlines_to_csv(
+                filename, skiprows=skip_head_lines, src_encoding=conf["encoding"], inplace=True
+            )
+        if conf["file_format"] == "CSV":
+            logger.info("converting CSV dialect and encoding if necessary...")
+            dialect_options = self._get_custom_csv_options(conf)
+            filename = ff.convert_csv_dialect(
+                filename,
+                src_dialect_options=dialect_options,
+                skiprows=skip_head_lines,
+                src_encoding=conf["encoding"],
+                inplace=True,
+            )
+        return filename
+
+    def transform_file(self, conf, filename):
+        transform_func_code = conf.get("transform_func", "").strip()
+        if not transform_func_code:
+            return filename
+
+        func = validate_transform(transform_func_code)
+        if not func:
+            return filename
+
+        logger.info("calling transform function with %s", (filename,))
+        result_file = func(filename)
+        if result_file is None or not (isinstance(result_file, str) and os.path.isabs(result_file)):
+            raise ValueError("transform must return an absolute filepath, got %s instead", result_file)
+        logger.info("got %s", result_file)
+        return result_file
+
+    def _get_custom_csv_options(self, conf):
+        rv = {
+            "delimiter": unescape_backslash(conf["csv_delimiter"]),
+            "lineterminator": unescape_backslash(conf["csv_lineterminator"]),
+        }
+        quoting = conf["csv_quoting"]
+        rv["quoting"] = {
+            "QUOTE_ALL": csv.QUOTE_ALL,
+            "QUOTE_MINIMAL": csv.QUOTE_MINIMAL,
+            "QUOTE_NONE": csv.QUOTE_NONE,
+            "QUOTE_NONNUMERIC": csv.QUOTE_NONNUMERIC,
+        }[quoting]
+        return rv
+
+    @classmethod
+    def config_schema(cls):
+        # get_choices_by_type = cls.get_connection_names_by_type
+        return {
+            "type": "object",
+            "properties": {
+                "data_source_name": {
+                    "type": "string",
+                    "title": _l("FTP Connection"),
+                    "ui:field": "ProjectConnectionSelectorField",
+                    "ui:options": {
+                        "supportTypes": [
+                            "ftp",
+                        ],
+                    },
+                },
+                "filepath": {
+                    "type": "string",
+                    "title": _l("Source File Path"),
+                    "description": _l(
+                        "Absolute path to the file on FTP server (e.g. /path/to/file.csv). Supports Jinja templating."
+                    ),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "decompress": {
+                    "type": "string",
+                    "title": _l("Decompression Method"),
+                    "description": _l("Decompress downloaded file using specified method"),
+                    "enum": ["None", "Gzip"],
+                    "enumNames": ["None", "Gzip"],
+                    "default": "None",
+                },
+                "file_format": {
+                    "type": "string",
+                    "title": _l("Input Format"),
+                    "description": _l("Format of the source file to be converted to CSV"),
+                    "enum": ["CSV", "Excel", "JSONLines"],
+                    "enumNames": ["CSV", "Excel", "JSONLines"],
+                    "default": "CSV",
+                },
+                "skip_head_lines": {
+                    "type": "number",
+                    "ui:options": {"controls": False},
+                    "title": _l("Skip Header Rows"),
+                    "description": _l("Number of rows to skip from the beginning of the file"),
+                    "default": 0,
+                    "minimum": 0,
+                },
+                "encoding": {
+                    "ui:hidden": '{{parentFormData.file_format !== "CSV"}}',
+                    "type": "string",
+                    "title": _l("File Encoding"),
+                    "description": _l("Character encoding of the CSV file (e.g. utf-8, gbk)"),
+                    "default": "utf-8",
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "csv_delimiter": {
+                    "ui:hidden": '{{parentFormData.file_format !== "CSV"}}',
+                    "type": "string",
+                    "title": _l("Field Delimiter"),
+                    "description": _l("Character used to separate fields in the CSV file"),
+                    "default": ",",
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "csv_lineterminator": {
+                    "ui:hidden": '{{parentFormData.file_format !== "CSV"}}',
+                    "type": "string",
+                    "title": _l("Line Ending"),
+                    "description": _l("Character sequence used to terminate lines"),
+                    "enum": [r"\n", r"\r\n"],
+                    "enumNames": [r"\n", r"\r\n"],
+                    "default": r"\r\n",
+                },
+                "csv_quoting": {
+                    "ui:hidden": '{{parentFormData.file_format !== "CSV"}}',
+                    "type": "string",
+                    "title": _l("Field Quoting"),
+                    "description": _l("Strategy for quoting fields in the CSV file"),
+                    "enum": ["QUOTE_ALL", "QUOTE_MINIMAL", "QUOTE_NONE", "QUOTE_NONNUMERIC"],
+                    "enumNames": ["QUOTE_ALL", "QUOTE_MINIMAL", "QUOTE_NONE", "QUOTE_NONNUMERIC"],
+                    "default": "QUOTE_MINIMAL",
+                },
+                "transform_func": {
+                    "type": "string",
+                    "title": _l("Custom Transformation"),
+                    "description": _l(
+                        "Python function to transform the downloaded file. Must accept a filepath argument and return "
+                        "the path to the transformed file. Runs after built-in transformations."
+                    ),
+                    "default": FILE_TRANSFORM_FUNC_DEFAULT_VALUE,
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "code",
+                        "lang": "python",
+                    },
+                },
+            },
+            "required": ["data_source_name", "filepath"],
+        }
+
+    @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()) != ("filename",):
+        raise jsonschema.ValidationError(
+            message="transform must accept and only accept filename as parameter", path=("transform_func",)
+        )
+    return func

recurvedata/operators/transfer_operator/dump_task_google_sheet.py

@@ -0,0 +1,66 @@
+import logging
+
+try:
+    import pandas as pd
+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 GoogleSheetDumpTask(SheetDumpTaskBase):
+    _AUTO_REGISTER = True
+    ds_name_fields = ("google_service_account",)
+    worker_install_require = [
+        "gspread",
+    ]
+
+    custom_config_schema_properties = {
+        "google_service_account": {
+            "type": "string",
+            "title": _l("Service Account"),
+            "description": _l("Google service account with permissions to access the spreadsheet"),
+            "ui:field": "ProjectConnectionSelectorField",
+            "ui:options": {
+                "supportTypes": [
+                    "google_service_account",
+                ],
+            },
+        },
+        "file_url": {
+            "type": "string",
+            "title": _l("Spreadsheet URL"),
+            "description": _l("URL of the Google spreadsheet (defaults to first sheet if no sheet ID specified)"),
+            "ui:field": "CodeEditorWithReferencesField",
+            "ui:options": {
+                "type": "plain",
+            },
+        },
+        "cell_range": {
+            "type": "string",
+            "title": _l("Data Range"),
+            "description": _l("Cell range in A1 notation (e.g. A1:B10). Reads entire sheet if empty"),
+            "ui:field": "CodeEditorWithReferencesField",
+            "ui:options": {
+                "type": "plain",
+            },
+        },
+    }
+    custom_config_schema_required = ["google_service_account", "file_url"]
+
+    def read_origin_df(self) -> "pd.DataFrame":
+        conf = self.rendered_config
+
+        ds = self.must_get_connection_by_name(conf.google_service_account)
+        service_account = ds.recurve_connector
+        spread_sheet_id, sheet_id = service_account.parse_sheet_url(conf.file_url)
+        logger.info(f"reading {conf.file_url}, gid {sheet_id}")
+
+        sheet = service_account.get_sheet(conf.file_url, sheet_id)
+        df = service_account.read_sheet_to_df(sheet, cell_range=conf.cell_range)
+        logger.info(f"original DataFrame shape {df.shape}, dtypes:\n{df.dtypes}")
+        logger.info(df.head())
+        return df

recurvedata/operators/transfer_operator/dump_task_mongodb.py

@@ -0,0 +1,168 @@
+import copy
+import json
+import logging
+
+try:
+    from bson import json_util
+
+    from recurvedata.pigeon.dumper.mongodb import MongoDBDumper
+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 date_time, extract_dict
+
+logger = logging.getLogger(__name__)
+
+
+class MongoDBDumpTask(DumpTask):
+    ds_name_fields = ("data_source_name",)
+    worker_install_require = ["pigeon[mongo]"]
+
+    @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()
+        # 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 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=["collection", "filter", "projection"])
+        dump_options.update({"connector": ds.connector, "handler_factories": [hf], "database": ds.database})
+
+        # projection 设置为 null 或 '' 都当作 None 处理，即包含所有字段
+        proj = dump_options.get("projection")
+        if proj:
+            dump_options["projection"] = json.loads(proj)
+        else:
+            dump_options["projection"] = None
+
+        if dump_options["filter"]:
+            flt = json_util.loads(dump_options["filter"])
+        else:
+            flt = {}
+        if not self.dag.is_once and self.config.incremental_by_time:
+            start, end = self.determine_time_range()
+            time_flt = {self.config.time_column: {"$gte": start, "$lt": end}}
+            flt.update(time_flt)
+
+        dump_options["filter"] = flt
+
+        logger.info("Dump options: %s", dump_options)
+        dumper = MongoDBDumper(**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('mongodb')
+        return {
+            "type": "object",
+            "properties": {
+                "data_source_name": {
+                    "type": "string",
+                    "title": _l("MongoDB Connection"),
+                    "ui:field": "ProjectConnectionSelectorField",
+                    "ui:options": {
+                        "supportTypes": [
+                            "mongodb",
+                        ],
+                    },
+                },
+                "collection": {
+                    "type": "string",
+                    "title": _l("MongoDB Collection"),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "filter": {
+                    "type": "string",
+                    "title": _l("Query Filter"),
+                    "default": "{}",
+                    "description": _l(
+                        "MongoDB query filter in JSON format. Will be deserialized using bson.json_util and passed to find() method. "
+                        "Supports MongoDB query operators like $gt, $lt, $in etc. See MongoDB documentation for details."
+                    ),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "code",
+                        "lang": "json",
+                    },
+                },
+                "projection": {
+                    "type": "string",
+                    "title": _l("Field Selection"),
+                    "description": _l(
+                        "Specify which fields to return in JSON format. Empty value returns all fields. Passed directly to MongoDB find() function."
+                    ),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "code",
+                        "lang": "json",
+                    },
+                },
+                "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_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", "collection"],
+        }