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

recurvedata/pigeon/dumper/ftp.py

@@ -0,0 +1,64 @@
+import os
+
+import humanize
+
+from recurvedata.pigeon.connector.ftp import FtpConnector
+from recurvedata.pigeon.meta import DumperMeta
+from recurvedata.pigeon.utils import LoggingMixin
+
+
+class FtpDumperMeta(DumperMeta):
+    def __init__(self):
+        super().__init__()
+        self.total_size = 0
+        self.dumped_files = []
+
+    def to_dict(self):
+        speed = self.total_size / self.duration.total_seconds()
+        return {
+            "time_start": self.time_start,
+            "time_finish": self.time_finish,
+            "time_duration": self.duration,
+            "total_size": self.total_size,
+            "total_size_human": humanize.naturalsize(self.total_size, gnu=True),
+            "download_speed": f"{humanize.naturalsize(speed, gnu=True)}/s",
+            "num_dumped_files": len(self.dumped_files),
+            "dumped_files": self.dumped_files,
+        }
+
+
+class FtpDumper(LoggingMixin):
+    def __init__(self, connector, src, dst):
+        assert isinstance(connector, FtpConnector)
+        self.connector = connector
+        self.src = src
+        self.dst = dst
+        self.meta = FtpDumperMeta()
+
+    def execute(self):
+        self.meta.mark_start()
+        self.execute_impl()
+        self.meta.mark_finish()
+        self.logger.info(f"dumper meta: {self.meta.to_json(indent=2)}")
+        return self.meta
+
+    def execute_impl(self):
+        if self.connector.is_ftp_dir(self.src):
+            for item in self.connector.list_dir(self.src):
+                if self.connector.is_ftp_dir(item):
+                    self.logger.warning(f"{item} may be a directory. Skip")
+                else:
+                    _, remote_file = os.path.split(item)
+                    dst = os.path.join(self.dst, remote_file)
+                    self.connector.download_file(item, dst)
+                    self.collect_meta(dst)
+        else:
+            self.connector.download_file(self.src, self.dst)
+            self.collect_meta(self.dst)
+
+    def collect_meta(self, filepath):
+        if not os.path.exists(filepath):
+            return None
+        file_size = os.stat(filepath).st_size
+        self.meta.dumped_files.append({"filepath": filepath, "size": file_size})
+        self.meta.total_size += file_size

recurvedata/pigeon/dumper/mongodb.py

@@ -0,0 +1,103 @@
+from bson import json_util
+
+from recurvedata.pigeon.connector.mongodb import MongoDBConnector
+from recurvedata.pigeon.dumper.base import BaseDumper
+from recurvedata.pigeon.row_factory import ordered_dict_factory
+from recurvedata.pigeon.schema import Schema
+from recurvedata.pigeon.utils.timing import TimeCounter
+
+
+class MongoDBDumper(BaseDumper):
+    _row_factory = staticmethod(ordered_dict_factory)
+
+    def __init__(self, connector, database, collection, filter=None, projection=None, handler_factories=None):
+        """MongoDBDumper 用于从 MongoDB 导出数据
+        :param connector: MongoDBConnector 对象
+        :param database: database 名字
+        :param collection: collection 名字
+        :param filter: 查询条件，用于 find 函数。
+                       注意，如果传入字符串格式的 filter，将被当作 json 字符串，用 bson.json_util.loads 反序列化。
+        :param projection: 控制返回的字段，用于 find 函数
+        :param handler_factories:
+        """
+        super().__init__(handler_factories=handler_factories)
+
+        if not isinstance(connector, MongoDBConnector):
+            raise TypeError("connector should be instance of MongoDBConnector")
+
+        self.connector = connector
+        self.database = database
+        self.collection = collection
+
+        self.filter = filter or {}
+        if isinstance(self.filter, str):
+            self.filter = json_util.loads(self.filter)
+
+        self.projection = projection
+
+        self.meta.context = {
+            "database": database,
+            "collection": collection,
+            "filter": filter,
+            "projection": projection,
+        }
+
+    @property
+    def row_factory(self):
+        return ordered_dict_factory
+
+    @row_factory.setter
+    def row_factory(self, factory):
+        raise ValueError(f"{self.__class__.__name__}.row_factory is dict_factory, and is readonly")
+
+    def execute(self):
+        self.meta.mark_start()
+        self.execute_impl()
+        self.meta.mark_finish()
+        self.logger.info("dumper meta: %s", self.meta.to_json(indent=2))
+        self.handle_schema()
+        return self.meta
+
+    def execute_impl(self):
+        handlers = self.create_handlers()
+
+        self.logger.info("execute with context")
+        self.logger.info("  filter: %s", self.filter)
+        self.logger.info("  projection: %s", self.projection)
+
+        # MongoDB 没有模式，以第一条结果的字段和值来推导 schema
+        schema = Schema()
+        field_names = []
+        client = self.connector.connect()
+        if self.collection not in client[self.database].list_collection_names():
+            raise RuntimeError(f"collection '{self.collection}' does not exist")
+        total_count = client[self.database][self.collection].count_documents(self.filter)
+        cursor = client[self.database][self.collection].find(self.filter, self.projection)
+
+        counter = TimeCounter(name="", log_threshold=10000, logger=self.logger, total=total_count)
+        # Use projection field order as the base order
+        field_names = list(self.projection.keys()) if self.projection else []
+
+        for doc in cursor:
+            counter.incr(1)
+
+            fixed_doc = doc
+            # if projection is not None, then use projection to filter the fields and fill the missing fields with None
+            if field_names:
+                # Use projection field order, and fill missing fields with None
+                fixed_doc = self.row_factory(field_names, [doc.get(x, None) for x in field_names])
+
+            for h in handlers:
+                h.handle(fixed_doc)
+
+        counter.show_stat()
+        self.meta.schema = schema
+        self.meta.num_dumped_rows = counter.count
+
+        for hf, h in zip(self.handler_factories, handlers):
+            hf.meta.update(h.meta)
+        self.meta.handlers_meta = [x.meta for x in self.handler_factories]
+
+        for h in handlers:
+            h.close()
+        self.join_handlers()

recurvedata/pigeon/handler/__init__.py

@@ -0,0 +1,4 @@
+from recurvedata.pigeon.handler.base import Handler, HandlerFactory, NullHandler
+from recurvedata.pigeon.handler.csv_handler import CSVFileHandler, HiveCSVFileHandler
+
+null_factory = HandlerFactory(NullHandler)

recurvedata/pigeon/handler/base.py

@@ -0,0 +1,153 @@
+import cytoolz as toolz
+
+from recurvedata.pigeon.meta import HandlerFactoryMeta, HandlerMeta
+from recurvedata.pigeon.transformer import Transformer
+from recurvedata.pigeon.utils import LoggingMixin
+
+_default_transformer = Transformer()
+
+
+class Handler(LoggingMixin):
+    ERROR_HANDLE_PARAMS = [
+        "max_error_rate",
+        "min_sample_rows",
+        "check_error_rate_on_finish",
+        "max_continuous_error_log_rows",
+        "max_continuous_error_log_size",
+        "error_log_cycle_interval",
+    ]
+
+    def __init__(
+        self,
+        transformer=_default_transformer,
+        max_error_rate=0.2,
+        min_sample_rows=1000,
+        check_error_rate_on_finish=True,
+        max_continuous_error_log_rows=10,
+        max_continuous_error_log_size=500000,
+        error_log_cycle_interval=100,
+        *args,
+        **kwargs,
+    ):
+        """记录处理逻辑，接收一行输入，调用转换逻辑，输出处理结果。
+        :param transformer: 定义数据处理逻辑，Transformer 或子类对象
+        :param max_error_rate: 最大错误率，超出该错误率可能会（结合样本大小）抛出异常中断程序
+        :param min_sample_rows: 最小采样数量，避免样本太小导致误报
+        :param check_error_rate_on_finish: 在程序结束时检查错误率，达到阈值会抛出异常，避免因为数据量太小达不到样本容量而淹没错误。
+        :param max_continuous_error_log_rows: 单个 handler 连续打印的报错 row 的最大行数，防止错误日志太多 log 太大
+        :param max_continuous_error_log_size: 单个 handler 连续打印的报错 row 字符串的最大字符数，默认 50w, 防止错误日志太多 log 太大
+        "param error_log_cycle_interval: 当达到单个 handler 连续打印最大行数或者最大字符数后，仍然间隔 error_log_cycle_interval 打印一次报错日志
+        """
+        self.transformer = transformer
+        self.max_error_rate = max_error_rate
+        self.min_sample_rows = min_sample_rows
+        self.check_error_rate_on_finish = check_error_rate_on_finish
+        self.max_continuous_error_log_rows = max_continuous_error_log_rows
+        self.max_continuous_error_log_size = max_continuous_error_log_size
+        self.error_log_cycle_interval = error_log_cycle_interval
+
+        self.meta = HandlerMeta()
+        self.meta.schema = self.transformer.output_schema
+
+    def set_transformer(self, transformer):
+        self.transformer = transformer
+
+    def transform(self, row):
+        return self.transformer.transform(row)
+
+    def set_input_schema(self, schema):
+        self.transformer.input_schema = schema
+
+    def close(self):
+        if self.check_error_rate_on_finish:
+            self.check_error_rate()
+
+    def emit(self, row):
+        raise NotImplementedError("emit must be implemented by subclasses")
+
+    def handle(self, row):
+        self.meta.num_input_rows += 1
+        try:
+            rv = self.transform(row)
+            if rv:
+                num_rows = self.emit(rv)
+                self.meta.num_output_rows += num_rows
+        except (KeyboardInterrupt, SystemExit):
+            raise
+        except Exception:
+            self.handle_error(row)
+
+        if self.meta.num_input_rows >= self.min_sample_rows:
+            self.check_error_rate()
+
+    @property
+    def error_rate(self):
+        if self.meta.num_input_rows == 0:
+            return 0
+        return self.meta.num_error_rows / self.meta.num_input_rows
+
+    def check_error_rate(self):
+        if self.error_rate > self.max_error_rate:
+            raise RuntimeError(
+                f"max_error_rate reached,"
+                f" #input: {self.meta.num_input_rows},"
+                f" #error: {self.meta.num_error_rows},"
+                f" error_rate: {self.error_rate},"
+                f" threshold: {self.max_error_rate}"
+            )
+
+    def handle_error(self, row):
+        self.meta.num_error_rows += 1
+        self.meta.error_log_size += len(str(row))
+        if (
+            self.meta.num_error_rows <= self.max_continuous_error_log_rows
+            and self.meta.error_log_size <= self.max_continuous_error_log_size
+        ):
+            self.logger.exception("failed to handle row: %s", row)
+        elif self.meta.num_error_rows % self.error_log_cycle_interval == 0:
+            self.logger.exception(
+                "current handler total %s error rows, failed to handle row: %s", self.meta.num_error_rows, row
+            )
+
+    def reset(self):
+        """reset all states"""
+        self.meta.reset()
+
+
+class NullHandler(Handler):
+    def transform(self, row):
+        pass
+
+    def emit(self, row):
+        pass
+
+    def handle(self, row):
+        return 0
+
+
+class HandlerFactory(LoggingMixin):
+    def __init__(self, handler_class, transformer=_default_transformer, **handler_options):
+        self.handler_class = handler_class
+        self.transformer = transformer
+        self.handler_options = handler_options
+        self.handlers = []
+
+        self.meta = HandlerFactoryMeta(name=self.meta_name())
+
+    def set_transformer(self, transformer):
+        self.transformer = transformer
+
+    def create_handler(self, **kwargs):
+        h = self.handler_class(**toolz.merge(self.handler_options, kwargs))
+        h.set_transformer(self.transformer)
+        self.handlers.append(h)
+        return h
+
+    def join(self):
+        pass
+
+    def meta_name(self):
+        return f"<{self.__class__.__name__}>"
+
+    def handle_dumper_schema(self, schema):
+        pass

recurvedata/pigeon/handler/csv_handler.py

@@ -0,0 +1,290 @@
+import csv
+import glob
+import logging
+import os
+
+import cytoolz as toolz
+
+from recurvedata.pigeon import const
+from recurvedata.pigeon.csv import CSV
+from recurvedata.pigeon.handler.base import Handler, HandlerFactory
+from recurvedata.pigeon.row_factory import get_row_keys
+from recurvedata.pigeon.schema import Schema
+from recurvedata.pigeon.utils import escape, extract_dict, fs
+from recurvedata.pigeon.utils.json import json_dumps
+
+logger = logging.getLogger(__name__)
+
+
+class CSVFileHandler(Handler):
+    def __init__(
+        self,
+        filename=None,
+        encoding=None,
+        write_header=False,
+        null=None,
+        delimiter=",",
+        quoting=csv.QUOTE_ALL,
+        escapechar=None,
+        doublequote=True,
+        **kwargs,
+    ):
+        """Transform rows and write result as CSV file.
+
+        :param filename: result filename
+        :param encoding: encoding
+        :param write_header: writer headers or not
+        """
+        if filename is None:
+            filename = fs.new_tempfile(suffix=".csv")
+        self.filename = os.path.abspath(filename)
+        if os.path.exists(self.filename):
+            logger.info(f"found {filename}, remove before writing")
+            os.unlink(self.filename)
+
+        self.encoding = encoding
+        self.null = null
+        self._fd = None
+        self._writer = None
+        self.writer_header = write_header
+        self._field_names = None
+
+        self.csv_options = {k: v for k, v in kwargs.items() if k not in self.ERROR_HANDLE_PARAMS}
+        self.csv_options.update(
+            {"delimiter": delimiter, "quoting": quoting, "escapechar": escapechar, "doublequote": doublequote}
+        )
+        super().__init__(**extract_dict(kwargs, self.ERROR_HANDLE_PARAMS))
+
+    def flush(self):
+        if self._fd is not None:
+            self._fd.flush()
+
+    def close(self):
+        super().close()
+
+        if self._fd is not None:
+            self._fd.close()
+
+    def reset(self):
+        super().reset()
+        if self._fd is not None:
+            self._fd.seek(0)
+            self._fd.truncate(0)
+            self._fd.close()
+            self._fd = self._writer = None
+
+    def _open_writer(self, row):
+        self._fd = open(self.filename, "w", newline="", encoding=self.encoding)
+
+        self._determine_header(row)
+
+        if isinstance(row, dict):
+            self._writer = csv.DictWriter(self._fd, fieldnames=self._field_names, **self.csv_options)
+            if self.writer_header:
+                self._writer.writeheader()
+        else:
+            self._writer = csv.writer(self._fd, **self.csv_options)
+            if self.writer_header:
+                self._writer.writerow(self._field_names)
+
+    def _determine_header(self, row):
+        logger.info("try to get schema from row (%s)", type(row))
+        field_names = get_row_keys(row)
+        if not field_names:
+            logger.info("try to get schema from transformer")
+            if self.transformer.output_schema is not None:
+                field_names = [x.name for x in self.transformer.output_schema]
+        logger.info("header: %s", field_names)
+        self._field_names = field_names
+
+    def emit(self, row):
+        if not isinstance(row, list):
+            row = [row]
+
+        if self._fd is None:
+            self._open_writer(row[0])
+
+        for r in row:
+            self._writerow(r)
+
+        return len(row)
+
+    def _writerow(self, row):
+        if isinstance(row, dict):
+            row = toolz.valmap(self._escape_item, row)
+        else:
+            row = [self._escape_item(x) for x in row]
+        self._writer.writerow(row)
+
+    def _escape_item(self, v):
+        if v is None:
+            return self.null
+
+        # Handle dict, tuple, set and list
+        if isinstance(v, (dict, tuple, set, list)):
+            v = json_dumps(v, ensure_ascii=False)
+
+        if isinstance(v, str):
+            return escape.escape_string(v)
+        return v
+
+    def __str__(self):
+        return "<{} (filename={})>".format(self.__class__.__name__, self.filename)
+
+
+class HiveCSVFileHandler(CSVFileHandler):
+    """
+    The default file format of Hive is not CSV, but only delimiter-ed text file.
+    """
+
+    def __init__(
+        self,
+        filename=None,
+        encoding=None,
+        write_header=False,
+        null=const.HIVE_NULL,
+        delimiter=const.HIVE_FIELD_DELIMITER,
+        quoting=csv.QUOTE_NONE,
+        **csv_options,
+    ):
+        super().__init__(filename, encoding, write_header, null, delimiter, quoting, **csv_options)
+        self.delimiter = delimiter
+
+    def _writerow(self, row):
+        if isinstance(row, dict):
+            line = self.format_line(row.values())
+        else:
+            line = self.format_line(row)
+
+        self._fd.write(line)
+        self._fd.write("\n")
+
+    def _escape_item(self, v):
+        v = super()._escape_item(v)
+        return str(v)
+
+    def format_line(self, row):
+        values = map(self._escape_item, row)
+        return self.delimiter.join(values)
+
+
+class CSVFileHandlerFactory(HandlerFactory):
+    def __init__(
+        self, handler_class=CSVFileHandler, filename=None, encoding=None, write_header=False, merge_files=True, **kwargs
+    ):
+        self.filename = filename or fs.new_tempfile(".csv")
+        self.merge_files = merge_files
+        options = {"filename": self.filename, "encoding": encoding, "write_header": write_header}
+        options.update(kwargs)
+        super().__init__(handler_class=handler_class, **options)
+
+        self._saved_schema = False
+
+    def meta_name(self):
+        return f"<{self.__class__.__name__} ({self.filename})>"
+
+    def create_handler(self, **kwargs):
+        filename = f'{self.handler_options["filename"]}.{len(self.handlers)}'
+        return super().create_handler(filename=filename)
+
+    def join(self):
+        if not self.handlers:
+            return
+
+        files = [h.filename for h in self.handlers if not fs.is_file_empty(h.filename)]
+        if not files:
+            return
+
+        if self.merge_files:
+            has_header = self.handler_options.get("write_header", False)
+            logger.info("files have header: %s", has_header)
+            logger.info("merging files %s into %s", files, self.filename)
+            if not has_header:
+                fs.merge_files(files, self.filename)
+            else:
+                base_file = files[0]  # keep header
+                target = fs.merge_files(files[1:], num_skip_lines=1)
+                fs.merge_files([base_file, target], self.filename)
+
+        self.save_output_schema()
+
+    def _determine_output_schema(self):
+        if self.transformer.output_schema is not None:
+            # list of Fields
+            if isinstance(self.transformer.output_schema, list):
+                return Schema(self.transformer.output_schema)
+            return self.transformer.output_schema
+
+        if self.handler_options.get("write_header", False):
+            csv_options = self.handlers[0].csv_options
+            if self.merge_files:
+                f = self.filename
+            else:
+                # use the first non-empty file
+                f = toolz.first(x for x in glob.glob(f"{self.filename}.[0-9]*") if os.path.getsize(x))
+            csv_proxy = CSV(path=f, encoding=self.handler_options["encoding"], **csv_options)
+            return csv_proxy.infer_schema()
+        return None
+
+    def save_output_schema(self):
+        schema = self._determine_output_schema()
+        if not schema:
+            logger.warning("could not able to infer output schema")
+            return
+
+        filename = fs.schema_filename(self.filename)
+        logger.info("saving output schema to %s", filename)
+        schema.dump(filename)
+        self._saved_schema = True
+        return filename
+
+    def handle_dumper_schema(self, schema):
+        filename = fs.schema_filename(self.filename)
+        if self._saved_schema and os.path.exists(filename):
+            logger.info("file %s already exists, pass", filename)
+            return
+
+        if not isinstance(schema, Schema):
+            raise TypeError(f"got {type(schema)}")
+
+        logger.info("saving dumper schema to %s", filename)
+        schema.dump(filename)
+        return filename
+
+
+def create_csv_file_handler_factory(
+    filename=None, encoding=None, write_header=False, hive=False, transformer=None, merge_files=True, **kwargs
+):
+    if hive:
+        handler_class = HiveCSVFileHandler
+    else:
+        handler_class = CSVFileHandler
+
+    factory = CSVFileHandlerFactory(
+        handler_class=handler_class,
+        filename=filename,
+        encoding=encoding,
+        write_header=write_header,
+        merge_files=merge_files,
+        **kwargs,
+    )
+    if transformer is not None:
+        factory.set_transformer(transformer)
+    return factory
+
+
+def convert_csv_to_hive_textfile(filename, folder=None, replace=False, has_header=False, **csv_options):
+    new_name = fs.new_tempfile(prefix=os.path.basename(filename), dir=folder)
+    handler = HiveCSVFileHandler(filename=new_name)
+    with open(filename, newline="") as fd:
+        if has_header:
+            fd.readline()
+
+        reader = csv.reader(fd, **csv_options)
+        for row in reader:
+            handler.handle(tuple(row))
+    handler.close()
+
+    if replace:
+        os.rename(new_name, filename)
+    return new_name

recurvedata/pigeon/loader/__init__.py

@@ -0,0 +1,87 @@
+from recurvedata.pigeon.connector import (
+    new_azure_synapse_connector,
+    new_elasticsearch_connector,
+    new_google_bigquery_connector,
+    new_postgresql_connector,
+    new_redshift_connector,
+    new_tidb_connector,
+)
+from recurvedata.pigeon.loader.csv_to_azure_synapse import CSVToAzureSynapseLoader
+from recurvedata.pigeon.loader.csv_to_clickhouse import CSVToClickHouseLoader
+from recurvedata.pigeon.loader.csv_to_es import CSVToElasticSearchLoader
+from recurvedata.pigeon.loader.csv_to_google_bigquery import CSVToGoogleBigqueryLoader
+from recurvedata.pigeon.loader.csv_to_hive import CSVToHiveLoader
+from recurvedata.pigeon.loader.csv_to_mysql import CSVToMySQLLoader
+from recurvedata.pigeon.loader.csv_to_postgresql import CSVToPostgresqlLoader
+from recurvedata.pigeon.loader.csv_to_redshift import CSVToRedshiftLoader
+from recurvedata.pigeon.loader.csv_to_starrocks import CSVToStarRocksLoader
+
+
+def new_csv_to_hive_loader(table, filename, database, **kwargs):
+    params = kwargs.copy()
+    params.update(table=table, filename=filename, database=database)
+    loader = CSVToHiveLoader(**params)
+    return loader
+
+
+def new_csv_to_mysql_loader(table, filename, database, **kwargs):
+    params = kwargs.copy()
+    params.update(table=table, filename=filename, database=database)
+    loader = CSVToMySQLLoader(**params)
+    return loader
+
+
+def new_csv_to_starrocks_loader(table, filename, database, **kwargs):
+    params = kwargs.copy()
+    params.update(table=table, filename=filename, database=database)
+    loader = CSVToStarRocksLoader(**params)
+    return loader
+
+
+def new_csv_to_tidb_loader(table, filename, database, **kwargs):
+    params = kwargs.copy()
+    params.update(table=table, filename=filename, database=database)
+    loader = CSVToMySQLLoader(**params)
+    return loader
+
+
+def new_csv_to_redshift_loader(table, filename, database, **kwargs):
+    params = kwargs.copy()
+    params.update(table=table, filename=filename, database=database)
+    loader = CSVToRedshiftLoader(**params)
+    return loader
+
+
+def new_csv_to_postgresql_loader(table, filename, database, **kwargs):
+    params = kwargs.copy()
+    params.update(table=table, filename=filename, database=database)
+    loader = CSVToPostgresqlLoader(**params)
+    return loader
+
+
+def new_csv_to_azure_synapse_loader(table, filename, **kwargs):
+    params = kwargs.copy()
+    params.update(table=table, filename=filename)
+    loader = CSVToAzureSynapseLoader(**params)
+    return loader
+
+
+def new_csv_to_clickhouse_loader(table, filename, database, **kwargs):
+    params = kwargs.copy()
+    params.update(table=table, filename=filename, database=database)
+    loader = CSVToClickHouseLoader(**params)
+    return loader
+
+
+def new_csv_to_elasticsearch_loader(index, filename, **kwargs):
+    params = kwargs.copy()
+    params.update(index=index, filename=filename)
+    loader = CSVToElasticSearchLoader(**params)
+    return loader
+
+
+def new_csv_to_google_bigquery_loader(table, filename, **kwargs):
+    params = kwargs.copy()
+    params.update(table=table, filename=filename)
+    loader = CSVToGoogleBigqueryLoader(**params)
+    return loader