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

recurvedata/pigeon/sync.py

@@ -0,0 +1,283 @@
+import logging
+import os.path
+from urllib.parse import urlparse
+
+from slugify import slugify
+
+from recurvedata.pigeon.connector import get_connector, get_connector_class
+from recurvedata.pigeon.const import LOAD_APPEND, LOAD_MERGE
+from recurvedata.pigeon.handler.csv_handler import create_csv_file_handler_factory
+from recurvedata.pigeon.loader import CSVToHiveLoader, CSVToMySQLLoader, CSVToRedshiftLoader
+from recurvedata.pigeon.utils import ensure_list, fs
+
+logger = logging.getLogger(__name__)
+
+
+def _parse_db_table(db_table):
+    t_db_table = db_table.split(".")
+    if len(t_db_table) == 2:
+        db, table = t_db_table
+        schema = None
+    elif len(t_db_table) == 3:
+        db, schema, table = t_db_table
+    else:
+        raise ValueError(f"Invalid database and table {db_table!r}")
+    return {
+        "database": db,
+        "schema": schema,
+        "table": table,
+    }
+
+
+class Location:
+    """
+    Syntax:
+        完整URL形式：{protocol}://{user}:{password}@{host}:{port}/{database}.{table}
+        简写形式：   {dbconf}:{database}.{table}
+        本地文件：   file://{path}
+        本地文件简写：file:{path}
+    Example:
+        mysql://dev:pass@172.16.24.93:3306/testdb.test
+        tidb:testdb.test
+        file:///tmp/result.csv
+        file:/tmp/result.csv
+    """
+
+    ATTRS = ["protocol", "user", "password", "host", "port", "database", "schema", "table", "dbconf", "path"]
+
+    def __init__(self, spec=None, **kwargs):
+        if not kwargs:
+            if not spec:
+                raise ValueError("Location spec is required")
+            if spec.startswith("file:"):
+                self._from_file(spec)
+            elif "://" in spec:
+                self._from_url(spec)
+            else:
+                self._from_simple(spec)
+        else:
+            if spec is not None:
+                kwargs["dbconf"] = spec
+            self._from_dict(kwargs)
+
+    def __str__(self):
+        if self.protocol == "file":
+            return self._format_file()
+        if self.dbconf:
+            return self._format_simple()
+        else:
+            return self._format_url()
+
+    def _format_file(self):
+        return f"{self.protocol}:{self.path}"
+
+    def _format_url(self, hide_password=True):
+        ret = [f"{self.protocol}://"]
+        if self.user:
+            ret.append(str(self.user))
+            if self.password:
+                password = "***" if hide_password else self.passsword
+                ret.append(f":{password}")
+            ret.append("@")
+        if self.host:
+            if ":" in self.host:
+                ret.append(f"[{self.host}]")
+            else:
+                ret.append(str(self.host))
+            if self.port:
+                ret.append(f":{self.port}")
+        if self.database:
+            ret.append(f"/{self.database}")
+            if self.schema:
+                ret.append(f".{self.schema}")
+            if self.table:
+                ret.append(f".{self.table}")
+        return "".join(ret)
+
+    def _format_simple(self):
+        ret = [f"{self.dbconf}:"]
+        if self.database:
+            ret.append(f"{self.database}")
+            if self.schema:
+                ret.append(f".{self.schema}")
+            if self.table:
+                ret.append(f".{self.table}")
+        return "".join(ret)
+
+    def __repr__(self):
+        return f"<{type(self).__name__} {str(self)}>"
+
+    def _from_file(self, url):
+        """
+        Syntax:
+            file://{path}
+            file:{path}
+        """
+        protocol, path = url.split(":", maxsplit=1)
+        if path.startswith("//"):
+            path = path[2:]
+        params = {"protocol": protocol, "path": path}
+        self._from_dict(params)
+
+    def _from_url(self, url):
+        """
+        Syntax:
+            {protocol}://{user}:{password}@{host}:{port}/{database}.{table}
+        """
+        parsed = urlparse(url)
+        params = {
+            "protocol": parsed.scheme,
+            "user": parsed.username,
+            "password": parsed.password,
+            "host": parsed.hostname,
+            "port": parsed.port,
+        }
+        params.update(_parse_db_table(parsed.path.strip("/")))
+        self._from_dict(params)
+
+    def _from_simple(self, spec):
+        """
+        Syntax:
+            {dbconf}:{database}.{table}
+        """
+        dbconf, db_table = spec.strip().split(":")
+        params = {"dbconf": dbconf}
+        params.update(_parse_db_table(db_table))
+        self._from_dict(params)
+
+    def _from_dict(self, params):
+        unknown_params = set(params) - set(self.ATTRS)
+        if unknown_params:
+            raise ValueError(f"Unknown params {unknown_params}")
+        if not params.get("protocol") and not params.get("dbconf"):
+            raise ValueError("protocol or dbconf is required")
+        if params.get("protocol") == "file" and not params.get("path"):
+            raise ValueError("path is required")
+        for k in self.ATTRS:
+            setattr(self, k, params.get(k))
+
+    def to_dict(self):
+        ret = {}
+        for k in self.ATTRS:
+            v = getattr(self, k, None)
+            if v is not None and v != "":
+                ret[k] = v
+        return ret
+
+    @property
+    def is_local(self):
+        return self.dbconf == "file" or self.protocol == "file"
+
+
+def _get_dumper_cls(dtype):
+    # XXX: so ugly here
+    if dtype in ["cassandra"]:
+        from recurvedata.pigeon.dumper.cass import CassandraDumper
+
+        return CassandraDumper
+
+    from recurvedata.pigeon.dumper.dbapi import DBAPIDumper
+
+    return DBAPIDumper
+
+
+def _get_connector(location):
+    if location.protocol:
+        # only support dbapi
+        _conn_cls = get_connector_class(location.protocol)
+        connector = _conn_cls(
+            host=location.host,
+            port=location.port,
+            user=location.user,
+            password=location.passsword,
+            database=location.database,
+        )
+    else:
+        connector = get_connector(location.dbconf, database=location.database)
+    return connector
+
+
+def _dump(src, handler_factory):
+    connector = _get_connector(src)
+    if src.schema:
+        table = f"{src.schema}.{src.table}"
+    else:
+        table = src.table
+    dumper = _get_dumper_cls(src.protocol or src.dbconf)(
+        connector,
+        table=table,
+        handler_factories=[handler_factory],
+    )
+    logger.info("Dump start".center(40, "="))
+    dumper.execute()
+
+
+_loader_config = {
+    "hive": {
+        "cls": CSVToHiveLoader,
+        "connector": "hive_connector",
+    },
+    "redshift": {
+        "cls": CSVToRedshiftLoader,
+        "connector": "redshift_connector",
+    },
+    "mysql": {
+        "cls": CSVToMySQLLoader,
+        "connector": "connector",
+    },
+}
+_loader_config["tidb"] = _loader_config["mysql"]
+
+
+def _load(dst, filename, mode, merge_keys=()):
+    connector = _get_connector(dst)
+
+    cfg = _loader_config[dst.protocol or dst.dbconf]
+    cls, connector_name = cfg["cls"], cfg["connector"]
+    kwargs = {
+        "database": dst.database,
+        "table": dst.table,
+        "filename": filename,
+        connector_name: connector,
+    }
+    if mode.upper() == LOAD_MERGE:
+        logger.info(f"Primary keys: {merge_keys} in {dst}")
+        kwargs.update({"mode": LOAD_MERGE, "primary_keys": ensure_list(merge_keys)})
+    elif mode.upper() == LOAD_APPEND:
+        kwargs.update({"mode": LOAD_APPEND})
+    loader = cls(**kwargs)
+    logger.info("Load start".center(40, "="))
+    loader.execute()
+
+
+def _get_stage_filename(src, dst):
+    if src.is_local:
+        return src.path
+    if dst.is_local:
+        if not os.path.isabs(dst.path):
+            return os.path.abspath(dst.path)
+        return dst.path
+    tmpdir = f"{src.protocol or src.dbconf}_to_{dst.protocol or dst.dbconf}"
+    new_stagefile = fs.new_stagefile_factory(tmpdir)
+    return new_stagefile(slugify(f"{src}_to_{dst}") + ".txt")
+
+
+def sync(src, dst, mode, merge_keys=()):
+    """同步一个表"""
+    if not isinstance(src, Location):
+        src = Location(src)
+    if not isinstance(dst, Location):
+        dst = Location(dst)
+
+    if dst.protocol and dst.protocol != "file":
+        raise NotImplementedError("暂不支持URL形式的目标")
+
+    filename = _get_stage_filename(src, dst)
+    logger.info(f"Dump to file: {filename}")
+
+    if not src.is_local:
+        for_hive = (dst.protocol or dst.dbconf) in ["impala", "hive"]
+        handler_factory = create_csv_file_handler_factory(filename=filename, hive=for_hive)
+        _dump(src, handler_factory)
+    if not dst.is_local:
+        _load(dst, filename, mode=mode, merge_keys=merge_keys)

recurvedata/pigeon/transformer.py

@@ -0,0 +1,146 @@
+import struct
+import zlib
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+try:
+    # use ujson for better performance
+    import ujson as json
+except ImportError:
+    import json
+
+from recurvedata.pigeon import const
+from recurvedata.pigeon.schema import Schema
+
+_Row = Union[Tuple, Dict[str, Any]]
+
+
+class Transformer:
+    _input_schema: Optional[Schema] = None
+    _use_input_schema_as_output: bool = False
+
+    @property
+    def input_schema(self) -> Optional[Schema]:
+        """Returns the schema of input data"""
+        return self._input_schema
+
+    @input_schema.setter
+    def input_schema(self, schema: Schema):
+        """Should be called by the handler"""
+        assert isinstance(schema, Schema)
+        self._input_schema = schema
+
+    @property
+    def output_schema(self) -> Optional[Schema]:
+        """Subclasses that change the rows schema should provide the output schema.
+
+        These operations will change the output schema:
+        - Add or remove fields
+        - Change the name of fields
+        - Change the type of fields
+
+        An example of valid schema:
+
+        from recurvedata.pigeon.schema import Schema, Field, types
+
+        Schema([
+            Field(name='id', type=types.INT32),
+            Field(name='name', type=types.STRING, size=64),
+            Field(name='snapshot_time', type=types.DATETIME, comment='snapshot_time in UTC'),
+            Field(name='is_active', type=types.BOOLEAN)
+        ])
+
+        Allowed types:
+
+        - 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'
+        """
+        if self._use_input_schema_as_output:
+            return self._input_schema
+        return None
+
+    def transform(self, row: _Row, *args, **kwargs) -> Union[_Row, List[_Row]]:
+        """This is the method called by Handler.
+
+        It internally calls `transform_impl` to do the real transform logic.
+        Subclasses should implement `transform_impl` but not this method.
+
+        :param row: a Row (namedtuple) object contains a row record fetched from database
+        :returns: returns one (tuple) or multiple (list of tuple) rows
+        """
+        return self.transform_impl(row, *args, **kwargs)
+
+    def transform_impl(self, row: _Row, *args, **kwargs) -> Union[_Row, List[_Row]]:
+        return row
+
+    @staticmethod
+    def convert_json_to_hive_map(data: Union[str, bytes]) -> str:
+        if not data:
+            return const.HIVE_NULL
+
+        d = json.loads(data)
+        items = []
+        for key, value in d.items():
+            key = str(key).strip()
+            value = str(value).strip()
+            item = '{0}{1}{2}'.format(key, const.HIVE_MAP_KV_DELIMITER, value)
+            items.append(item)
+        return const.HIVE_MAP_ITEM_DELIMITER.join(items)
+
+    @staticmethod
+    def convert_json_to_hive_array(data: Union[str, bytes]) -> str:
+        if not data:
+            return const.HIVE_NULL
+
+        items = json.loads(data)
+        return const.HIVE_ARRAY_DELIMITER.join(items)
+
+    @staticmethod
+    def mysql_uncompress(value: bytes, return_str: bool = False) -> Union[bytes, str]:
+        """A Python implementation of UNCOMPRESS function of MySQL.
+
+        Used to decompress result of COMPRESS function.
+
+        https://dev.mysql.com/doc/refman/5.7/en/encryption-functions.html#function_compress
+
+        :param value: the compressed data in bytes
+        :param return_str: the return value should be unicode
+        :type return_str: bool
+        :rtype: bytes | str
+        """
+
+        # Empty strings are stored as empty strings.
+        # Nonempty strings are stored as a 4-byte length of the uncompressed string
+        if not value or len(value) < 4:
+            return value
+
+        rv = zlib.decompress(value[4:])
+
+        if return_str:
+            rv = rv.decode()
+        return rv
+
+    @staticmethod
+    def mysql_compress(value: Optional[str]) -> Optional[bytes]:
+        if value is None:
+            return None
+        if value == '':
+            return b''
+        size = struct.pack('I', len(value))
+        data = zlib.compress(value.encode())
+        return size + data
+
+    @staticmethod
+    def json_loads(*args, **kwargs) -> Any:
+        return json.loads(*args, **kwargs)
+
+    @staticmethod
+    def json_dumps(*args, **kwargs) -> str:
+        return json.dumps(*args, **kwargs)

recurvedata/pigeon/utils/__init__.py

@@ -0,0 +1,134 @@
+import hashlib
+import logging
+import time
+from contextlib import contextmanager
+from typing import Dict, List, Set, Tuple, TypeVar, Union
+from uuid import uuid4
+
+import cytoolz as toolz
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T")
+
+
+def safe_int(v: Union[str, int, float], default: int = 0) -> int:
+    try:
+        return int(v)
+    except Exception:
+        return default
+
+
+def trim_prefix(s: str, sub: str) -> str:
+    if not s.startswith(sub):
+        return s
+    return s[len(sub) :]
+
+
+def trim_suffix(s: str, sub: str) -> str:
+    if not s.endswith(sub):
+        return s
+    return s[: -len(sub)]
+
+
+class LoggingMixin(object):
+    @property
+    def logger(self) -> logging.Logger:
+        try:
+            return self._logger
+        except AttributeError:
+            self._logger = logging.root.getChild(self.__class__.__module__ + "." + self.__class__.__name__)
+            return self._logger
+
+
+def init_logging(
+    level_name="info",
+    fmt="%(asctime)s - %(levelname)s - %(name)s - %(filename)s:%(lineno)d - [%(process)d:%(threadName)s] - %(message)s",
+    silent_cassandra=True,
+):
+    level = logging.INFO
+    if level_name == "info":
+        level = logging.INFO
+    elif level_name == "warning":
+        level = logging.WARNING
+    elif level_name == "error":
+        level = logging.ERROR
+    elif level_name == "debug":
+        level = logging.DEBUG
+    logging.basicConfig(level=level, format=fmt)
+
+    if silent_cassandra:
+        # cassandra is too noisy
+        logging.getLogger("cassandra.cluster").setLevel(logging.WARNING)
+
+
+def ensure_list(v: Union[T, Tuple[T], List[T], Set[T]]) -> List[T]:
+    if isinstance(v, (tuple, set, list)):
+        return list(v)
+    return [v]
+
+
+def ensure_str_list(v: str, sep: str = ",", strip: bool = True) -> List[str]:
+    if v is None:
+        return []
+
+    if isinstance(v, str):
+        if not v:
+            return []
+        if strip:
+            return [x.strip() for x in v.split(sep)]
+        else:
+            return v.split(sep)
+
+    if isinstance(v, (tuple, set, list)):
+        return list(v)
+    raise TypeError(f'unsupported type "{type(v)}"')
+
+
+def ensure_query_list(v: Union[str, List[str]]) -> List[str]:
+    if not v:
+        return []
+    if isinstance(v, list):
+        return v
+    return list(filter(None, map(lambda x: x.strip(), v.split(";"))))
+
+
+def extract_dict(d: Dict, keys: List) -> Dict:
+    return {k: v for k, v in d.items() if k in keys}
+
+
+@contextmanager
+def silent(*_excs, excs=None):
+    excs = excs or _excs or (Exception,)
+    try:
+        yield
+    except excs as e:
+        logging.exception("silent %s", type(e).__name__)
+
+
+def replace_null_values(row: Union[List, Tuple, Dict], null_values: List, replacer=None):
+    def _f(v):
+        if v in null_values:
+            return replacer
+        return v
+
+    if isinstance(row, list):
+        return list(map(_f, row))
+    if isinstance(row, tuple):
+        return tuple(map(_f, row))
+    if isinstance(row, dict):
+        return toolz.valmap(_f, row)
+    raise TypeError(f"only list, tuple or dict type is supported, got {repr(type(row))}")
+
+
+def md5hash(v: Union[str, bytes]) -> str:
+    if isinstance(v, str):
+        v = v.encode()
+    if not isinstance(v, bytes):
+        v = str(v).encode()
+    return hashlib.md5(v).hexdigest()
+
+
+def randomized_suffix() -> str:
+    pure_time_str = str(time.time()).replace(".", "")
+    return pure_time_str[-1] + uuid4().hex[:6]