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

recurvedata/operators/transfer_operator/dump_task_s3.py

@@ -0,0 +1,270 @@
+import csv
+import inspect
+import logging
+import os
+import shutil
+
+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
+from recurvedata.utils.files import merge_files
+
+logger = logging.getLogger(__name__)
+
+
+class S3DumpTask(DumpTask, HiveTextfileConverterMixin):
+    ds_name_fields = ("data_source_name",)
+
+    def execute_impl(self, *args, **kwargs):
+        tmp_dirname = f"{self.filename}_dir"
+        if os.path.exists(tmp_dirname):
+            shutil.rmtree(tmp_dirname)
+        os.makedirs(tmp_dirname)
+
+        conf = self.rendered_config.copy()
+
+        ds = self.must_get_connection_by_name(conf["data_source_name"])
+        ds_extra_bucket = ds.extra.get("bucket")
+        config_bucket = conf.get("bucket_name")
+        bucket = config_bucket if config_bucket else ds_extra_bucket
+
+        conf_keys = conf["keys"]
+        keys = self.get_keys(ds=ds, bucket=bucket, prefix=conf_keys)
+        logger.info(f"[start] downloading keys:{keys} from s3")
+        local_files = []
+        for key in keys:
+            ds.connector.download(bucket, key, folder=tmp_dirname)
+            local_files.append(os.path.join(tmp_dirname, os.path.basename(key)))
+        logger.info(f"[finish] downloading files to {local_files}")
+
+        filename = self.process_file(conf=conf, files=local_files)
+        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()
+
+        shutil.rmtree(tmp_dirname)
+        return None
+
+    @staticmethod
+    def get_keys(ds, bucket, prefix):
+        keys = ds.connector.get_keys(bucket_name=bucket, prefix=prefix)
+        return keys
+
+    def process_file(self, conf, files):
+        filename = self.filename
+
+        for f in files:
+            if conf["decompress"] == "Gzip":
+                logger.info("decompressing %s using gzip", f)
+                ff.gzip_decompress(f, inplace=True)
+            if conf["decompress"] == "Zip":
+                logger.info("decompressing %s using gzip", f)
+                ff.zip_decompress(f, inplace=True)
+
+            skip_head_lines = conf.get("skip_head_lines", 0)
+            if conf["file_format"] == "Excel":
+                logger.info("converting Excel to CSV...")
+                ff.convert_excel_to_csv(f, skiprows=skip_head_lines, inplace=True)
+            if conf["file_format"] == "JSONLines":
+                logger.info("converting JSON lines to CSV...")
+                ff.convert_jsonlines_to_csv(f, 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)
+                ff.convert_csv_dialect(
+                    f,
+                    src_dialect_options=dialect_options,
+                    skiprows=skip_head_lines,
+                    src_encoding=conf["encoding"],
+                    inplace=True,
+                )
+        if files:
+            merge_files(files=files, filename=filename)
+        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
+
+    @staticmethod
+    def _get_custom_csv_options(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("S3 Connection"),
+                    "ui:field": "ProjectConnectionSelectorField",
+                    "ui:options": {
+                        "supportTypes": [
+                            "s3",
+                        ],
+                    },
+                },
+                "bucket_name": {
+                    "type": "string",
+                    "title": _l("S3 Bucket Name"),
+                    "description": _l(
+                        "Name of the S3 bucket to download from. Required if not configured in data source."
+                    ),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "keys": {
+                    "type": "string",
+                    "title": _l("S3 Object Keys"),
+                    "description": _l("Object key or prefix pattern to download. Supports Jinja templating syntax."),
+                    "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", "Zip"],
+                    "enumNames": ["None", "Gzip", "Zip"],
+                    "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", "keys"],
+        }
+
+    @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_sftp.py

@@ -0,0 +1,229 @@
+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 SFTPDumpTask(DumpTask, HiveTextfileConverterMixin):
+    ds_name_fields = ("data_source_name",)
+
+    def execute_impl(self, *args, **kwargs):
+        conf = self.rendered_config
+
+        ds = self.must_get_connection_by_name(conf["data_source_name"])
+        ds.connector.download_file(conf["filepath"], self.filename)
+
+        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 None
+
+    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):
+        return {
+            "type": "object",
+            "properties": {
+                "data_source_name": {
+                    "type": "string",
+                    "title": _l("SFTP Connection"),
+                    "ui:field": "ProjectConnectionSelectorField",
+                    "ui:options": {
+                        "supportTypes": [
+                            "sftp",
+                        ],
+                    },
+                },
+                "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/load_task_aliyun_oss.py

@@ -0,0 +1,107 @@
+import logging
+import os
+
+try:
+    from recurvedata.pigeon.connector.aliyun_oss import OSSBucketConnector
+    from recurvedata.pigeon.utils import fs
+except ImportError:
+    pass
+
+from recurvedata.core.translation import _l
+from recurvedata.operators.transfer_operator.task import LoadTask
+from recurvedata.utils import extract_dict
+
+logger = logging.getLogger(__name__)
+
+
+class AliyunOSSLoadTask(LoadTask):
+    ds_name_fields = ("data_source_name",)
+    ds_types = ("oss",)
+    should_write_header = True
+    worker_install_require = ["pigeon[aliyun]"]
+
+    def execute_impl(self, *args, **kwargs):
+        if fs.is_file_empty(self.filename):
+            logger.warning("file %s not exists or has no content, skip.", self.filename)
+            return
+
+        ds = self.must_get_connection_by_name(self.config["data_source_name"])
+
+        opt_keys = ["access_key_id", "access_key_secret", "endpoint", "bucket_name", "proxies"]
+        oss = OSSBucketConnector(**extract_dict(ds.extra, opt_keys))
+
+        load_options = self.rendered_config.copy()
+
+        # 文件压缩
+        compress_mode = load_options["compress_mode"]
+        if compress_mode != "None" and not load_options["key"].endswith(("/",)):
+            target_filename = os.path.join(os.path.dirname(self.filename), os.path.basename(load_options["key"]))
+        else:
+            target_filename = None
+        file_upload, ext = self.compress_file(
+            filename=self.filename, target_filename=target_filename, compress_mode=compress_mode
+        )
+        if compress_mode != "None" and not load_options["key"].endswith(("/", ext)):
+            load_options["key"] = f"{load_options['key']}{ext}"
+
+        # 根据 key 的内容创建 upload 方法需要的 key, folder 参数
+        upload_conf = {
+            "filename": file_upload,
+            "overwrite": load_options["overwrite"],
+        }
+        if load_options["key"].endswith("/"):
+            upload_conf.update({"folder": load_options["key"]})
+        elif load_options["key"]:
+            upload_conf.update({"key": load_options["key"]})
+        else:
+            upload_conf.update({"key": os.path.basename(file_upload)})
+
+        logger.info("uploading...")
+        oss.upload(**upload_conf)
+        return fs.remove_files_safely([self.filename, file_upload])
+
+    @classmethod
+    def config_schema(cls):
+        # get_choices_by_type = cls.get_connection_names_by_type
+        # dss = get_choices_by_type(cls.ds_types)
+        schema = {
+            "type": "object",
+            "properties": {
+                "data_source_name": {
+                    "type": "string",
+                    "title": _l("OSS Connection"),
+                    "ui:field": "ProjectConnectionSelectorField",
+                    "ui:options": {
+                        "supportTypes": cls.ds_types,
+                    },
+                },
+                "key": {
+                    "type": "string",
+                    "title": _l("Upload Path"),
+                    "description": _l(
+                        "Target path in the bucket. Can be an object key or folder path (ending with /). "
+                        "Supports Jinja templating."
+                    ),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "compress_mode": {
+                    "type": "string",
+                    "title": _l("Compression Method"),
+                    "description": _l("Compress file before uploading using specified method"),
+                    "enum": ["None", "Gzip", "Zip"],
+                    "enumNames": ["None", "Gzip", "Zip"],
+                    "default": "None",
+                },
+                "overwrite": {
+                    "type": "boolean",
+                    "title": _l("Overwrite Existing"),
+                    "description": _l("Whether to overwrite if target object already exists"),
+                    "default": True,
+                },
+            },
+            "required": ["compress_mode", "data_source_name"],
+        }
+        return schema