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

recurvedata/pigeon/loader/csv_to_starrocks.py

@@ -0,0 +1,233 @@
+import time
+from typing import TYPE_CHECKING, Any, List, Optional
+
+from recurvedata.pigeon import const
+from recurvedata.pigeon.loader.csv_to_mysql import CSVToMySQLLoader
+from recurvedata.pigeon.utils import md5hash, randomized_suffix
+from recurvedata.pigeon.utils.sql import bak_table_of, reconcile_table_of, staging_table_of
+
+if TYPE_CHECKING:
+    from recurvedata.pigeon.connector.starrocks import StarRocksConnector
+
+allowed_modes = (const.LOAD_OVERWRITE, const.LOAD_MERGE, const.LOAD_APPEND)
+
+
+class CSVToStarRocksLoader(CSVToMySQLLoader):
+    def __init__(
+        self,
+        database: str,
+        table: str,
+        filename: str,
+        connector: Optional["StarRocksConnector"] = None,
+        create_table_ddl: Optional[str] = None,
+        mode: str = const.LOAD_OVERWRITE,
+        primary_keys: Optional[List[str]] = None,
+        skiprows: int = 0,
+        columns: Optional[List[str]] = None,
+        using_insert: bool = False,
+        insert_batch_size: int = 1000,
+        insert_concurrency: int = 1,
+        delete_file: bool = False,
+        pre_queries: Optional[List[str]] = None,
+        post_queries: Optional[List[str]] = None,
+        load_strict_mode: bool = False,
+        *args: Any,
+        **kwargs: Any,
+    ):
+        if not connector:
+            raise ValueError(f"connector is required for {self.__class__.__name__}")
+        self.load_strict_mode: bool = load_strict_mode
+        connector.load_strict_mode = load_strict_mode
+        self.logger.info(f"load_strict_mode: {load_strict_mode}")
+        # the same filename incoming, the same intermediate tables will be generated
+        # and the previous fail intermediate tables will be cleaned in a new try
+        table_suffix: str = md5hash(filename)[:6] if filename is not None else randomized_suffix()
+        self.__staging_table: str = staging_table_of(table) + "_" + table_suffix
+        self.__reconcile_table: str = reconcile_table_of(table) + "_" + table_suffix
+        self.__bak_table: str = bak_table_of(table) + "_" + table_suffix
+        if any(
+            [
+                len(self.__staging_table) > 64,
+                len(self.__reconcile_table) > 64,
+                len(self.__bak_table) > 64,
+            ]
+        ):
+            self.logger.error(
+                f"table name {self.__staging_table} 's length: {len(self.__staging_table)}\n"
+                f"table name {self.__reconcile_table}'s length: {len(self.__reconcile_table)}\n"
+                f"table name {self.__bak_table}'s length: {len(self.__bak_table)}\n"
+            )
+            raise ValueError("length of intermediate table name is greater than 64!")
+        super().__init__(
+            database=database,
+            table=table,
+            filename=filename,
+            connector=connector,
+            create_table_ddl=create_table_ddl,
+            mode=mode,
+            primary_keys=primary_keys,
+            skiprows=skiprows,
+            columns=columns,
+            using_insert=using_insert,
+            insert_batch_size=insert_batch_size,
+            insert_concurrency=insert_concurrency,
+            delete_file=delete_file,
+            pre_queries=pre_queries,
+            post_queries=post_queries,
+            *args,
+            **kwargs,
+        )
+
+    @property
+    def staging_table(self) -> str:
+        """
+        overwrite method, return a table name with randomized postfix
+        """
+        return self.__staging_table
+
+    def _merge_into_target_table(self) -> None:
+        queries = []
+        if self.mode == const.LOAD_MERGE:
+            queries.extend(self._ingest_by_merging())
+        elif self.mode == const.LOAD_OVERWRITE:
+            # bak_table = bak_table_of(self.table)
+            bak_table = self.__bak_table
+            queries.extend(
+                [
+                    f"DROP TABLE IF EXISTS {bak_table}",
+                    f"ALTER TABLE {self.table} RENAME {bak_table}",
+                    f"ALTER TABLE {self.staging_table} RENAME {self.table}",
+                    f"DROP TABLE IF EXISTS {bak_table}",
+                ]
+            )
+        else:
+            # special process at `APPEND` mode, cuz an occasional error happens:
+            # ================================== ERROR MSG START ======================================
+            # pymysql.err.ProgrammingError: (1064, 'Unexpected exception: Failed to drop table {table_name}.
+            # msg: There are still some transactions in the COMMITTED state waiting to be completed.
+            # The table {table_name} cannot be dropped. If you want to forcibly drop(cannot be recovered),
+            # please use "DROP TABLE <table> FORCE".')
+            # ================================== ERROR MSG END ========================================
+            # Here's the optimization: commit insert statement first, make it blocking until finished.
+            queries.append(f"INSERT INTO {self.table} SELECT * FROM {self.staging_table}")
+            self.connector.execute(self.pre_queries + queries, autocommit=True, commit_on_close=False)
+
+            queries.clear()
+            queries.append(f"DROP TABLE {self.staging_table}")
+            self.connector.execute(queries + self.post_queries, autocommit=True, commit_on_close=False)
+            return
+
+        queries = self.pre_queries + queries + self.post_queries
+        self.logger.info("running MySQL queries within a transaction")
+        self.connector.execute(queries, autocommit=False, commit_on_close=True)
+
+    def _ingest_by_merging(self) -> List[str]:
+        """Merge with deduplication based on primary keys using StarRocks-compatible syntax"""
+        # First, deduplicate staging table based on primary keys using window function
+        pk_columns = ", ".join(self.primary_keys)
+
+        # Get all columns from staging table (excluding the rn column we'll add)
+        cols = self.connector.get_columns(self.staging_table)
+        cols_str = ", ".join(self.connector.quote_identifier(x) for x in cols)
+
+        # Create a temporary table with deduplicated data
+        tmp_table = f"{self.staging_table}_dedup"
+        dedup_sql = f"""
+            DROP TABLE IF EXISTS {tmp_table};
+            CREATE TABLE {tmp_table} LIKE {self.staging_table};
+            INSERT INTO {tmp_table}
+            SELECT {cols_str} FROM (
+                SELECT *, ROW_NUMBER() OVER(PARTITION BY {pk_columns} ORDER BY {pk_columns}) AS rn
+                FROM {self.staging_table}
+            ) t
+            WHERE rn = 1;
+        """
+
+        # Replace staging table with deduplicated data
+        replace_sql = f"""
+            DROP TABLE {self.staging_table};
+            ALTER TABLE {tmp_table} RENAME {self.staging_table};
+        """
+
+        # Simple merge: backup target table, then merge and deduplicate in one step
+        bak = self.__bak_table
+        table = self.connector.quote_identifier(self.table)
+        staging = self.connector.quote_identifier(self.staging_table)
+        bak = self.connector.quote_identifier(bak)
+
+        # Simple and efficient merge: backup + merge + deduplicate in one operation
+        merge_sql = f"""
+            -- Backup target table
+            DROP TABLE IF EXISTS {bak};
+            ALTER TABLE {table} RENAME {bak};
+
+            -- Create new target table and insert deduplicated merged data in one step
+            CREATE TABLE {table} AS
+            SELECT {cols_str} FROM (
+                SELECT *, ROW_NUMBER() OVER(PARTITION BY {pk_columns} ORDER BY {pk_columns}) AS rn
+                FROM (
+                    SELECT * FROM {bak}
+                    UNION ALL
+                    SELECT * FROM {staging}
+                ) combined
+            ) t WHERE rn = 1;
+
+            -- Clean up
+            DROP TABLE {bak};
+            DROP TABLE {staging};
+        """
+
+        return [dedup_sql, replace_sql, merge_sql]
+
+    def execute(self) -> None:
+        """
+        overwrite method, implemented try...catch...
+        """
+        self.before_execute()
+        try:
+            self.execute_impl()
+        except Exception as e:
+            self.handle_exception()
+            raise e
+        self.after_execute()
+
+    def _prepare_staging_table(self):
+        queries = """
+            DROP TABLE IF EXISTS {staging};
+            CREATE TABLE {staging} LIKE {table};
+        """.format(
+            staging=self.staging_table, table=self.table
+        )
+        self.connector.execute(queries, autocommit=True)
+        time.sleep(5)  # wait for table to be created and visible
+
+    def handle_exception(self) -> None:
+        """
+        ensure all intermediate tables are cleaned safely after catch the exception
+        """
+        qry_exists_sql = """
+        SELECT 1 FROM information_schema.tables
+        WHERE table_schema = '{database}' AND table_name = '{table}';
+        """
+        is_table_exists = self.connector.fetchall(qry_exists_sql.format(database=self.database, table=self.table))
+        is_bak_exists = self.connector.fetchall(qry_exists_sql.format(database=self.database, table=self.__bak_table))
+        if is_table_exists:
+            # clean intermediate tables directly.
+            queries = [
+                f"DROP TABLE IF EXISTS {self.__bak_table}",
+                f"DROP TABLE IF EXISTS {self.__staging_table}",
+                f"DROP TABLE IF EXISTS {self.__reconcile_table}",
+            ]
+        elif is_bak_exists:
+            # rollback from bak_table
+            queries = [
+                f"ALTER TABLE {self.__bak_table} RENAME {self.table}",
+                f"DROP TABLE IF EXISTS {self.__staging_table}",
+                f"DROP TABLE IF EXISTS {self.__reconcile_table}",
+            ]
+        else:
+            queries = [
+                f"DROP TABLE IF EXISTS {self.__staging_table}",
+                f"DROP TABLE IF EXISTS {self.__reconcile_table}",
+            ]
+        self.connector.execute(queries, autocommit=False, commit_on_close=True)

recurvedata/pigeon/meta.py

@@ -0,0 +1,116 @@
+import datetime
+import json
+
+import cytoolz as toolz
+
+from recurvedata.pigeon.schema import Schema
+
+
+class JSONEncoder(json.JSONEncoder):
+    def default(self, o):
+        if isinstance(o, datetime.date):
+            return o.isoformat()
+        if isinstance(o, datetime.timedelta):
+            return str(o)
+        if isinstance(o, Schema):
+            return o.to_list()
+        return super().default(o)
+
+
+class Meta(object):
+    def to_dict(self):
+        raise NotImplementedError()
+
+    def to_json(self, **kwargs):
+        params = toolz.merge({"sort_keys": True, "ensure_ascii": False, "cls": JSONEncoder}, kwargs)
+        return json.dumps(self.to_dict(), **params)
+
+
+class HandlerMeta(Meta):
+    def __init__(self):
+        self.reset()
+
+    def reset(self):
+        self.num_input_rows = 0
+        self.num_output_rows = 0
+        self.num_error_rows = 0
+        self.error_log_size = 0  # 当前打印的报错 row 字符串的字符数
+
+    def to_dict(self):
+        return {
+            "num_input_rows": self.num_input_rows,
+            "num_output_rows": self.num_output_rows,
+            "num_error_rows": self.num_error_rows,
+            "error_log_size": self.error_log_size,
+        }
+
+
+class HandlerFactoryMeta(HandlerMeta):
+    def __init__(self, name):
+        self.name = name
+        super().__init__()
+
+    def update(self, handler_meta):
+        self.num_input_rows += handler_meta.num_input_rows
+        self.num_output_rows += handler_meta.num_output_rows
+        self.num_error_rows += handler_meta.num_error_rows
+        self.error_log_size += handler_meta.error_log_size
+
+    def to_dict(self):
+        d = super().to_dict()
+        d["name"] = self.name
+        return d
+
+
+class DumperWorkerMeta(Meta):
+    def __init__(self):
+        self.num_dumped_rows = 0
+        self.schema = None
+        self.handlers_meta = None
+
+    def to_dict(self):
+        return {
+            "num_dumped_rows": self.num_dumped_rows,
+            "schema": self.schema,
+            "handlers_meta": [x.to_dict() for x in self.handlers_meta],
+        }
+
+
+class DumperMeta(Meta):
+    def __init__(self, context=None):
+        self.time_start = None
+        self.time_finish = None
+        self.num_dumped_rows = 0
+        self.context = context
+        self.schema = None
+        self.handlers_meta = []
+
+    def mark_start(self):
+        self.time_start = datetime.datetime.now()
+
+    def mark_finish(self):
+        self.time_finish = datetime.datetime.now()
+
+    @property
+    def rows_per_second(self):
+        return self.num_dumped_rows / (self.time_finish - self.time_start).total_seconds()
+
+    @property
+    def duration(self):
+        if not self.time_start:
+            return None
+        if not self.time_finish:
+            return datetime.datetime.now() - self.time_start
+        return self.time_finish - self.time_start
+
+    def to_dict(self):
+        return {
+            "time_start": self.time_start,
+            "time_finish": self.time_finish,
+            "time_duration": self.duration,
+            "num_dumped_rows": self.num_dumped_rows,
+            "rows_per_second": self.rows_per_second,
+            "context": self.context,
+            "schema": self.schema,
+            "handlers_meta": [x.to_dict() for x in self.handlers_meta],
+        }

recurvedata/pigeon/row_factory.py

@@ -0,0 +1,42 @@
+import logging
+from collections import OrderedDict
+
+from recurvedata.pigeon.utils.keyed_tuple import KeyedTuple
+
+logger = logging.getLogger(__name__)
+
+
+def tuple_factory(colnames, row):
+    """Returns each row as a tuple"""
+    return row
+
+
+def keyed_tuple_factory(colnames, row):
+    return KeyedTuple(row, colnames)
+
+
+def dict_factory(colnames, row):
+    return dict(zip(colnames, row))
+
+
+def ordered_dict_factory(colnames, row):
+    return OrderedDict(zip(colnames, row))
+
+
+def get_row_keys(row):
+    if isinstance(row, dict):
+        # created by dict_factory or ordered_dict_factory
+        return list(row.keys())
+    if hasattr(row, "_fields"):
+        # created by keyed_tuple_factory
+        return list(row._fields)
+    else:
+        # created by tuple_factory, which is not able to know the keys
+        return None
+
+
+def get_row_values(row):
+    if isinstance(row, dict):
+        # created by dict_factory or ordered_dict_factory
+        return list(row.values())
+    return list(row)

recurvedata/pigeon/schema/__init__.py

@@ -0,0 +1,124 @@
+import json
+
+import dateutil.parser
+
+from recurvedata.pigeon.schema import types
+
+
+class Field(object):
+    def __init__(self, name, type, size=None, comment=None, extra=None):
+        self.name = name
+        self.type = type
+        self.size = size
+        self.comment = comment
+        self.extra = extra
+
+        self._cast_func = {
+            types.INT8: self._cast_to_int,
+            types.INT16: self._cast_to_int,
+            types.INT32: self._cast_to_int,
+            types.INT64: self._cast_to_int,
+            types.FLOAT32: self._cast_to_float,
+            types.FLOAT64: self._cast_to_float,
+            types.BOOLEAN: self._cast_to_boolean,
+            types.DATETIME: self._cast_to_datetime,
+            types.DATE: self._cast_to_date,
+            types.JSON: self._cast_to_json,
+        }.get(self.type, self._cast_pass)
+
+    def cast(self, value):
+        if value is None:
+            return None
+        if value == "NULL":
+            return None
+        return self._cast_func(value)
+
+    def _cast_pass(self, value):
+        return value
+
+    def _cast_to_int(self, value: str):
+        if value == "":
+            return 0
+        return int(value)
+
+    def _cast_to_float(self, value: str):
+        if value == "":
+            return 0.0
+        return float(value)
+
+    def _cast_to_boolean(self, value: str):
+        if value.lower() in ("", "0", "false"):
+            return False
+        return True
+
+    def _cast_to_datetime(self, value: str):
+        if value == "":
+            return None
+        return dateutil.parser.parse(value)
+
+    def _cast_to_date(self, value: str):
+        if value == "":
+            return None
+        return dateutil.parser.parse(value).date()
+
+    def _cast_to_json(self, value: str):
+        if value in ("",):
+            # 正常情况下不会有 ''，很可能是从 CSV 文件读到了空字符，当作 None 处理
+            return None
+        return json.loads(value)
+
+    def to_dict(self):
+        return {
+            "name": self.name,
+            "type": self.type,
+            "size": self.size,
+            "comment": self.comment,
+            "extra": self.extra,
+        }
+
+    def __str__(self):
+        return f'<Field ("{self.name}", "{self.type}")>'
+
+    def __repr__(self):
+        return f'<Field ("{self.name}", "{self.type}")>'
+
+
+class Schema(object):
+    def __init__(self, fields=None):
+        self.fields = fields or []
+
+    def add_field(self, field):
+        # TODO(liyangliang): clean field names, handle special characters and duplications
+        self.fields.append(field)
+
+    def add_field_by_attrs(self, name, type, size=None, comment=None, extra=None):
+        self.add_field(Field(name, type, size, comment, extra))
+
+    def remove_field(self, name):
+        self.fields = [x for x in self.fields if x.name != name]
+
+    def keep_fields(self, names):
+        self.fields = [x for x in self.fields if x.name in names]
+
+    @property
+    def field_names(self):
+        return [x.name for x in self.fields]
+
+    def __iter__(self):
+        return iter(self.fields)
+
+    def to_list(self):
+        return [x.to_dict() for x in self.fields]
+
+    def to_json(self):
+        return json.dumps(self.to_list())
+
+    def dump(self, filename):
+        with open(filename, "w") as f:
+            json.dump(self.to_list(), f, indent=2)
+
+    @classmethod
+    def load(cls, filename):
+        with open(filename) as f:
+            data = json.load(f)
+        return cls([Field(**item) for item in data])

recurvedata/pigeon/schema/types.py

@@ -0,0 +1,13 @@
+INT8 = "INT8"  # 1-byte (8-bit) signed integers
+INT16 = "INT16"  # 2-byte (16-bit) signed integers
+INT32 = "INT32"  # 4-byte (32-bit) signed integers
+INT64 = "INT64"  # 8-byte (64-bit) signed integers
+FLOAT32 = "FLOAT32"  # 4-byte (32-bit) single-precision floating
+FLOAT64 = "FLOAT64"  # 8-byte (64-bit) double-precision floating
+BOOLEAN = "BOOLEAN"
+
+DATETIME = "DATETIME"
+DATE = "DATE"
+
+STRING = "STRING"
+JSON = "JSON"