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

recurvedata/operators/transfer_operator/dump_task_oss.py

@@ -0,0 +1,285 @@
+import csv
+import inspect
+import logging
+import os
+import shutil
+from copy import deepcopy
+
+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.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 AliyunOSSDumpTask(DumpTask):
+    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"])
+
+        object_path = conf["object_path"]
+        delimiter = "/" if object_path.endswith("/") else ""
+        keys = self.get_keys(ds, prefix=object_path, delimiter=delimiter)
+        logger.info(f"[start] downloading keys:{keys} from oss")
+        local_files = []
+        for key in keys:
+            ds.connector.download(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)
+
+        shutil.rmtree(tmp_dirname)
+        return None
+
+    @staticmethod
+    def get_keys(ds, prefix, delimiter=""):
+        keys = ds.connector.get_keys(prefix=prefix, delimiter=delimiter)
+        return keys
+
+    @staticmethod
+    def _infer_compression(filename: str, default_compression: str) -> str:
+        """infer compression method from filename"""
+        ext = os.path.splitext(filename)[1].lower()
+        compression_map = {".gz": "Gzip", ".zip": "Zip"}
+        return compression_map.get(ext, default_compression)
+
+    @staticmethod
+    def _infer_file_format(filename: str, default_format: str) -> str:
+        """infer file format from filename"""
+        ext = os.path.splitext(filename)[1].lower()
+        format_map = {
+            ".xlsx": "Excel",
+            ".xls": "Excel",
+        }
+        return format_map.get(ext, default_format)
+
+    def process_file(self, conf, files):
+        filename = self.filename
+
+        for f in files:
+            compression = self._infer_compression(f, conf["decompress"])
+
+            file_format = self._infer_file_format(f, conf["file_format"])
+
+            if compression == "Gzip":
+                logger.info("decompressing %s using gzip", f)
+                ff.gzip_decompress(f, inplace=True)
+            elif compression == "Zip":
+                logger.info("decompressing %s using zip", f)
+                ff.zip_decompress(f, inplace=True)
+
+            skip_head_lines = conf.get("skip_head_lines", 0)
+
+            if file_format == "Excel":
+                logger.info("converting Excel to CSV...")
+                ff.convert_excel_to_csv(f, skiprows=skip_head_lines, inplace=True)
+            elif 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)
+            elif file_format == "CSV":
+                logger.info("converting CSV dialect and encoding if necessary...")
+                dialect_options = self._get_custom_csv_options(conf)
+                src_dialect_options = deepcopy(dialect_options)
+                src_dialect_options.pop("quoting")
+                src_dialect_options.pop("doublequote")
+                src_dialect_options.pop("escapechar")
+                ff.convert_csv_dialect(
+                    f,
+                    src_dialect_options=src_dialect_options,
+                    dst_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 = self._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"]),
+            "quotechar": '"',
+            "doublequote": False,
+            "escapechar": "'",
+        }
+        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("Aliyun OSS Connection"),
+                    "ui:field": "ProjectConnectionSelectorField",
+                    "ui:options": {
+                        "supportTypes": [
+                            "oss",
+                        ],
+                    },
+                },
+                "object_path": {
+                    "type": "string",
+                    "title": _l("OSS Object Path"),
+                    "description": _l("Object path 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", "object_path"],
+        }
+
+    @staticmethod
+    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
+
+    @classmethod
+    def validate(cls, configuration):
+        conf = super().validate(configuration)
+
+        transform_func_code = conf.get("transform_func", "").strip()
+        if transform_func_code:
+            cls._validate_transform(transform_func_code)
+        return conf

recurvedata/operators/transfer_operator/dump_task_python.py

@@ -0,0 +1,212 @@
+import copy
+import inspect
+import logging
+import re
+
+import jsonschema
+
+from recurvedata.config import RECURVE_EXECUTOR_PYENV_NAME
+from recurvedata.core.translation import _l
+from recurvedata.operators.python_operator.operator import PythonRequirementsMixin
+from recurvedata.operators.transfer_operator.mixin import HiveTextfileConverterMixin
+from recurvedata.operators.transfer_operator.task import DumpTask
+
+logger = logging.getLogger(__name__)
+
+_SOURCE_SKELETON = _l(
+    '''
+def execute(filename: str, *args, **kwargs):
+    """
+    The execute function must be implemented as the entry point for ReOrc.
+
+    Args:
+        filename: Required. Output data to this file in CSV format. This file will be used as input for the Loader.
+
+    Data Source Parameters:
+        For database configurations, use ReOrc's Data Sources instead of hardcoding credentials in the code.
+        When defining the execute function, use special parameter names to specify required data sources.
+        ReOrc will pass corresponding pigeon connector objects that can be used for database interactions.
+
+        Parameter naming convention:
+        - Must have 'datasource_' prefix, e.g. datasource_xxx
+        - Example: datasource_mysql='my_mysql_default'
+          At runtime, ReOrc will pass a pigeon.connector.MysqlConnector object
+
+        Example usage:
+            def execute(filename, datasource_mysql='my_mysql_default'):
+                df = datasource_mysql.get_pandas_df('SELECT * FROM my_database.my_table')
+                df.to_csv(filename, header=False)
+    """
+    pass
+'''
+)
+
+
+# FIXME: record all supported template variables, find a way to keep consistent with `get_template_context` method
+_TEMPLATE_VARIABLES = {
+    "dt",
+    "yesterday",
+    "yesterday_dt",
+    "tomorrow",
+    "tomorrow_dt",
+    "logical_date",
+    "data_interval_start",
+    "data_interval_end",
+    "data_interval_start_dt",
+    "data_interval_end_dt",
+}
+
+
+class PythonCodeRunner(object):
+    # test page: https://regex101.com/r/p8YCQc/1
+    _JINJA2_VAR_PATTERN = re.compile(r"^{{\s*([^\d\W]\w*)\s*}}$")
+
+    def __init__(self, source):
+        self.source = source
+
+        self.__namespace = {}
+        self.__parameters = {}
+        self.__datasource_params = {}
+        self.__jinja2_variables_params = {}
+        self.__ready_for_execution = False
+        self.__compiled = False
+
+    @property
+    def entrypoint(self):
+        if not self.__compiled:
+            raise ValueError("entrypoint is not ready, inspect first")
+        return self.__namespace.get("execute")
+
+    def inspect(self):
+        logger.info("compiling source code\n%s", self.source)
+        code = compile(self.source, "", "exec")
+        exec(code, self.__namespace)
+        self.__compiled = True
+        entrypoint = self.entrypoint
+
+        if not (entrypoint and inspect.isfunction(entrypoint)):
+            raise jsonschema.ValidationError(message="execute function is required", path=("source",))
+
+        sig = inspect.signature(entrypoint)
+        for name, param in sig.parameters.items():
+            value = param.default
+            logger.info("found parameter %s=%s", name, value)
+
+            # special naming for data source parameters: `datasource_xxx`
+            if self.is_datasource_param(name):
+                if self._is_empty(value):
+                    raise jsonschema.ValidationError(message=f"{name} must be known data source name", path=("source",))
+                ds = DumpTask.get_connection_by_name(value)
+                if not ds:
+                    raise jsonschema.ValidationError(message=f"Unknown data source {repr(name)}", path=("source",))
+                self.__datasource_params[name] = value
+
+            # jinja2 template `{{ dt }}`, no Jinja2 rendering, directly replace
+            elif self.is_jinja2_variable(value):
+                variable = self._JINJA2_VAR_PATTERN.search(value).groups()[0]
+                # unsupported variables
+                if variable not in _TEMPLATE_VARIABLES:
+                    raise jsonschema.ValidationError(
+                        message=f"Unsupport template variable {repr(value)}", path=("source",)
+                    )
+                self.__jinja2_variables_params[name] = variable
+
+            else:
+                # keep default value, data source and template variable parameters are injected at runtime by calling `bind_parameters`
+                self.__parameters[name] = value
+        self.__parameters.update(self.__datasource_params)
+        self.__parameters.update(self.__jinja2_variables_params)
+
+    def is_datasource_param(self, name: str) -> bool:
+        return name.startswith("datasource_")
+
+    def is_jinja2_variable(self, name: str) -> bool:
+        return isinstance(name, str) and self._JINJA2_VAR_PATTERN.match(name)
+
+    @staticmethod
+    def _is_empty(obj) -> bool:
+        return obj is inspect.Signature.empty
+
+    def bind_parameters(self, filename, template_context, **kwargs):
+        params = copy.deepcopy(kwargs)
+        params["filename"] = filename
+
+        logger.info("binding data source connectors %s", self.__datasource_params)
+        for param_name, ds_name in self.__datasource_params.items():
+            params[param_name] = DumpTask.get_connection_by_name(ds_name).connector
+
+        logger.info("binding jinja2 variables %s", self.__jinja2_variables_params)
+        for param_name, variable in self.__jinja2_variables_params.items():
+            params[param_name] = template_context[variable]
+
+        # bind other parameters, or override default parameters
+        for k, v in params.items():
+            if k in self.__parameters:
+                self.__parameters[k] = v
+
+        # check if there are any parameters not passed
+        for name, value in self.__parameters.items():
+            if name not in ["args", "kwargs"] and self._is_empty(value):
+                raise TypeError(f"parameter {repr(name)} is not bound")
+
+        logger.info("bounded parameters %s", self.__parameters)
+        self.__ready_for_execution = True
+
+    def execute(self):
+        if not self.__ready_for_execution:
+            raise RuntimeError("must call inspect and bind_parameters before calling execute")
+        logger.info("calling entrypoint %s with parameters %s...", self.entrypoint, self.__parameters)
+        self.entrypoint(**self.__parameters)
+        logger.info("done.")
+
+
+class PythonDumpTask(DumpTask, HiveTextfileConverterMixin, PythonRequirementsMixin):
+    no_template_fields = ("source",)
+
+    def execute_impl(self, *args, **kwargs):
+        config = self.rendered_config.copy()
+        runner = PythonCodeRunner(config["source"])
+
+        # Get and install requirements if any
+        py_conn_configs = self.client.get_py_conn_configs()
+        if py_conn_configs and isinstance(py_conn_configs, dict):
+            requirements = "\n".join(py_conn_configs.get("requirements", []))
+            self._install_requirements(requirements, RECURVE_EXECUTOR_PYENV_NAME)
+
+        runner.inspect()
+        context = self.get_template_context()
+        runner.bind_parameters(filename=self.filename, template_context=context)
+        runner.execute()
+
+        self.convert_csv_to_hive_text_if_needed()
+        return None
+
+    @classmethod
+    def config_schema(cls):
+        return {
+            "type": "object",
+            "properties": {
+                "source": {
+                    "type": "string",
+                    "title": _l("Python Source Code"),
+                    "description": _l(
+                        "Python code that extracts data and writes to a CSV file. Must implement an execute() function that takes a filename parameter. Note: The Load step must specify a Create Table DDL when using PythonDump."
+                    ),
+                    "default": _SOURCE_SKELETON,
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "code",
+                        "lang": "python",
+                    },
+                },
+            },
+            "required": ["source"],
+        }
+
+    @classmethod
+    def validate(cls, configuration):
+        config = super().validate(configuration)
+
+        runner = PythonCodeRunner(config["source"])
+        runner.inspect()
+        return config