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

recurvedata/operators/transfer_operator/load_task_email.py

@@ -0,0 +1,188 @@
+import os
+
+try:
+    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.email_util import send_email
+
+_default_html = """<div><includetail><div style="font:Verdana normal 14px;color:#000;">
+<div style="position:relative;"><blockquote style="margin-Top: 0px; margin-Bottom: 0px; margin-Left: 0.5em">
+<div class="FoxDiv20190108121908737768">
+<div id="mailContentContainer" style=" font-size: 14px; padding: 0px; height: auto; min-height: auto ; ">
+<div>Dear all,</div>
+<div><br></div>
+
+<div>————————</div>
+<div>RecurveData Automatic Reporting</div>
+</div>
+</div></blockquote>
+</div></div>"""
+
+
+class EmailLoadTask(LoadTask):
+    ds_name_fields = ("data_source_name",)
+    ds_types = ("mail",)
+    worker_install_require = ["pigeon"]
+
+    def execute_impl(self, *args, **kwargs):
+        load_options = self.rendered_config
+        ds = self.must_get_connection_by_name(load_options["data_source_name"])
+        smtp_config = {
+            "host": ds.host,
+            "port": ds.port,
+            "ssl": ds.extra.get("ssl", True),
+            "user": ds.user,
+            "password": ds.password,
+            "timeout": ds.extra.get("timeout", 60),
+        }
+
+        remove_files = [self.filename]
+        filename = load_options.get("filename")
+        default_file_ext = ".csv"
+        if filename and "." not in filename:
+            filename = f"{filename}{default_file_ext}"
+        if filename and self.filename:
+            # 文件压缩
+            uncompress_filename = filename
+            if filename.endswith((".zip", ".gz")):
+                uncompress_filename = ".".join(filename.split(".")[:-1])
+            new_filename = os.path.join(os.path.dirname(self.filename), uncompress_filename)
+            os.rename(self.filename, new_filename)
+            compress_mode = load_options["compress_mode"]
+            file_upload, ext = self.compress_file(filename=new_filename, compress_mode=compress_mode)
+            if compress_mode != "None" and not load_options["filename"].endswith(ext):
+                filename = f"{filename}{ext}"
+
+            files = {filename: file_upload}
+            remove_files = [new_filename, file_upload]
+        else:
+            files = None
+
+        ok = send_email(
+            mail_to=self.parse_email_list(load_options["mail_to"]),
+            subject=load_options["subject"],
+            html=load_options["html"],
+            cc=self.parse_email_list(load_options.get("cc")),
+            bcc=self.parse_email_list(load_options.get("bcc")),
+            files=files,
+            mail_from=load_options["mail_from"],
+            reply_to=load_options.get("reply_to"),
+            smtp_config=smtp_config,
+        )
+        assert ok, "Failed to send email"
+        fs.remove_files_safely(remove_files)
+
+    @staticmethod
+    def parse_email_list(obj, separator=";"):
+        if not obj:
+            return None
+        return obj.split(separator)
+
+    @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("SMTP Server"),
+                    "ui:field": "ProjectConnectionSelectorField",
+                    "ui:options": {
+                        "supportTypes": cls.ds_types,
+                    },
+                },
+                "subject": {
+                    "type": "string",
+                    "title": _l("Email Subject"),
+                    "description": _l("Subject line of the email"),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "mail_to": {
+                    "type": "string",
+                    "title": _l("Recipients"),
+                    "description": _l("Email recipients (separate multiple addresses with semicolons)"),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "mail_from": {
+                    "type": "string",
+                    "title": _l("Sender Name"),
+                    "default": "RecurveData SERVICE",
+                    "description": _l("Display name that appears as the email sender"),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "cc": {
+                    "type": "string",
+                    "title": _l("CC Recipients"),
+                    "description": _l("Carbon copy recipients (separate multiple addresses with semicolons)"),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "bcc": {
+                    "type": "string",
+                    "title": _l("BCC Recipients"),
+                    "description": _l("Blind carbon copy recipients (separate multiple addresses with semicolons)"),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "html": {
+                    "type": "string",
+                    "title": _l("Email Body"),
+                    "description": _l("HTML content of the email body."),
+                    "default": _default_html,
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "code",
+                        "lang": "python",
+                    },
+                },
+                "filename": {
+                    "type": "string",
+                    "title": _l("Attachment Name"),
+                    "description": _l(
+                        "Name of the email attachment. Supports template variables. Leave empty for no attachment. Default extension is .csv if none specified."
+                    ),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "compress_mode": {
+                    "type": "string",
+                    "title": _l("Compression Method"),
+                    "enum": ["None", "Gzip", "Zip"],
+                    "enumNames": ["None", "Gzip", "Zip"],
+                    "default": "None",
+                    "description": _l("Compression method for attachments"),
+                },
+                "reply_to": {
+                    "type": "string",
+                    "title": _l("Reply-To Address"),
+                    "description": _l("Email address that will receive replies to this email"),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+            },
+            "required": ["data_source_name", "subject", "mail_to", "mail_from", "html"],
+        }
+        return schema

recurvedata/operators/transfer_operator/load_task_es.py

@@ -0,0 +1,86 @@
+try:
+    from recurvedata.pigeon.loader.csv_to_es import CSVToElasticSearchLoader
+except ImportError:
+    pass
+
+from recurvedata.core.translation import _l
+from recurvedata.operators.transfer_operator.task import LoadTask
+from recurvedata.utils import extract_dict
+
+
+class ElasticSearchLoadTask(LoadTask):
+    ds_name_fields = ("data_source_name",)
+    ds_types = ("elasticsearch",)
+    should_write_header = True
+    worker_install_require = ["pigeon[elasticsearch]"]
+
+    def execute_impl(self, *args, **kwargs):
+        es_ds = self.must_get_connection_by_name(self.config["data_source_name"])
+        load_options = extract_dict(self.rendered_config, keys=["index", "doc_type", "id_field", "generate_id"])
+        load_options.update(
+            {
+                "connector": es_ds.connector,
+                "filename": self.filename,
+                "delete_file": True,
+            }
+        )
+        loader = CSVToElasticSearchLoader(**load_options)
+        return loader.execute()
+
+    @classmethod
+    def config_schema(cls):
+        # get_choices_by_type = cls.get_connection_names_by_type
+        # dss = get_choices_by_type(cls.ds_types)
+        return {
+            "type": "object",
+            "properties": {
+                "data_source_name": {
+                    "type": "string",
+                    "title": _l("Elasticsearch Data Source"),
+                    "ui:field": "ProjectConnectionSelectorField",
+                    "ui:options": {
+                        "supportTypes": cls.ds_types,
+                    },
+                    # 'default': cls.first_or_default(dss, ''),
+                },
+                "index": {
+                    "type": "string",
+                    "title": _l("Elasticsearch Index"),
+                    "description": _l("Name of the Elasticsearch index to load data into"),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "doc_type": {
+                    "type": "string",
+                    "title": _l("Document Type"),
+                    "description": _l("Type of document to create in Elasticsearch"),
+                    "default": "_doc",
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "id_field": {
+                    "type": "string",
+                    "title": _l("Document ID Field"),
+                    "description": _l(
+                        "Field from the input data to use as the document ID. Leave empty to auto-generate IDs"
+                    ),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "generate_id": {
+                    "type": "boolean",
+                    "title": _l("Generate Document IDs"),
+                    "description": _l(
+                        "Automatically generate unique document IDs based on record content. Takes precedence over ID Field if both are specified"
+                    ),
+                    "default": False,
+                },
+            },
+            "required": ["data_source_name", "index"],
+        }

recurvedata/operators/transfer_operator/load_task_filebrowser.py

@@ -0,0 +1,151 @@
+import datetime
+import logging
+import mimetypes
+import os
+import urllib.parse
+
+import requests
+
+try:
+    from recurvedata.pigeon.utils import fs, trim_suffix
+except ImportError:
+    pass
+
+from recurvedata.operators.transfer_operator.task import LoadTask
+
+logger = logging.getLogger(__name__)
+
+
+class FileBrowserLoadTask(LoadTask):
+    ds_name_fields = ("data_source_name",)
+    ds_types = ("filebrowser",)
+    should_write_header = True
+    worker_install_require = ["requests", "pigeon"]
+
+    def execute_impl(self, *args, **kwargs):
+        conf = self.rendered_config
+        ds = self.must_get_connection_by_name(conf["data_source_name"])
+        file_to_upload = self.filename
+        remote_filename = conf["filename"]
+
+        cm = conf["compress_method"]
+        if cm != "None":
+            logger.info("compressing file using %s...", cm)
+            if cm == "Gzip":
+                ext = ".gz"
+                file_to_upload = fs.gzip_compress(self.filename, using_cmd=True)
+            elif cm == "Zip":
+                ext = ".zip"
+                arcname = trim_suffix(os.path.basename(remote_filename), ext)
+                file_to_upload = fs.zip_compress(self.filename, using_cmd=False, arcname=arcname)
+            else:
+                # won't reach here
+                raise ValueError(f"compress method {cm} is not supported")
+
+            if not remote_filename.endswith(ext):
+                remote_filename = f"{remote_filename}{ext}"
+
+        client = FileBrowserClient(ds.host, ds.user, ds.password)
+
+        remote_filename = os.path.join(conf["directory"], remote_filename)
+        logger.info("uploading %s to %s", file_to_upload, remote_filename)
+        client.upload(file_to_upload, remote_filename, override=True)
+        fs.remove_files_safely([self.filename, file_to_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": "FileBrowser Data Source",
+                    "ui:field": "ProjectConnectionSelectorField",
+                    "ui:options": {
+                        "supportTypes": cls.ds_types,
+                    },
+                    # 'default': cls.first_or_default(dss, ''),
+                },
+                "directory": {
+                    "type": "string",
+                    "title": "Directory",
+                    "description": "要上传到的文件夹",
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "filename": {
+                    "type": "string",
+                    "title": "Filename",
+                    "description": "上传后的文件名，支持模板变量",
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "compress_method": {
+                    "type": "string",
+                    "title": "Compress Method",
+                    "enum": ["None", "Gzip", "Zip"],
+                    "enumNames": ["None", "Gzip", "Zip"],
+                    "default": "None",
+                    "description": "文件的压缩方式，默认不压缩。如果选择了压缩，会在文件名加上相应的后缀。",
+                },
+            },
+            "required": ["data_source_name", "directory", "filename", "compress_method"],
+        }
+        return schema
+
+
+class FileBrowserClient(object):
+    # token 有效期，实际上默认是 2 小时，这里只保留 1 小时
+    TOKEN_AGE = datetime.timedelta(seconds=1 * 60 * 60)
+
+    def __init__(self, host, username, password):
+        self.host = host
+        self.username = username
+        self.password = password
+        self._session = requests.Session()
+        self._token = None
+        self._token_expires_at = datetime.datetime.fromtimestamp(0)
+
+    def _request(self, method, url, params=None, data=None, json=None, auth=True, **kwargs):
+        full_url = urllib.parse.urljoin(self.host, url)
+        if auth:
+            headers = {"X-Auth": self.token}
+        else:
+            headers = {}
+        headers.update(kwargs.pop("headers", {}))
+        resp = self._session.request(method, full_url, params=params, data=data, json=json, headers=headers, **kwargs)
+        resp.raise_for_status()
+        return resp
+
+    @property
+    def token(self):
+        if self._token is None or self._token_expires_at <= datetime.datetime.now():
+            self._token = self.login()
+            self._token_expires_at = datetime.datetime.now() + self.TOKEN_AGE
+        return self._token
+
+    def login(self):
+        params = {
+            "username": self.username,
+            "password": self.password,
+            "recaptcha": "",
+        }
+        resp = self._request("POST", "/api/login", json=params, auth=False)
+        token = resp.text
+        return token
+
+    def upload(self, local_filename, remote_filename, override=True):
+        headers = {}
+        content_type, _ = mimetypes.guess_type(local_filename)
+        if content_type:
+            headers = {"Content-Type": content_type}
+        params = {"override": override and "true" or "false"}
+        url = f"/api/resources/{urllib.parse.quote(remote_filename)}"
+        with open(local_filename, "rb") as f:
+            self._request("POST", url, params=params, data=f, headers=headers)

recurvedata/operators/transfer_operator/load_task_ftp.py

@@ -0,0 +1,19 @@
+import logging
+
+from recurvedata.operators.transfer_operator.load_task_sftp import SFTPLoadTask
+
+logger = logging.getLogger(__name__)
+
+
+class FTPLoadTask(SFTPLoadTask):
+    ds_name_fields = ("data_source_name",)
+    ds_types = ("ftp",)
+
+    @staticmethod
+    def ensure_directory_exists(ds, conf):
+        ftp = ds.connector
+        try:
+            ftp.list_dir(conf["directory"])
+        except OSError:
+            logger.warning("failed to list directory %s, maybe not exists, try to make it", conf["directory"])
+            ftp.makedir(conf["directory"])

recurvedata/operators/transfer_operator/load_task_google_bigquery.py

@@ -0,0 +1,90 @@
+import copy
+
+try:
+    from recurvedata.pigeon.loader.csv_to_google_bigquery import CSVToGoogleBigqueryLoader
+except ImportError:
+    pass
+
+from recurvedata.core.translation import _l
+from recurvedata.operators.transfer_operator import utils
+from recurvedata.operators.transfer_operator.task import LoadTask
+
+
+class GoogleBigqueryLoadTask(LoadTask):
+    ds_name_fields = ("google_bigquery_data_source_name",)
+    ds_types = ("bigquery",)
+    default_dumper_handler_options = {}
+    worker_install_require = ["pigeon[google_bigquery]"]
+
+    def execute_impl(self, *args, **kwargs):
+        google_bigquery_ds = self.must_get_connection_by_name(self.config["google_bigquery_data_source_name"])
+        load_options = self.rendered_config.copy()
+        for k in ["google_bigquery_data_source_name"]:
+            load_options.pop(k, None)
+        load_options.update(
+            {
+                "filename": self.filename,
+                "google_bigquery_connector": google_bigquery_ds.connector,
+                "delete_file": True,
+                "dataset": google_bigquery_ds.data.get("database"),
+            }
+        )
+        loader = CSVToGoogleBigqueryLoader(**load_options)
+        return loader.execute()
+
+    @classmethod
+    def config_schema(cls):
+        # get_choices_by_type = cls.get_connection_names_by_type
+        # dws = get_choices_by_type(cls.ds_types)
+        schema = {
+            "type": "object",
+            "properties": {
+                "google_bigquery_data_source_name": {
+                    "type": "string",
+                    "title": _l("BigQuery Connection"),
+                    "description": _l("Select the BigQuery connection to use"),
+                    "ui:field": "ProjectConnectionSelectorField",
+                    "ui:options": {
+                        "supportTypes": cls.ds_types,
+                    },
+                    # 'default': cls.first_or_default(dws, ''),
+                },
+                # "dataset": {
+                #     "type": "string",
+                #     "title": "Dataset",
+                #     "ui:field": "CodeEditorWithReferencesField",
+                #     "ui:options": {
+                #         "type": "plain",
+                #     },
+                # },
+                "table": {
+                    "type": "string",
+                    "title": _l("Target Table"),
+                    "description": _l("Name of the table to load data into"),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "create_table_ddl": {
+                    "type": "string",
+                    "title": _l("Table Creation SQL"),
+                    "description": _l("SQL statement to create the destination table if it doesn't exist"),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "code",
+                        "lang": "sql",
+                        "sqlLang": "sql",
+                    },
+                },
+            },
+            "required": ["google_bigquery_data_source_name", "table", "mode"],
+        }
+        properties_schema = schema["properties"]
+        properties_schema.update(copy.deepcopy(utils.LOAD_COMMON))
+
+        # remove dedup
+        properties_schema.pop("dedup", None)
+        properties_schema.pop("dedup_uniq_keys", None)
+        properties_schema.pop("dedup_orderby", None)
+        return schema

recurvedata/operators/transfer_operator/load_task_google_cloud_storage.py

@@ -0,0 +1,127 @@
+import logging
+import os
+
+try:
+    from recurvedata.pigeon.connector.google_cloud_storage import GoogleCloudStorageConnector
+except ImportError:
+    pass
+
+from recurvedata.core.translation import _l
+from recurvedata.operators.transfer_operator.task import LoadTask
+from recurvedata.utils import extract_dict
+from recurvedata.utils.files import is_file_empty, remove_files_safely
+
+logger = logging.getLogger(__name__)
+
+
+class GoogleCloudStorageLoadTask(LoadTask):
+    ds_name_fields = ("data_source_name",)
+    ds_types = ("google_cloud_storage",)
+    should_write_header = True
+    worker_install_require = ["pigeon[google_bigquery]"]
+
+    def execute_impl(self, *args, **kwargs):
+        if 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"])
+
+        load_options = self.rendered_config.copy()
+        ds_extra_bucket = ds.extra.get("bucket")
+        config_bucket = load_options.get("bucket_name")
+        bucket_upload = config_bucket if config_bucket else ds_extra_bucket
+
+        opt_keys = ["key_path", "key_dict", "project", "proxies", "bucket_name"]
+        gcs = GoogleCloudStorageConnector(**extract_dict(ds.extra, opt_keys))
+
+        # 文件压缩
+        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 = {
+            "bucket_name": bucket_upload,
+            "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)})
+
+        if ext == ".zip":
+            upload_conf["filename"] += ext
+
+        logger.info("uploading...")
+        logger.info(upload_conf)
+        gcs.upload(**upload_conf)
+        return 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("GCS Connection"),
+                    "ui:field": "ProjectConnectionSelectorField",
+                    "ui:options": {
+                        "supportTypes": cls.ds_types,
+                    },
+                    # 'default': cls.first_or_default(dss, ''),
+                },
+                "bucket_name": {
+                    "type": "string",
+                    "title": _l("GCS Bucket Name"),
+                    "description": _l(
+                        "The name of the Google Cloud Storage bucket to upload files to. This is required if not already configured in the data source connection."
+                    ),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "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", "key"],
+        }
+        return schema