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

recurvedata/operators/transfer_operator/load_task_google_sheet.py

@@ -0,0 +1,130 @@
+import csv
+import logging
+from typing import Any
+
+from recurvedata.core.translation import _l
+from recurvedata.operators.transfer_operator import const
+from recurvedata.operators.transfer_operator.task import LoadTask
+from recurvedata.pigeon.utils import fs
+
+logger = logging.getLogger(__name__)
+GOOGLE_SHEET_MAX_ROWS = 1000000
+GOOGLE_SHEET_MAX_COLUMNS = 18278
+
+
+class GoogleSheetLoadTask(LoadTask):
+    ds_name_fields = ("google_service_account",)
+    should_write_header = True
+    worker_install_require = ["gspread"]
+
+    @staticmethod
+    def check_csv_content(filename: str) -> tuple[int, int]:
+        """Check if the CSV file row and column counts exceed the maximum allowed limits for Google Sheets."""
+        row_count = 0
+        col_count = 0
+        with open(filename, "r") as file:
+            reader = csv.reader(file)
+            for row in reader:
+                row_count += 1
+                if row_count == 1:
+                    col_count = len(row)
+                    if col_count > GOOGLE_SHEET_MAX_COLUMNS:
+                        raise ValueError(
+                            f"CSV file contains {col_count} columns, which exceeds the maximum allowed "
+                            f"{GOOGLE_SHEET_MAX_COLUMNS} columns in Google Sheets."
+                        )
+                if row_count > GOOGLE_SHEET_MAX_ROWS:
+                    raise ValueError(
+                        f"CSV file contains {row_count} rows, which exceeds the maximum allowed "
+                        f"{GOOGLE_SHEET_MAX_ROWS} rows in Google Sheets."
+                    )
+        return row_count, col_count
+
+    def execute_impl(self, *args: Any, **kwargs: Any) -> None:
+        import pandas as pd
+
+        if fs.is_file_empty(self.filename):
+            logger.warning("File %s does not exist or has no content, skipping.", self.filename)
+            return
+
+        ds = self.must_get_connection_by_name(self.config["google_service_account"])
+        service_account = ds.recurve_connector
+        _, sheet_id = service_account.parse_sheet_url(self.config["file_url"])
+        sheet = service_account.get_sheet(self.config["file_url"], sheet_id)
+
+        logger.info(f'Loading to {self.config["file_url"]}, gid {sheet_id}')
+
+        # Perform all necessary checks
+        csv_row_count, csv_col_count = self.check_csv_content(self.filename)
+        current_sheet_rows, current_sheet_cols = sheet.row_count, sheet.col_count
+
+        if self.config["mode"] == const.LOAD_APPEND:
+            csv_row_count += current_sheet_rows
+            csv_col_count = max(current_sheet_cols, csv_col_count)
+
+        if csv_row_count > GOOGLE_SHEET_MAX_ROWS:
+            raise ValueError(
+                f"Appending the CSV file will exceed the maximum allowed {GOOGLE_SHEET_MAX_ROWS} rows in Google Sheets."
+            )
+        if csv_col_count > GOOGLE_SHEET_MAX_COLUMNS:
+            raise ValueError(
+                f"Appending the CSV file will exceed the maximum allowed {GOOGLE_SHEET_MAX_COLUMNS} columns in Google Sheets."
+            )
+
+        # Load the CSV file into a DataFrame after checking the row count
+        df = pd.read_csv(self.filename, keep_default_na=False)
+        df.fillna("", inplace=True)
+
+        try:
+            service_account.load_df_to_sheet(df, sheet, self.config["mode"], value_input_option="USER_ENTERED")
+            logger.info(
+                f'Data loaded successfully into {self.config["file_url"]}, mode: {self.config["mode"]}, '
+                f"rows: {csv_row_count}, cols: {csv_col_count}"
+            )
+        except Exception as e:
+            logger.error(f'Failed to load data into {self.config["file_url"]}: {e}')
+            raise
+
+    @classmethod
+    def config_schema(cls) -> dict[str, Any]:
+        schema = {
+            "type": "object",
+            "properties": {
+                "google_service_account": {
+                    "type": "string",
+                    "title": _l("Google Service Account Connection"),
+                    "description": _l(
+                        "Select the Google Service Account connection with write permissions to the target spreadsheet"
+                    ),
+                    "ui:field": "ProjectConnectionSelectorField",
+                    "ui:options": {
+                        "supportTypes": ["google_service_account"],
+                    },
+                },
+                "file_url": {
+                    "type": "string",
+                    "title": _l("Google Sheet URL"),
+                    "description": _l(
+                        "URL of the target Google Sheet in format: "
+                        "https://docs.google.com/spreadsheets/d/{Spreadsheet ID}/edit#gid={Sheet GID}. "
+                        "If no sheet GID is specified, the first sheet will be used."
+                    ),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "mode": {
+                    "type": "string",
+                    "title": _l("Import Mode"),
+                    "enum": [const.LOAD_OVERWRITE, const.LOAD_APPEND],
+                    "enumNames": [const.LOAD_OVERWRITE, const.LOAD_APPEND],
+                    "default": const.LOAD_OVERWRITE,
+                    "description": _l(
+                        "OVERWRITE: Replace existing data with new data. " "APPEND: Add new data after existing data."
+                    ),
+                },
+            },
+            "required": ["google_service_account", "file_url", "mode"],
+        }
+        return schema

recurvedata/operators/transfer_operator/load_task_hive.py

@@ -0,0 +1,158 @@
+import copy
+import glob
+import json
+import os
+
+try:
+    from recurvedata.pigeon.loader.csv_to_hive import CSVToHiveLoader
+    from recurvedata.pigeon.utils import fs
+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 HiveLoadTask(LoadTask):
+    ds_name_fields = ("hive_data_source_name",)
+    ds_types = ("hive",)
+    default_dumper_handler_options = {
+        "hive": True,
+        "merge_files": False,  # do not merge intermediate files, pass in file pattern
+    }
+    worker_install_require = ["pigeon[hive_impala]"]
+
+    def execute_impl(self, *args, **kwargs):
+        hive_ds = self.must_get_connection_by_name(self.config["hive_data_source_name"])
+
+        load_options = self.rendered_config.copy()
+        for k in ["hive_data_source_name", "impala_data_source_name"]:
+            load_options.pop(k, None)
+
+        partition = load_options.pop("partition", None)
+        if partition:
+            load_options["partition"] = json.loads(partition)
+
+        sub_files = glob.glob(f"{self.filename}.[0-9]*")
+        if os.path.exists(self.filename) and not sub_files:
+            # dumper merged file
+            filename = self.filename
+        else:
+            # dump without merging, use pattern
+            # if upstream dump result is empty, sub_files is empty array, force to [self.filename] to ensure array is not empty
+            if all([fs.is_file_empty(x) for x in sub_files]):
+                sub_files = [self.filename]
+            filename = sub_files
+
+        load_options.update(
+            {
+                "filename": filename,
+                "hive_connector": hive_ds.connector,
+                "delete_file": True,
+            }
+        )
+        impala_ds = self.get_connection_by_name(self.config["impala_data_source_name"])
+        if impala_ds:
+            load_options.update({"impala_connector": impala_ds.connector})
+        loader = CSVToHiveLoader(**load_options)
+        return loader.execute()
+
+    @classmethod
+    def config_schema(cls):
+        # hive_dss = cls.get_connection_names_by_type(cls.ds_types)
+        # impala_dss = cls.get_connection_names_by_type('impala')
+        schema = {
+            "type": "object",
+            "properties": {
+                "hive_data_source_name": {
+                    "type": "string",
+                    "title": _l("Hive Connection"),
+                    "ui:field": "ProjectConnectionSelectorField",
+                    "ui:options": {
+                        "supportTypes": cls.ds_types,
+                    },
+                    # 'default': cls.first_or_default(hive_dss, ''),
+                },
+                "impala_data_source_name": {
+                    "type": "string",
+                    "title": _l("Impala Connection"),
+                    "description": _l("Optional Impala connection for faster data loading"),
+                    "ui:field": "ProjectConnectionSelectorField",
+                    "ui:options": {
+                        "supportTypes": [
+                            "impala",
+                        ],
+                    },
+                    # 'default': cls.first_or_default(impala_dss, ''),
+                },
+                "database": {
+                    "type": "string",
+                    "title": _l("Database Name"),
+                    "description": _l("Name of the Hive database to load data into. Supports template variables."),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "table": {
+                    "type": "string",
+                    "title": _l("Table Name"),
+                    "description": _l("Name of the Hive table to load data into. Supports template variables."),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "create_table_ddl": {
+                    "type": "string",
+                    "title": _l("Table Creation SQL"),
+                    "description": _l(
+                        "SQL statement to create the table if it doesn't exist. "
+                        "PARQUET storage format is recommended for better performance. "
+                    ),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "code",
+                        "lang": "sql",
+                        "sqlLang": "hive",
+                    },
+                },
+                "partition": {
+                    "type": "string",
+                    "title": _l("Partition Specification"),
+                    "description": _l(
+                        "JSON object specifying the partition to load data into. "
+                        "For T+1 tasks, use {'dt': '{{ yesterday_dt }}'} to load yesterday's partition. "
+                        "Supports template variables."
+                    ),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "plain",
+                    },
+                },
+                "compression_codec": {
+                    "type": "string",
+                    "title": _l("Compression Method"),
+                    "enum": ["snappy", "none", "gzip"],
+                    "enumNames": ["snappy", "none", "gzip"],
+                    "description": _l(
+                        "Data compression format. 'none' for no compression, 'gzip' for maximum compression, "
+                        "'snappy' for balanced compression/performance."
+                    ),
+                    "default": "snappy",
+                },
+            },
+            # NOTE: frontend uses vue-json-schema, which has a bug where enum fields must be required...
+            "required": [
+                "hive_data_source_name",
+                "impala_data_source_name",
+                "database",
+                "table",
+                "mode",
+                "compression_codec",
+            ],
+        }
+        properties_schema = schema["properties"]
+        properties_schema.update(copy.deepcopy(utils.LOAD_COMMON))
+        return schema

recurvedata/operators/transfer_operator/load_task_microsoft_fabric.py

@@ -0,0 +1,105 @@
+import logging
+from typing import Any
+
+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
+
+try:
+    from recurvedata.pigeon.loader.csv_to_microsoft_fabric import CSVToMsFabricLoader
+except ImportError:
+    pass
+
+logger = logging.getLogger(__name__)
+
+
+class MicrosoftFabricLoadTask(LoadTask):
+    ds_name_fields = ("data_source_name",)
+    ds_types = ("microsoft_fabric",)
+    worker_install_require = ["pigeon[azure]"]
+
+    def execute_impl(self, *args: Any, **kwargs: Any) -> None:
+        ds = self.must_get_connection_by_name(self.config["data_source_name"])
+        load_options: dict[str, Any] = self.rendered_config.copy()
+        for k in ["data_source_name"]:
+            load_options.pop(k, None)
+        columns = load_options.get("columns", "")
+        columns = [x.strip() for x in columns.split(",")] if columns.strip(" ,") else []
+        load_options["lineterminator"] = "\r\n" if self.dump_task_type == "PythonDumpTask" else "0x0D0A"
+        load_options.update(
+            {
+                "filename": self.filename,
+                "connector": ds.connector,
+                "delete_file": True,
+                "using_insert": False,
+                "columns": columns,
+                "database": ds.database,
+                "schema": ds.data.get("schema"),
+                "compress": True,  # Enable compression for better performance
+                "blob_options": ds.data.get("blob_options", {}),
+            }
+        )
+        logger.info(load_options)
+        loader = CSVToMsFabricLoader(**load_options)
+        return loader.execute()
+
+    @classmethod
+    def config_schema(cls):
+        schema = {
+            "type": "object",
+            "properties": {
+                "data_source_name": {
+                    "type": "string",
+                    "title": _l("Microsoft Fabric Connection"),
+                    "description": _l("The Microsoft Fabric data source to load data into"),
+                    "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"),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "code",
+                        "lang": "sql",
+                        "sqlLang": "sql",
+                    },
+                },
+                "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",
+                    },
+                },
+            },
+            "required": ["data_source_name", "table"],
+        }
+        return schema

recurvedata/operators/transfer_operator/load_task_mssql.py

@@ -0,0 +1,153 @@
+import logging
+
+try:
+    from recurvedata.pigeon.loader.csv_to_mssql import CSVToMsSQLLoader
+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
+
+logger = logging.getLogger(__name__)
+
+
+class MsSQLLoadTask(LoadTask):
+    ds_name_fields = ("data_source_name",)
+    ds_types = ("mssql", "azure_mssql")
+    worker_install_require = ["pigeon[azure]"]
+
+    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)
+        columns = load_options.get("columns", "")
+        columns = [x.strip() for x in columns.split(",")] if columns.strip(" ,") else []
+        load_options.update(
+            {
+                "filename": self.filename,
+                "connector": ds.connector,
+                "delete_file": True,
+                "using_insert": False,  # 自动推导，优先使用批量加载文件
+                "columns": columns,
+                "database": ds.database,
+                "schema": ds.data.get("schema"),
+            }
+        )
+        logger.info(load_options)
+        loader = CSVToMsSQLLoader(**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("MSSQL Connection"),
+                    "description": _l("The MSSQL data source to load data into"),
+                    "ui:field": "ProjectConnectionSelectorField",
+                    "ui:options": {
+                        "supportTypes": cls.ds_types,
+                    },
+                },
+                # "database": {
+                #     "type": "string",
+                #     "title": _l("Target Database"),
+                #     "description": _l("Name of the database to load data into"),
+                #     "ui:field": "CodeEditorWithReferencesField",
+                #     "ui:options": {
+                #         "type": "plain",
+                #     },
+                # },
+                # "schema": {
+                #     "type": "string",
+                #     "title": _l("Database Schema"),
+                #     "description": _l("Schema name in the target database (default: dbo)"),
+                #     "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"),
+                    "ui:field": "CodeEditorWithReferencesField",
+                    "ui:options": {
+                        "type": "code",
+                        "lang": "sql",
+                        "sqlLang": "sql",
+                    },
+                },
+                "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",
+                    },
+                },
+                # 'insert_batch_size': {
+                #     'type': 'number',
+                #      "ui:options": {"controls": False},
+                #     'title': 'INSERT Batch Size',
+                #     'default': 500,
+                #     'minimum': 1,
+                #     'maximum': 2000,
+                #     'description': '如果不支持批量加载 INSERT 导入数据，该参数设定 batch 大小'
+                # }
+                # "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", "table"],
+        }
+        return schema