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

recurvedata/operators/transfer_operator/load_task_azure_blob.py

@@ -0,0 +1,115 @@
+import logging
+import os
+
+import jsonschema
+
+try:
+    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 AzureBlobStorageLoadTask(LoadTask):
+    ds_name_fields = ("data_source_name",)
+    ds_types = ("azure_blob",)
+    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
+
+        azure_blob_ds = self.must_get_connection_by_name(self.config["data_source_name"])
+
+        azure_blob = azure_blob_ds.connector
+        # 是否自动创建 container
+        config_container = self.rendered_config["container"] or azure_blob_ds.extra.get("container")
+        # if not self.rendered_config['auto_create_container'] and config_container:
+        #     if not azure_blob.exists(container_name=config_container):
+        #         raise ValueError(f'{config_container} not exists')
+
+        # compress or not
+        compress_mode = self.rendered_config["compress_mode"]
+        file_upload, _ = self.compress_file(filename=self.filename, compress_mode=compress_mode)
+
+        blob_name = self.rendered_config["blob"] or os.path.basename(self.filename)
+        logger.info(f"uploading {file_upload} to {config_container}/{blob_name}...")
+        config = {
+            "blob_name": blob_name,
+            "overwrite": self.rendered_config["overwrite"],
+            "local_file_path": file_upload,
+            "container_name": config_container,
+        }
+        azure_blob.upload(**config)
+        fs.remove_files_safely([self.filename, file_upload])
+
+    @classmethod
+    def validate(cls, configuration):
+        config = super().validate(configuration)
+        must_get_by_name = cls.must_get_connection_by_name
+
+        if not config.get("container"):
+            azure_blob = must_get_by_name(configuration["data_source_name"])
+            if not azure_blob.extra.get("container"):
+                # ensure container
+                raise jsonschema.ValidationError(message="Unknown Container", path=("container",))
+        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("Azure Blob Connection"),
+                    "ui:field": "ProjectConnectionSelectorField",
+                    "ui:options": {
+                        "supportTypes": cls.ds_types,
+                    },
+                    # 'default': cls.first_or_default(dss, ''),
+                },
+                "container": {
+                    "type": "string",
+                    "title": _l("Container"),
+                    "description": _l("Container name, required if not set in data source"),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "blob": {
+                    "type": "string",
+                    "title": _l("Blob Name"),
+                    "description": _l("Blob name in the container. Jinja templating is supported."),
+                    "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", "Bzip2"],
+                    "enumNames": ["None", "Gzip", "Zip", "Bzip2"],
+                    "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_azure_synapse.py

@@ -0,0 +1,90 @@
+import copy
+
+try:
+    from recurvedata.pigeon.loader.csv_to_azure_synapse import CSVToAzureSynapseLoader
+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 AzureSynapseLoadTask(LoadTask):
+    ds_name_fields = ("azure_synapse_data_source_name",)
+    ds_types = ("azure_synapse",)
+    default_dumper_handler_options = {}
+    worker_install_require = ["pigeon"]
+
+    def execute_impl(self, *args, **kwargs):
+        azure_synapse_ds = self.must_get_connection_by_name(self.config["azure_synapse_data_source_name"])
+        load_options = self.rendered_config.copy()
+        for k in ["azure_synapse_data_source_name"]:
+            load_options.pop(k, None)
+        load_options.update(
+            {
+                "filename": self.filename,
+                "azure_synapse_connector": azure_synapse_ds.connector,
+                "delete_file": True,
+                "compress": True,
+            }
+        )
+        loader = CSVToAzureSynapseLoader(**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": {
+                "azure_synapse_data_source_name": {
+                    "type": "string",
+                    "title": _l("Azure Synapse Connection"),
+                    "ui:field": "ProjectConnectionSelectorField",
+                    "ui:options": {
+                        "supportTypes": cls.ds_types,
+                    },
+                    # 'default': cls.first_or_default(dws, ''),
+                },
+                "schema": {
+                    "type": "string",
+                    "title": _l("Database Schema"),
+                    "description": _l("Schema name in Azure Synapse database"),
+                    "default": "dbo",
+                    "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://learn.microsoft.com/en-us/sql/t-sql/statements/create-table-azure-sql-data-warehouse'>"
+                        "Azure Synapse Docs</a> for syntax."
+                    ),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "code",
+                        "lang": "sql",
+                        "sqlLang": "sql",
+                    },
+                },
+            },
+            "required": ["azure_synapse_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_clickhouse.py

@@ -0,0 +1,167 @@
+import csv
+
+try:
+    from recurvedata.pigeon.loader.csv_to_clickhouse import CSVToClickHouseLoader
+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 ClickHouseLoadTask(LoadTask):
+    ds_name_fields = ("data_source_name",)
+    ds_types = ("clickhouse",)
+    default_dumper_handler_options = {
+        "null": r"\N",
+        "quoting": csv.QUOTE_MINIMAL,
+    }
+    worker_install_require = ["pigeon[clickhouse]"]
+
+    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})
+        loader = CSVToClickHouseLoader(**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": {
+                "data_source_name": {
+                    "type": "string",
+                    "title": _l("ClickHouse Connection"),
+                    "ui:field": "ProjectConnectionSelectorField",
+                    "ui:options": {
+                        "supportTypes": cls.ds_types,
+                    },
+                    # 'default': cls.first_or_default(dss, ''),
+                },
+                "database": {
+                    "type": "string",
+                    "title": _l("Target Database"),
+                    "description": _l("Name of the database to load data into"),
+                    "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://clickhouse.com/docs/en/sql-reference/statements/create/table'>"
+                        "ClickHouse Docs</a> for syntax."
+                    ),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "code",
+                        "lang": "sql",
+                        "sqlLang": "sql",
+                    },
+                },
+                "table_engine": {
+                    "type": "string",
+                    "title": _l("Table Engine"),
+                    "description": _l(
+                        "Storage engine for the target table. Ignored if Table Creation SQL is provided. See "
+                        "<a target='_blank' href='https://clickhouse.com/docs/en/engines/table-engines'>"
+                        "ClickHouse Docs</a> for options."
+                    ),
+                    "default": "Log",
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "mode": {
+                    "type": "string",
+                    "title": _l("Load Mode"),
+                    "description": _l("How to handle existing data in the target table"),
+                    "enum": list(allowed_modes),
+                    "enumNames": 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": "Using INSERT",
+                #     "default": False,
+                #     "description": "默认使用 `clickhouse-client` 导入数据，出错时会回退到用 INSERT 语句批量导入数据",
+                # },
+                "insert_batch_size": {
+                    # "ui:hidden": "{{!parentFormData.using_insert}}",
+                    "type": "number",
+                    "ui:options": {"controls": False},
+                    "title": _l("Batch Size"),
+                    "description": _l("Number of rows to insert in each batch"),
+                    "default": 10000,
+                    "minimum": 1000,
+                    "maximum": 100000,
+                },
+                "insert_concurrency": {
+                    # "ui:hidden": "{{!parentFormData.using_insert}}",
+                    "type": "number",
+                    "ui:options": {"controls": False},
+                    "title": _l("Concurrent Inserts"),
+                    "default": 1,
+                    "minimum": 1,
+                    "maximum": 5,
+                    "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": "sql",
+                #     },
+                # },
+                # "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": "sql",
+                #     },
+                # },
+            },
+            "required": ["data_source_name", "database", "table", "insert_batch_size"],
+        }
+        return schema

recurvedata/operators/transfer_operator/load_task_doris.py

@@ -0,0 +1,164 @@
+try:
+    from recurvedata.pigeon.loader.csv_to_doris import CSVToDorisLoader
+except ImportError:
+    pass
+
+from typing import TYPE_CHECKING, Any, List, Tuple
+
+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
+
+if TYPE_CHECKING:
+    from recurvedata.connectors.pigeon import DataSource
+
+
+class DorisLoadTask(LoadTask):
+    ds_name_fields: Tuple[str] = ("data_source_name",)
+    ds_types: Tuple[str] = ("doris",)
+    default_dumper_handler_options = {
+        "null": r"\N",
+        "lineterminator": "\n",
+        "escapechar": "'",
+        "doublequote": False,
+    }
+    worker_install_require: List[str] = [
+        "pigeon[doris]",
+    ]
+
+    def execute_impl(self, *args, **kwargs) -> Any:
+        """Execute the Doris load task by loading CSV data into a Doris table."""
+        # Get the Doris data source connection
+        ds: "DataSource" = self.must_get_connection_by_name(self.config["data_source_name"])
+
+        # Copy and prepare the load options
+        load_options: dict = self.rendered_config.copy()
+        load_options.pop("data_source_name", None)
+
+        # Update with required loader options
+        load_options.update(
+            {
+                "filename": self.filename,
+                "connector": ds.connector,
+                "delete_file": True,  # Clean up CSV file after loading
+                "load_strict_mode": self.config.get("load_strict_mode", False),
+                "max_filter_ratio": self.config.get("max_filter_ratio", 0),
+                "database": ds.database,
+            }
+        )
+
+        # Initialize and execute the loader
+        loader = CSVToDorisLoader(**load_options)
+        return loader.execute()
+
+    @classmethod
+    def config_schema(cls) -> dict[str, Any]:
+        schema = {
+            "type": "object",
+            "properties": {
+                "data_source_name": {
+                    "type": "string",
+                    "title": _l("Doris Connection"),
+                    "ui:field": "ProjectConnectionSelectorField",
+                    "ui:options": {
+                        "supportTypes": cls.ds_types,
+                    },
+                },
+                "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://doris.apache.org/docs/sql-manual/sql-statements/table-and-view/table/CREATE-TABLE'>"
+                        "Doris Docs</a> for syntax."
+                    ),
+                    "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"),
+                },
+                "max_filter_ratio": {
+                    "type": "number",
+                    "ui:options": {"controls": False},
+                    "title": _l("Max Filter Ratio"),
+                    "default": 0,
+                    "minimum": 0,
+                    "maximum": 1,
+                    "description": _l(
+                        "The maximum tolerated ratio of filterable (e.g., non-compliant) data. Default is zero tolerance. Value range: 0~1. If the error rate during import exceeds this value, the import will fail."
+                    ),
+                    "ui:hidden": "{{parentFormData.using_insert}}",
+                },
+            },
+            "required": ["data_source_name", "table"],
+        }
+        return schema