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

recurvedata/pigeon/connector/doris.py

@@ -0,0 +1,166 @@
+import json
+import subprocess
+import time
+
+from recurvedata.pigeon.connector._registry import register_connector_class
+from recurvedata.pigeon.connector.mysql import MySQLConnector
+from recurvedata.pigeon.schema import types
+
+_canonical_type_to_doris_type = {
+    types.BOOLEAN: "TINYINT",
+    types.INT8: "TINYINT",
+    types.INT16: "SMALLINT",
+    types.INT32: "INT",
+    types.INT64: "BIGINT",
+    types.FLOAT32: "FLOAT",
+    types.FLOAT64: "DOUBLE",
+    types.DATE: "DATE",
+    types.DATETIME: "DATETIME",
+    types.STRING: "STRING",
+    types.JSON: "STRING",
+}
+
+
+@register_connector_class(["doris"])
+class DorisConnector(MySQLConnector):
+    _sqla_driver = "doris+pymysql"
+    _default_port = 9030
+    _default_fe_http_port = 8030
+
+    def __init__(self, host, port=None, http_port=None, database=None, user=None, password=None, *args, **kwargs):
+        self.http_port = http_port or self._default_fe_http_port
+        super().__init__(host=host, port=port, database=database, user=user, password=password, *args, **kwargs)
+
+    @property
+    def load_strict_mode(self):
+        if not hasattr(self, "_load_strict_mode"):
+            return False
+        return self._load_strict_mode
+
+    @load_strict_mode.setter
+    def load_strict_mode(self, mode: bool):
+        self._load_strict_mode = mode
+
+    @property
+    def max_filter_ratio(self):
+        if not hasattr(self, "_max_filter_ratio"):
+            return 0
+        return self._max_filter_ratio
+
+    def has_table(self, table, database=None, cursor=None, **kwargs):
+        retry_num = 6
+        for attempt in range(retry_num):
+            if super().has_table(table, database, cursor, **kwargs):
+                return True
+            if attempt < retry_num - 1:
+                wait_time = (attempt + 1) ** 2
+                time.sleep(wait_time)  # wait for table to be created and visible
+        return False
+
+    @max_filter_ratio.setter
+    def max_filter_ratio(self, ratio: float):
+        if ratio < 0:
+            self._max_filter_ratio = 0
+        elif ratio > 1:
+            self._max_filter_ratio = 1
+        else:
+            self._max_filter_ratio = ratio
+
+    def _load_csv_mysql(
+        self,
+        table,
+        filename,
+        columns=None,
+        delimiter=",",
+        quotechar='"',
+        lineterminator="\r\n",
+        escapechar=None,
+        skiprows=0,
+        **kwargs,
+    ):
+        """
+        stream load data from csv file into table
+        """
+
+        def _split_database_table(table_name: str):
+            tmp_lst = table_name.split(".")
+            if len(tmp_lst) == 1:
+                return self.database, table_name
+            return tmp_lst
+
+        db_name, table_name = _split_database_table(table)
+        shell_cmd = self._format_load_shell(filename, db_name, table_name)
+
+        # Set only authentication environment variables
+        _env = {}
+        if self.user is not None:
+            _env["DORIS_USER"] = self.user
+        if self.password is not None:
+            _env["DORIS_PASSWORD"] = self.password
+
+        output = subprocess.check_output(shell_cmd, env=_env, shell=True)
+        self._log(output)
+        res_txt = output.decode()
+        res = json.loads(res_txt)
+        self._log(res_txt)
+
+        if res["Status"] != "Success":
+            if "ErrorURL" not in res:
+                err_output = res["Message"]
+            else:
+                err_url = res["ErrorURL"]
+                err_output = subprocess.check_output(["curl", err_url])
+            self._log(f"error: {err_output}")
+            raise Exception("load csv failed")
+
+    def _format_load_shell(self, filename: str, db_name: str, table_name: str) -> str:
+        """Format the curl command for Doris stream load.
+
+        Args:
+            filename: Path to the CSV file to load
+            db_name: Target database name
+            table_name: Target table name
+
+        Returns:
+            Formatted curl command string for stream loading data
+        """
+
+        def __format_column_mapping(db_name: str, table_name: str) -> str:
+            columns = self.get_columns(table_name, db_name)
+            cols_txt = ",".join(columns)
+            return f"columns: {cols_txt}"
+
+        def __format_stream_load_url(db_name: str, table_name: str) -> str:
+            return f"http://{self.host}:{self.http_port}/api/{db_name}/{table_name}/_stream_load"
+
+        # Clean table and db names
+        db_name = db_name.strip("`")
+        table_name = table_name.strip("`")
+
+        # Build command components
+        url = __format_stream_load_url(db_name, table_name)
+        strict_mode = "true" if self.load_strict_mode else "false"
+        column_mapping = __format_column_mapping(db_name, table_name)
+
+        # Construct the full curl command with properly escaped quotes
+        return (
+            f"curl --location-trusted -u $DORIS_USER:$DORIS_PASSWORD "
+            f'-H "Expect:100-continue" '
+            f'-H "max_filter_ratio:{self.max_filter_ratio}" '
+            f'-H "column_separator:," '
+            f'-H "enclose:\\"" '
+            f'-H "trim_double_quotes:true" '
+            f'-H "strict_mode:{strict_mode}" '
+            f'-H "escape:\'" '
+            f'-H "{column_mapping}" '
+            f"-T {filename} -XPUT "
+            f"{url}"
+        ).strip()
+
+    @staticmethod
+    def from_canonical_type(canonical_type, size):
+        if canonical_type == types.STRING:
+            doris_type = "STRING"
+        else:
+            doris_type = _canonical_type_to_doris_type.get(canonical_type, "STRING")
+        return doris_type

recurvedata/pigeon/connector/es.py

@@ -0,0 +1,176 @@
+import base64
+import hashlib
+import pickle
+from collections import defaultdict
+
+from elasticsearch import Elasticsearch, helpers
+from elasticsearch.exceptions import NotFoundError
+
+from recurvedata.pigeon.connector._registry import register_connector_class
+from recurvedata.pigeon.csv import CSV
+from recurvedata.pigeon.schema import Schema, types
+from recurvedata.pigeon.utils import LoggingMixin, ensure_str_list, replace_null_values
+
+# https://www.elastic.co/guide/en/elasticsearch/reference/current/mapping.html#_field_datatypes
+_es_type_to_canonical_type = {
+    "boolean": types.BOOLEAN,
+    "byte": types.INT8,
+    "short": types.INT16,
+    "integer": types.INT32,
+    "long": types.INT64,
+    "half_float": types.FLOAT32,
+    "float": types.FLOAT32,
+    "double": types.FLOAT64,
+    "scaled_float": types.FLOAT64,
+    "date": types.DATETIME,
+    "text": types.STRING,
+    "keyword": types.STRING,
+    "ip": types.STRING,
+    "object": types.STRING,
+    "nested": types.STRING,
+}
+
+_canonical_type_to_es_type = {
+    types.BOOLEAN: "boolean",
+    types.INT8: "byte",
+    types.INT16: "short",
+    types.INT32: "integer",
+    types.INT64: "long",
+    types.FLOAT32: "float",
+    types.FLOAT64: "double",
+    types.DATETIME: "date",
+    types.STRING: "text",
+}
+
+
+@register_connector_class(["es", "elasticsearch"])
+class ElasticSearchConnector(LoggingMixin):
+    def __init__(self, host, **kwargs):
+        self.host = host
+        self._es = Elasticsearch(self.host, **kwargs)
+
+    def scan(self, query=None, index=None, doc_type=None, fields=None, **search_kwargs):
+        if isinstance(query, str):
+            real_query = {"query": {"query_string": {"query": query}}}
+        else:
+            real_query = query
+
+        search_kwargs = search_kwargs.copy()
+        search_kwargs.update({"index": index, "doc_type": doc_type})
+        if fields:
+            search_kwargs["_source_include"] = fields
+        return helpers.scan(self._es, query=real_query, **search_kwargs)
+
+    def get_mapping(self, index, doc_type):
+        try:
+            result = self._es.indices.get_mapping(index=index, doc_type=doc_type)
+        except NotFoundError as e:
+            self.logger.error(str(e))
+            return None
+
+        mappings = list(result.values())[0]["mappings"]
+        if doc_type is not None:
+            properties = mappings[doc_type]["properties"]
+        else:
+            properties = list(mappings.values())[0]["properties"]
+        return properties
+
+    def get_schema(self, index, doc_type):
+        mapping = self.get_mapping(index, doc_type)
+        schema = Schema()
+        for name, attrs in mapping.items():
+            es_type = attrs.get("type", "text").lower()
+            schema.add_field_by_attrs(name, self.to_canonical_type(es_type))
+        return schema
+
+    @staticmethod
+    def get_meta_field_type(name):
+        return {
+            "_index": types.STRING,
+            "_type": types.STRING,
+            "_id": types.STRING,
+            "_score": types.FLOAT64,
+        }[name]
+
+    @staticmethod
+    def to_canonical_type(es_type):
+        return _es_type_to_canonical_type.get(es_type, types.STRING)
+
+    @staticmethod
+    def from_canonical_type(canonical_type):
+        return _canonical_type_to_es_type[canonical_type]
+
+    def load_csv(
+        self,
+        filename,
+        index,
+        doc_type="_doc",
+        schema=None,
+        id_field=None,
+        generate_id=False,
+        null_values=("NULL", r"\N"),
+        null_replacer=None,
+        **csv_options,
+    ):
+        csv_proxy = CSV(filename, **csv_options)
+        if not csv_proxy.has_header:
+            raise ValueError(f"missing header in CSV file {filename}")
+
+        # ensure id fields are present in header
+        if id_field:
+            fields = ensure_str_list(id_field)
+            if not all(x in csv_proxy.header for x in fields):
+                raise ValueError(f"{id_field} is invalid, only {csv_proxy.header} are support")
+        else:
+            fields = None
+
+        if schema is not None:
+            typed_fields = {x.name: x for x in schema.fields}
+        else:
+            typed_fields = {}
+
+        def actions_generator():
+            counters = defaultdict(int)
+            with csv_proxy.reader(as_dict=True) as reader:
+                for doc in reader:
+                    doc = replace_null_values(doc, null_values, null_replacer)
+                    doc = self.values_hook(doc, typed_fields)
+
+                    action = {"_index": index, "_type": doc_type, "_source": doc}
+
+                    if fields:
+                        # fields = ensure_str_list(id_field)
+                        if len(fields) == 1:
+                            action["_id"] = doc[fields[0]]
+                        else:
+                            action["_id"] = self.encode_id([doc[x] for x in fields])
+                    if generate_id:
+                        action["_id"] = self.encode_id(doc.values())
+
+                    counters["rows_read"] += 1
+                    counters["rows_yield"] += 1
+                    if counters["rows_yield"] % 10000 == 0:
+                        self.logger.info("progress: %s", counters)
+
+                    yield action
+
+        # 消费生成器
+        for _ in helpers.parallel_bulk(
+            self._es, actions=actions_generator(), thread_count=8, chunk_size=1000, queue_size=8
+        ):
+            pass
+
+    @staticmethod
+    def encode_id(values):
+        content = pickle.dumps(tuple(values))
+        return base64.urlsafe_b64encode(hashlib.sha1(content).digest()).decode()
+
+    @staticmethod
+    def values_hook(doc: dict, typed_fields: dict):
+        for k, v in doc.items():
+            field = typed_fields.get(k)
+            if field is None or field.type in [types.STRING]:
+                continue
+
+            doc[k] = typed_fields[k].cast(v)
+        return doc