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

recurvedata/operators/transfer_operator/load_task_redshift.py

@@ -0,0 +1,95 @@
+import copy
+
+try:
+    from recurvedata.pigeon.loader.csv_to_redshift import CSVToRedshiftLoader
+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 RedshiftLoadTask(LoadTask):
+    ds_name_fields = ("redshift_data_source_name",)
+    ds_types = ("redshift",)
+    worker_install_require = ["pigeon[redshift]"]
+
+    def execute_impl(self, *args, **kwargs):
+        redshift_ds = self.must_get_connection_by_name(self.config["redshift_data_source_name"])
+
+        load_options = self.rendered_config.copy()
+        for k in ["redshift_data_source_name"]:
+            load_options.pop(k, None)
+
+        load_options.update(
+            {
+                "filename": self.filename,
+                "redshift_connector": redshift_ds.connector,
+                "delete_file": True,
+                "database": redshift_ds.database,
+                "schema": redshift_ds.data.get("schema"),
+            }
+        )
+        loader = CSVToRedshiftLoader(**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)
+        schema = {
+            "type": "object",
+            "properties": {
+                "redshift_data_source_name": {
+                    "type": "string",
+                    "title": _l("Redshift Connection"),
+                    "ui:field": "ProjectConnectionSelectorField",
+                    "ui:options": {
+                        "supportTypes": cls.ds_types,
+                    },
+                    # 'default': cls.first_or_default(dss, ''),
+                },
+                # "database": {
+                #     "type": "string",
+                #     "title": "Database",
+                #     "ui:field": "CodeEditorWithReferencesField",
+                #     "ui:options": {
+                #         "type": "plain",
+                #     },
+                # },
+                # "schema": {
+                #     "type": "string",
+                #     "title": "Schema",
+                #     "default": "public",
+                #     "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 target table if it doesn't exist"),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "code",
+                        "lang": "sql",
+                        "sqlLang": "redshift",
+                    },
+                },
+            },
+            "required": ["redshift_data_source_name", "table", "mode"],
+        }
+        properties_schema = schema["properties"]
+        properties_schema.update(copy.deepcopy(utils.LOAD_COMMON))
+        return schema

recurvedata/operators/transfer_operator/load_task_s3.py

@@ -0,0 +1,150 @@
+import json
+import logging
+import os
+
+import jsonschema
+
+try:
+    from recurvedata.pigeon.connector import new_s3_connector
+    from recurvedata.pigeon.utils import fs
+except ImportError:
+    pass
+
+from recurvedata.core.translation import _l
+from recurvedata.operators.transfer_operator.task import LoadTask
+
+logger = logging.getLogger(__name__)
+
+
+class S3LoadTask(LoadTask):
+    ds_name_fields = ("data_source_name",)
+    ds_types = ("s3",)
+    should_write_header = True
+    worker_install_require = ["pigeon"]
+
+    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
+
+        s3_ds = self.must_get_connection_by_name(self.config["data_source_name"])
+
+        load_options = self.rendered_config.copy()
+        ds_extra_bucket = s3_ds.extra.get("bucket")
+        config_bucket = load_options.get("bucket_name")
+        bucket_upload = config_bucket if config_bucket else ds_extra_bucket
+
+        # 文件压缩
+        compress_mode = load_options["compress_mode"]
+        file_upload, ext = self.compress_file(filename=self.filename, compress_mode=compress_mode)
+        if compress_mode != "None" and not load_options["key"].endswith(("/", ext)):
+            load_options["key"] = f"{load_options['key']}{ext}"
+
+        # s3 connector 配置项
+        s3_conf = s3_ds.extra.copy()
+        if load_options.get("proxies"):
+            s3_conf["proxies"] = json.loads(load_options["proxies"])
+
+        # 创建 connector，如果不选择自动创建 Bucket，不存在则报错
+        s3 = new_s3_connector(conf=s3_conf)
+        if not load_options.get("auto_create_bucket") and config_bucket:
+            if not s3.has_bucket(bucket_name=config_bucket):
+                raise ValueError("Bucket not exists")
+
+        # 根据 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)})
+
+        logger.info("uploading ...")
+        s3.upload(**upload_conf)
+        return fs.remove_files_safely([self.filename, file_upload])
+
+    @classmethod
+    def validate(cls, configuration):
+        config = super().validate(configuration)
+
+        if not config.get("bucket_name"):
+            s3_ds = cls.must_get_connection_by_name(configuration["data_source_name"])
+            if not s3_ds.extra.get("bucket"):
+                raise jsonschema.ValidationError(message="Unknown Bucket", path=("bucket_name",))
+        return config
+
+    @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("S3 Connection"),
+                    "ui:field": "ProjectConnectionSelectorField",
+                    "ui:options": {
+                        "supportTypes": cls.ds_types,
+                    },
+                    # 'default': cls.first_or_default(dss, ''),
+                },
+                "bucket_name": {
+                    "type": "string",
+                    "title": _l("Bucket"),
+                    "description": _l("S3 bucket name, required if not set in data source"),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                # "auto_create_bucket": {
+                #     "type": "boolean",
+                #     "title": "Auto Create Bucket",
+                #     "default": False,
+                #     "description": "如果 Bucket 不存在，是否需要按输入创建 Bucket，命名规则见 https://amzn.to/2HL8VDX",
+                # },
+                "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",
+                    },
+                },
+                "proxies": {
+                    "type": "string",
+                    "title": _l("Proxies"),
+                    "description": _l('HTTP/HTTPS proxy to use. Format: {"https": "http://example.com:3128"}'),
+                    "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

recurvedata/operators/transfer_operator/load_task_sftp.py

@@ -0,0 +1,90 @@
+import logging
+import os
+
+from recurvedata.core.translation import _l
+from recurvedata.operators.transfer_operator.task import LoadTask
+from recurvedata.utils.files import remove_files_safely
+
+logger = logging.getLogger(__name__)
+
+
+class SFTPLoadTask(LoadTask):
+    ds_name_fields = ("data_source_name",)
+    ds_types = ("sftp",)
+    should_write_header = True
+    worker_install_require = ["pigeon[sftp]"]
+
+    def execute_impl(self, *args, **kwargs):
+        conf = self.rendered_config
+        ds = self.must_get_connection_by_name(conf["data_source_name"])
+
+        cm = conf["compress_method"]
+        file_to_upload, ext = self.compress_file(filename=self.filename, compress_mode=cm)
+        remote_filename = conf["filename"]
+        if cm != "None" and not remote_filename.endswith(ext):
+            remote_filename = f"{remote_filename}{ext}"
+
+        self.ensure_directory_exists(ds, conf)
+        remote_filename = os.path.join(conf["directory"], remote_filename)
+
+        logger.info("uploading %s to %s", file_to_upload, remote_filename)
+        ds.connector.upload_file(file_to_upload, remote_filename)
+        remove_files_safely([self.filename, file_to_upload])
+
+    @staticmethod
+    def ensure_directory_exists(ds, conf):
+        sftp = ds.connector
+        # 确保目录存在
+        try:
+            sftp.sftp.listdir(conf["directory"])
+        except OSError:
+            logger.warning("failed to list directory %s, maybe not exists, try to make it", conf["directory"])
+            # 这一步可以抛出异常，需要人工介入
+            sftp.sftp.mkdir(conf["directory"])
+
+    @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("SFTP Connection"),
+                    "ui:field": "ProjectConnectionSelectorField",
+                    "ui:options": {
+                        "supportTypes": cls.ds_types,
+                    },
+                    # 'default': cls.first_or_default(dss, ''),
+                },
+                "directory": {
+                    "type": "string",
+                    "title": _l("Upload Path"),
+                    "description": _l("The directory to upload the file to, Jinja template supported"),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "filename": {
+                    "type": "string",
+                    "title": _l("Filename"),
+                    "description": _l("The filename to save as, Jinja template supported"),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "compress_method": {
+                    "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",
+                },
+            },
+            "required": ["data_source_name", "filename", "compress_method"],
+        }
+        return schema

recurvedata/operators/transfer_operator/load_task_starrocks.py

@@ -0,0 +1,169 @@
+try:
+    from recurvedata.pigeon.loader.csv_to_starrocks import CSVToStarRocksLoader
+except ImportError:
+    pass
+
+from recurvedata.core.translation import _l
+from recurvedata.operators.transfer_operator import const
+from recurvedata.operators.transfer_operator.task import LoadTask
+from recurvedata.operators.transfer_operator.utils import allowed_modes
+
+
+class StarRocksLoadTask(LoadTask):
+    ds_name_fields = ("data_source_name",)
+    ds_types = ("starrocks",)
+    default_dumper_handler_options = {
+        "null": r"\N",
+        "lineterminator": "\n",
+        "escapechar": "'",
+        "doublequote": False,
+    }
+    worker_install_require = [
+        "pigeon[starrocks]",
+    ]
+
+    def execute_impl(self, *args, **kwargs):
+        ds = self.must_get_connection_by_name(self.config["data_source_name"])
+        load_options = self.rendered_config.copy()
+        for k in ["data_source_name"]:
+            load_options.pop(k, None)
+        load_options.update(
+            {
+                "filename": self.filename,
+                "connector": ds.connector,
+                "delete_file": True,
+                "load_strict_mode": self.config.get("load_strict_mode", False),
+                "database": ds.database,
+            }
+        )
+        loader = CSVToStarRocksLoader(**load_options)
+        return loader.execute()
+
+    @classmethod
+    def config_schema(cls):
+        schema = {
+            "type": "object",
+            "properties": {
+                "data_source_name": {
+                    "type": "string",
+                    "title": _l("StarRocks Connection"),
+                    "ui:field": "ProjectConnectionSelectorField",
+                    "ui:options": {
+                        "supportTypes": cls.ds_types,
+                    },
+                },
+                # "database": {
+                #     "type": "string",
+                #     "title": "Database",
+                #     "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 target table if it doesn't exist. See "
+                        "<a target='_blank' href='https://docs.starrocks.io/docs/sql-reference/sql-statements/table_bucket_part_index/CREATE_TABLE'>"
+                        "StarRocks Docs</a> for details"
+                    ),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "code",
+                        "lang": "sql",
+                        "sqlLang": "mysql",
+                    },
+                },
+                "mode": {
+                    "type": "string",
+                    "title": _l("Load Mode"),
+                    "description": _l("How to handle existing data in the target table"),
+                    "enum": list(allowed_modes),
+                    "default": const.LOAD_OVERWRITE,
+                },
+                "primary_keys": {
+                    "ui:hidden": '{{parentFormData.mode !== "MERGE"}}',
+                    "type": "string",
+                    "title": _l("Primary Keys"),
+                    "description": _l(
+                        "Comma-separated list of columns used for deduplication in MERGE mode. "
+                        "Should be primary or unique key columns."
+                    ),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "using_insert": {
+                    "type": "boolean",
+                    "title": _l("Use INSERT Mode"),
+                    "description": _l("By default Stream Load is used. Enable to use INSERT statements instead."),
+                    "default": False,
+                    "ui:hidden": True,
+                },
+                "load_strict_mode": {
+                    "type": "boolean",
+                    "title": _l("Enable Strict Mode"),
+                    "default": False,
+                    "description": _l(
+                        "When enabled, validates that data matches target table schema before loading. "
+                        "Raises error if validation fails."
+                    ),
+                    "ui:hidden": "{{parentFormData.using_insert}}",
+                },
+                "insert_batch_size": {
+                    "ui:hidden": "{{!parentFormData.using_insert}}",
+                    "type": "number",
+                    "ui:options": {"controls": False},
+                    "title": _l("Batch Size"),
+                    "default": 500,
+                    "minimum": 1,
+                    "maximum": 2000,
+                    "description": _l("Number of rows to insert in each batch"),
+                },
+                "insert_concurrency": {
+                    "ui:hidden": "{{!parentFormData.using_insert}}",
+                    "type": "number",
+                    "ui:options": {"controls": False},
+                    "title": _l("Concurrent Inserts"),
+                    "default": 1,
+                    "minimum": 1,
+                    "maximum": 10,
+                    "description": _l("Number of parallel insert operations"),
+                },
+                # "pre_queries": {
+                #     "type": "string",
+                #     "title": "Queries Ran Before Loading",
+                #     "description": '新数据导入前运行的 SQL，多条 SQL 用 `;` 分隔；支持传入变量，详见 <a target="_blank" href="http://bit.ly/2JMutjn">文档</a>',
+                #     "ui:field": "CodeEditorWithReferencesField",
+                #     "ui:options": {
+                #         "type": "code",
+                #         "lang": "sql",
+                #         "sqlLang": "mysql",
+                #     },
+                # },
+                # "post_queries": {
+                #     "type": "string",
+                #     "title": "Queries Ran After Loading",
+                #     "description": '新数据导入后运行的 SQL，多条 SQL 用 `;` 分隔；支持传入变量，详见 <a target="_blank" href="http://bit.ly/2JMutjn">文档</a>',
+                #     "ui:field": "CodeEditorWithReferencesField",
+                #     "ui:options": {
+                #         "type": "code",
+                #         "lang": "sql",
+                #         "sqlLang": "mysql",
+                #     },
+                # },
+            },
+            "required": ["data_source_name", "table"],
+        }
+        return schema

recurvedata/operators/transfer_operator/load_task_yicrowds.py

@@ -0,0 +1,97 @@
+import logging
+import os
+
+import jsonschema
+import pandas as pd
+
+from recurvedata.pigeon.schema import Schema
+from recurvedata.pigeon.utils import fs
+
+try:
+    import pyhive.hive  # noqa see https://github.com/cloudera/impyla/issues/277
+    from yicrowds_client.core import YiCrowdsClient
+except ImportError:
+    pass
+from recurvedata.operators.transfer_operator.task import LoadTask
+
+logger = logging.getLogger(__name__)
+
+
+class YiCrowdsLoadTask(LoadTask):
+    def _infer_column(self):
+        schema_file = fs.schema_filename(self.filename)
+        logger.info("infer column by schema file %s", schema_file)
+        if not os.path.exists(schema_file):
+            logger.error("file not exists, not able to infer column")
+        schema = Schema.load(schema_file)
+        columns = []
+        for f in schema:
+            columns.append(f.name)
+        return columns
+
+    def add_column_to_file(self, columns):
+        df = pd.read_csv(self.filename, header=None, names=columns)
+        df.to_csv(self.filename, index=False)
+
+    def execute_impl(self, *args, **kwargs):
+        conf = self.rendered_config
+        self.add_column_to_file(self._infer_column())
+        environment = conf.get("environment", "prod")
+        project_id = conf.get("yc_project_id")
+        tag = conf.get("yc_tag")
+        no_error = YiCrowdsClient(env=environment).upload_data_local(project_id=project_id, tags=tag, fp=self.filename)
+        if not no_error:
+            raise Exception("upload_data_local上传异常")
+
+    @classmethod
+    def validate(cls, configuration):
+        conf = super().validate(configuration)
+        environment = conf.get("environment", "prod")
+        project_id = conf.get("yc_project_id")
+        tag = conf.get("yc_tag")
+
+        yca = YiCrowdsClient(env=environment)
+        if not yca.if_project_id_exists(project_id=project_id):
+            raise jsonschema.ValidationError(
+                message=f"YiCrowds Project ID: {project_id} does not exist", path=("yc_project_id",)
+            )
+
+        if yca.if_tag_exists(project_id=project_id, tags=tag, is_regex=conf.get("is_regex")):
+            raise jsonschema.ValidationError(
+                message=f"YiCrowds Tag: {tag} exist", path=("yc_tag", "yc_project_id", "is_regex")
+            )
+
+    @classmethod
+    def config_schema(cls):
+        return {
+            "type": "object",
+            "properties": {
+                "environment": {
+                    "type": "string",
+                    "title": "Environment",
+                    "enum": ["prod", "dev"],
+                    "enumNames": ["prod", "dev"],
+                    "default": "prod",
+                    "description": "YiCrowds 有正式环境和测试环境，默认正式环境",
+                },
+                "yc_project_id": {
+                    "type": "string",
+                    "title": "YiCrowds Project ID",
+                    "description": "YiCrowds项目ID，可以通过URL地址获取",
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "yc_tag": {
+                    "type": "string",
+                    "title": "YiCrowds Tag",
+                    "description": "需要写入数据的tag",
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+            },
+            "required": ["yc_project_id", "yc_tag"],
+        }

recurvedata/operators/transfer_operator/mixin.py

@@ -0,0 +1,31 @@
+import csv
+import logging
+import os
+
+try:
+    from recurvedata.pigeon.handler.csv_handler import HiveCSVFileHandler
+except ImportError:
+    pass
+logger = logging.getLogger(__name__)
+
+
+class HiveTextfileConverterMixin(object):
+    def convert_csv_to_hive_text_if_needed(self):
+        """把 CSV 文件转换成 Hive 的文本文件"""
+        if not self.handler_options.get("hive"):
+            return
+
+        if not os.path.exists(self.filename):
+            logger.warning("%s is not exists", self.filename)
+            return
+
+        new_name = f"{self.filename}.hive"
+        handler = HiveCSVFileHandler(filename=new_name)
+        with open(self.filename, newline="") as inf:
+            reader = csv.reader(inf)
+            for row in reader:
+                handler.handle(tuple(row))
+        handler.close()
+
+        if os.path.exists(new_name):
+            os.rename(new_name, self.filename)