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

recurvedata/connectors/pigeon.py

@@ -0,0 +1,172 @@
+"""
+使用 pigeon connector 作为底层实现
+返回 pigeon connector
+"""
+
+import copy
+import os
+
+from recurvedata.connectors.datasource import DataSourceBase
+
+
+class DataSource(DataSourceBase):
+    """
+    pigeon 里 DataSource 等同于 Connection，这里保留两种叫法
+    """
+
+    def __post_init__(self):
+        super().__post_init__()
+        keyword_renames: dict = self.PIGEON_KEYWORD_MAPPING.get(self.connection_type, {})
+        for recurve_keyword, pigeon_keyword in keyword_renames.items():
+            val = self.data.pop(recurve_keyword, None)
+            if pigeon_keyword:
+                self.data[pigeon_keyword] = val
+                self.extra[pigeon_keyword] = self.extra.get(recurve_keyword)
+        if "passwd" in self.data:
+            self.data["password"] = self.data["passwd"]
+            self.extra["password"] = self.extra["passwd"]
+        if self.connection_type == "bigquery":
+            self.data["dataset"] = self.data.get("database")
+            if "key_dict" not in self.data:
+                # some are not saved through the web page, have key_dict field, no need to modify
+                key_dict = copy.deepcopy(self.data)
+                key_dict.pop("location", None)
+                key_dict.pop("proxies", None)
+                self.data["key_dict"] = key_dict
+                self.extra["key_dict"] = key_dict
+        if self.connection_type == "sftp":
+            pk_path = self.data.get("rsa_private_key_file")
+            if pk_path == "":
+                self.data.pop("rsa_private_key_file")
+            elif pk_path:
+                self.data["rsa_private_key_file"] = os.path.expanduser(pk_path)
+                self.extra["rsa_private_key_file"] = os.path.expanduser(pk_path)
+
+    PIGEON_KEYWORD_MAPPING = {
+        "s3": {"access_key_id": "aws_access_key_id", "secret_access_key": "aws_secret_access_key", "bucket": None},
+        "filebrowser": {"url": "host"},
+        "kafka": {
+            "bootstrap_servers": "bootstrap.servers",
+            "value_deserializer": "value.deserializer",
+            "value_schema_registry_client_url": "schema_registry_client_url",
+        },
+        "oss": {"secret_access_key": "access_key_secret", "access_key_id": "access_key_id", "bucket": "bucket_name"},
+        "sftp": {"user": "username", "private_key_path": "rsa_private_key_file"},
+        "mongodb": {"user": "username"},
+        "bigquery": {"database": "dataset"},
+    }
+
+    PIGEON_TYPE_CLS_MAPPING = {
+        "s3": "recurvedata.pigeon.connector.awss3",
+        "oss": "recurvedata.pigeon.connector.aliyun_oss",
+        "azure_blob": "recurvedata.pigeon.connector.azure_blob",
+        "azure_synapse": "recurvedata.pigeon.connector.azure_synapse",
+        "clickhouse": "recurvedata.pigeon.connector.clickhouse_native",
+        "es": "recurvedata.pigeon.connector.es",
+        "ftp": "recurvedata.pigeon.connector.ftp",
+        "google_bigquery": "recurvedata.pigeon.connector.google_bigquery",
+        "phoenix": "recurvedata.pigeon.connector.hbase_phoenix",
+        "hdfs": "recurvedata.pigeon.connector.hdfs",
+        "hive": "recurvedata.pigeon.connector.hive_impala",
+        "impala": "recurvedata.pigeon.connector.hive_impala",
+        "mongodb": "recurvedata.pigeon.connector.mongodb",
+        "mssql": "recurvedata.pigeon.connector.mssql",
+        "mysql": "recurvedata.pigeon.connector.mysql",
+        "postgresql": "recurvedata.pigeon.connector.postgresql",
+        "postgres": "recurvedata.pigeon.connector.postgresql",
+        "qcloud_cos": "recurvedata.pigeon.connector.qcloud_cos",
+        "cos": "recurvedata.pigeon.connector.qcloud_cos",
+        "redshift": "recurvedata.pigeon.connector.redshift",
+        "sftp": "recurvedata.pigeon.connector.sftp",
+        "owncloud": "recurvedata.pigeon.connector.owncloud",
+        "starrocks": "recurvedata.pigeon.connector.starrocks",
+        "doris": "recurvedata.pigeon.connector.doris",
+        "microsoft_fabric": "recurvedata.pigeon.connector.microsoft_fabric",
+    }
+
+    PIGEON_TYPE_MAPPING = {
+        "tidb": "mysql",
+        "bigquery": "google_bigquery",
+    }
+
+    @property
+    def connector(self):
+        """
+        暂时保留之前 OneFlow 做法，返回 pigeon Connector 对象
+        """
+        return self.get_pigeon_connector()
+
+    def create_engine(self):
+        return self.connector.create_engine()
+
+    def get_pigeon_connector(self):
+        klass = self.get_pigeon_connector_class()
+        if not klass:
+            raise ValueError(f"{self.connection_type} has no pigeon class")
+
+        return klass(**self.data)
+
+    def get_pigeon_connector_class(self):
+        from recurvedata.pigeon.connector import get_connector_class
+
+        connection_type = self.PIGEON_TYPE_MAPPING.get(self.connection_type, self.connection_type)
+        try:
+            return get_connector_class(connection_type)
+        except KeyError:
+            pigeon_cls = self.PIGEON_TYPE_CLS_MAPPING[connection_type]
+            __import__(pigeon_cls)
+            return get_connector_class(connection_type)
+
+    @property
+    def host(self):
+        for key in ["url", "host"]:
+            if key in self.data:
+                return self.data[key]
+
+    @property
+    def ds_type(self):
+        # 兼容 oneflow lineage
+        return self.connection_type
+
+    @property
+    def database(self):
+        # used in postgres load
+        return self.data.get("database")
+
+    @property
+    def port(self):
+        # used in email load
+        return self.data.get("port")
+
+    @property
+    def user(self):
+        # used in email load
+        return self.data.get("user")
+
+    @property
+    def password(self):
+        # used in email load
+        return self.data.get("password")
+
+    @classmethod
+    def is_support_connection_type(cls, connection_type: str) -> bool:
+        connection_type = cls.PIGEON_TYPE_MAPPING.get(connection_type, connection_type)
+        return connection_type in cls.PIGEON_TYPE_CLS_MAPPING
+
+    @property
+    def recurve_connector(self):
+        """
+        和 pigeon connector 区分开
+        :return:
+        """
+        recurve_cls = self.recurve_connector_cls
+        if not recurve_cls:
+            raise ValueError(f"Unknown connection type: {self.connection_type}")
+        recurve_con = recurve_cls(self.extra)
+        return recurve_con
+
+    # todo: is_dbapi
+
+
+def get_pigeon_connector(connection_type: str, data: dict):
+    return DataSource(connection_type=connection_type, data=data).get_pigeon_connector()

recurvedata/connectors/proxy.py

@@ -0,0 +1,104 @@
+import socket
+import urllib.parse
+
+from recurvedata.connectors.base import RecurveConnectorBase
+from recurvedata.connectors.utils import EnvContextManager
+from recurvedata.core.translation import _l
+from recurvedata.utils.log import LoggingMixin
+
+logger = LoggingMixin.logger()
+
+HTTP_PROXY_CONFIG_SCHEMA = {
+    "proxies": {
+        "type": "object",
+        "title": _l("HTTP/HTTPS Proxy Settings"),
+        "description": _l("Configure proxy servers for HTTP and HTTPS connections"),
+        "properties": {
+            "https": {
+                "type": "string",
+                "title": _l("HTTPS Proxy"),
+                "description": _l("HTTPS proxy URL in format https://host:port, or http://host:port"),
+            },
+            "http": {
+                "type": "string",
+                "title": _l("HTTP Proxy"),
+                "description": _l("HTTP proxy URL in format http://host:port"),
+            },
+        },
+        "order": ["https", "http"],
+    },
+}
+
+
+class ProxyEnvContextManager(EnvContextManager):
+    def __init__(self, proxy: str = None, http_proxy: str = None, https_proxy: str = None):
+        env_vars = {}
+        if http_proxy or proxy:
+            env_vars["http_proxy"] = http_proxy or proxy
+        if https_proxy or proxy:
+            env_vars["https_proxy"] = https_proxy or proxy
+        super().__init__(env_vars)
+
+
+class ProxyMixinBase:
+    pass
+
+
+class HttpProxyMixin(ProxyMixinBase):
+    @classmethod
+    def format_config_schema(cls):
+        config_schema = super(HttpProxyMixin, cls).format_config_schema()
+        return add_http_proxy_to_config_schema(config_schema)
+
+    def _init_proxy_manager(self):
+        proxies = self.proxies or {}
+        return ProxyEnvContextManager(http_proxy=proxies.get("http"), https_proxy=proxies.get("https"))
+
+    @classmethod
+    def preprocess_conf(cls, data: dict) -> dict:
+        data = RecurveConnectorBase.preprocess_conf(data)
+        proxies = data.get("proxies")
+        if proxies and not cls.check_proxy(proxies):
+            logger.warning(f"proxies {proxies} is not available, use direct connect")
+            data["proxies"] = None
+        return data
+
+    @classmethod
+    def check_proxy(cls, proxies: dict, timeout=10):
+        """
+        检查 proxy 是否可用
+        proxies example: {'http': 'http://proxy_host:proxy_port', 'https': 'https://proxy_host:proxy_port'}
+        """
+        if not proxies:
+            return False
+        return cls._check_proxy_connection(proxies.get("http"), timeout) or cls._check_proxy_connection(
+            proxies.get("https"), timeout
+        )
+
+    @classmethod
+    def _check_proxy_connection(cls, proxy_url, timeout):
+        parsed_url = urllib.parse.urlparse(proxy_url)
+        if not parsed_url.scheme:
+            proxy_url = f"http://{proxy_url}"
+            parsed_url = urllib.parse.urlparse(proxy_url)
+        hostname = parsed_url.hostname
+        port = parsed_url.port or (443 if parsed_url.scheme == "https" else 80)
+
+        try:
+            with socket.create_connection((hostname, port), timeout=timeout):
+                return True
+        except socket.timeout:
+            logger.info(f"Connection timed out while connecting to {hostname}:{port}")
+            return False
+        except socket.error as e:
+            logger.info(f"Error connecting to {hostname}:{port}: {e}")
+            return False
+
+
+def add_http_proxy_to_config_schema(config_schema: dict) -> dict:
+    if "proxies" in config_schema["properties"]:
+        return config_schema
+
+    config_schema["properties"].update(HTTP_PROXY_CONFIG_SCHEMA)
+    config_schema["order"].append("proxies")
+    return config_schema

recurvedata/connectors/service.py

@@ -0,0 +1,223 @@
+import hashlib
+import json
+import os
+import traceback
+from typing import Any
+
+from sqlalchemy.engine.url import URL
+
+from recurvedata.connectors._register import get_connection_class
+from recurvedata.connectors.config_schema import ALL_CONFIG_SCHEMA_DCT, get_complex_config_schema
+from recurvedata.connectors.const import (  # noqa
+    ALL_CONNECTION_SECRET_WORDS,
+    DBAPI_TYPES,
+    JUICE_SYNC_ABLE_DBAPI_TYPES,
+    SQL_OPERATOR_TYPES,
+)
+from recurvedata.connectors.datasource import DataSource, DataSourceWrapper
+from recurvedata.connectors.pigeon import DataSource as PigeonDataSource
+from recurvedata.consts import PROJECT_ID_KEY
+
+
+def list_config_schemas(only_enabled=True) -> list[dict[str, Any]]:
+    """
+    todo: 返回的类型 (看要不要用 pydantic)
+    :param only_enabled:
+    :return:
+    """
+    config_schemas = ALL_CONFIG_SCHEMA_DCT.values()
+    if only_enabled:
+        config_schemas = filter(lambda v: v["enabled"], config_schemas)
+    return list(config_schemas)
+
+
+def get_config_schema(connection_type: str):
+    complex_config_schema = get_complex_config_schema(connection_type)
+    if complex_config_schema:
+        return complex_config_schema["config_schema"]
+
+
+def get_connection_category(connection_type: str):
+    complex_config_schema = get_complex_config_schema(connection_type)
+    if complex_config_schema:
+        return complex_config_schema["category"]
+
+
+def get_connection_ui_category(connection_type: str):
+    complex_config_schema = get_complex_config_schema(connection_type)
+    if complex_config_schema:
+        return complex_config_schema["ui_category"]
+
+
+def get_connection_type(connection_type: str):
+    complex_config_schema = get_complex_config_schema(connection_type)
+    if complex_config_schema:
+        return complex_config_schema["type"]
+
+
+def get_connection_ui_type(connection_type: str):
+    complex_config_schema = get_complex_config_schema(connection_type)
+    if complex_config_schema:
+        return complex_config_schema["ui_type"]
+
+
+def get_connection_host(ui_type: str, connection_conf: dict):
+    """
+    前端页面列表页显示的 ui_type
+    :param ui_type:
+    :param connection_conf:
+    :return:
+    """
+    for keyword in ["host", "endpoint", "access_key_id", "url"]:  # todo: not so good
+        if keyword in connection_conf:
+            return connection_conf[keyword]
+    return ""
+
+
+def get_all_secret_keywords():
+    return ALL_CONNECTION_SECRET_WORDS
+
+
+def test_connection(connection_type: str, connection_conf: dict) -> tuple[bool, str]:
+    connection_cls = get_connection_class(connection_type)
+    try:
+        con = connection_cls(connection_conf)
+        con.test_connection()
+        return True, ""
+    except Exception:
+        return False, traceback.format_exc()
+
+
+def init_connector(connection_type: str, connection_conf: dict):
+    pass
+
+
+def list_dbapi_types():
+    return DBAPI_TYPES
+
+
+def list_juice_sync_able_dbapi_types():
+    return JUICE_SYNC_ABLE_DBAPI_TYPES
+
+
+def list_sql_operator_types():
+    return SQL_OPERATOR_TYPES
+
+
+def get_datasource_by_name(project_connection_name: str, project_id: int | None = None) -> DataSourceWrapper:
+    from recurvedata.executors.client import ExecutorClient
+    from recurvedata.executors.schemas import ConnectionItem
+
+    project_id = os.environ.get(PROJECT_ID_KEY) if project_id is None else project_id
+    if project_id is None:
+        raise ValueError("project id is not set")
+
+    recurve_client = ExecutorClient()
+    conn: ConnectionItem = recurve_client.get_connection(project_id=project_id, connection_name=project_connection_name)
+    if PigeonDataSource.is_support_connection_type(conn.type):
+        try:
+            return DataSourceWrapper(PigeonDataSource(connection_type=conn.type, name=conn.name, data=conn.data))
+        except ModuleNotFoundError:
+            pass
+    return DataSourceWrapper(DataSource(connection_type=conn.type, name=conn.name, data=conn.data))
+
+
+def get_datasource_by_config(
+    connection_type: str, config: dict, name: str = None, database: str = None, schema: str = None
+) -> DataSourceWrapper:
+    """
+    Get a DataSourceWrapper instance based on the connection type and configuration dictionary.
+
+    :param connection_type: The type of the connection (e.g., 'mysql', 'postgresql', 'snowflake')
+    :param config: A dictionary containing the connection configuration
+    :param database: project database name
+    :param schema: project schema name if have
+    :return: A DataSourceWrapper instance
+    """
+    update_dct = {}
+    if database:
+        update_dct.update({"database": database})
+    if schema:
+        update_dct.update({"schema": schema})
+    config.update(update_dct)
+    if PigeonDataSource.is_support_connection_type(connection_type):
+        try:
+            return DataSourceWrapper(PigeonDataSource(connection_type=connection_type, name=name, data=config))
+        except ModuleNotFoundError:
+            pass
+    return DataSourceWrapper(DataSource(connection_type=connection_type, name=name, data=config))
+
+
+def get_sqlalchemy_url_by_connection(connection_orm) -> URL:
+    ds = DataSourceWrapper(
+        DataSource(connection_type=connection_orm.type, name=connection_orm.name, data=connection_orm.data)
+    )
+    if not ds.is_dbapi:
+        raise ValueError(f"{ds.ds_type} is not dbapi, not support this function")
+    con = ds.recurve_connector
+    return con.sqlalchemy_url
+
+
+# todo: cache
+def list_column_data_types():
+    import sqlalchemy.sql.sqltypes
+
+    def _get_module_types(module):
+        types = set()
+        for cls_name, cls in module.__dict__.items():
+            if cls_name.startswith("_"):
+                continue
+            if not isinstance(cls, type):
+                continue
+            if not issubclass(cls, sqlalchemy.sql.sqltypes.TypeEngine):
+                continue
+            if not hasattr(cls, "__visit_name__"):
+                continue
+            if cls.__visit_name__ in ("TypeDecorator", "type_decorator"):
+                continue
+            types.add(cls.__visit_name__.lower())
+        return types
+
+    types = _get_module_types(sqlalchemy.sql.sqltypes)
+    try:
+        import sqlalchemy.dialects.mysql.types
+
+        types = types.union(_get_module_types(sqlalchemy.dialects.mysql.types))
+    except ImportError:
+        pass
+
+    try:
+        import clickhouse_sqlalchemy.types.common
+
+        types = types.union(_get_module_types(clickhouse_sqlalchemy.types))
+    except ImportError:
+        pass
+
+    return sorted(types)
+
+
+def convert_connection_to_dbt_profile(
+    connection_type: str, connection_config: dict, database: str, schema: str = None
+) -> dict:
+    ds = get_datasource_by_config(connection_type, connection_config, database=database, schema=schema)
+    if not ds.is_dbapi:
+        return {}
+    con = ds.recurve_connector
+
+    dct = con.convert_config_to_dbt_profile(database, schema)
+    return dct
+
+
+def convert_connection_to_cube_config(
+    connection_type: str, connection_config: dict, database: str, schema: str = None, masking: bool = None
+) -> dict:
+    ds = get_datasource_by_config(connection_type, connection_config, database=database, schema=schema)
+    con = ds.recurve_connector
+    dct = con.convert_config_to_cube_config(database, schema, ds)
+    if not dct or not isinstance(dct, dict):
+        return {}
+
+    masking = masking if masking is not None else not ds.is_dbapi
+    if masking:
+        return {k: hashlib.md5(json.dumps(dct[k]).encode("utf-8")).hexdigest() for k in dct}
+    return dct

recurvedata/connectors/utils.py

@@ -0,0 +1,47 @@
+import os
+import re
+from urllib.parse import quote
+
+
+class EnvContextManager:
+    def __init__(self, env_vars: dict):
+        self.env_vars = env_vars
+        self.old_env_vars = {}
+
+    def __enter__(self):
+        for key, value in self.env_vars.items():
+            self.old_env_vars[key] = os.environ.get(key)
+            os.environ[key] = value
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        for key, value in self.old_env_vars.items():
+            if value is None:
+                del os.environ[key]
+            else:
+                os.environ[key] = value
+
+
+def juice_sync_process_special_character_within_secret(secret_part: str) -> str:
+    """
+    `When you get "/" in ACCESS_KEY or SECRET_KEY strings,you need to replace "/" with "%2F".`
+    :return:
+    """
+    if "/" in secret_part:
+        secret_part = secret_part.replace("/", "%2F")
+    return secret_part
+
+
+def juice_sync_process_special_character_within_path(path: str) -> str:
+    """
+    1. 冒号需要在 juice sync 里处理两遍，第一遍转成 %3A, 第二遍再把 %3A quote 一下，用于 juice sync
+    2. 有些路径需要加引号
+    """
+    if not path:
+        return path
+    colon_quote = quote(":")
+    if colon_quote in path:
+        path = path.replace(colon_quote, quote(colon_quote))
+
+    if re.search("[ &]", path):
+        path = f'"{path}"'
+    return path

recurvedata/consts.py

@@ -0,0 +1,49 @@
+from enum import Enum
+
+# docker container environ key for recurve environment_id
+ENV_ID_KEY = "RECURVE__ENVIRONMENT_ID"
+PROJECT_ID_KEY = "RECURVE__PROJECT_ID"
+
+
+class ScheduleType(str, Enum):
+    crontab = "crontab"
+    customization = "customization"
+    manual = "manual"
+
+
+class Operator(str, Enum):
+    SQLOperator = "SQLOperator"
+    TransferOperator = "TransferOperator"
+    PythonOperator = "PythonOperator"
+    SparkOperator = "SparkOperator"
+    NotifyOperator = "NotifyOperator"
+    LinkOperator = "LinkOperator"
+    DBTOperator = "DBTOperator"
+    LinkModelPipelineOperator = "LinkModelPipelineOperator"
+    SensorOperator = "SensorOperator"
+
+    @classmethod
+    def is_link(cls, op: str):
+        return op in (cls.LinkOperator, cls.LinkModelPipelineOperator)
+
+    # todo(chenjingmeng): support dynamically added Operator
+
+
+class ETLExecutionStatus(str, Enum):
+    SUCCESS = "success"
+    FAILED = "failed"
+    RETRY = "retry"
+
+
+class ConnectionCategory(str, Enum):
+    DATABASE = "database"
+    STORAGE = "storage"
+    WAREHOUSE = "warehouse"
+    OTHERS = "others"
+    SERVICE = "service"
+
+
+class ConnectorGroup(str, Enum):
+    SOURCE = "source"
+    DESTINATION = "destination"
+    INTEGRATION = "integration"

recurvedata/core/config.py

@@ -0,0 +1,46 @@
+import os
+import pathlib
+
+
+class Config:
+    CONFIG_FOLDER_ENV_NAME = "RECURVE_CONFIG_FOLDER"
+    __instance = None
+
+    def __new__(cls, *args, **kwargs):
+        if cls.__instance is None:
+            cls.__instance = super().__new__(cls, *args, **kwargs)
+        return cls.__instance
+
+    def __init__(self, folder: str = None):
+        folder = folder or os.environ.get(self.CONFIG_FOLDER_ENV_NAME, "~/.recurve")
+        if not folder:
+            raise ValueError(f"config folder is required, got {repr(folder)}")
+
+        self._folder = pathlib.Path(folder).expanduser()
+        self._folder.mkdir(mode=0o755, parents=True, exist_ok=True)
+        self._config_file = self._folder / "config"
+
+        self.__defaults = self.load_defaults()
+
+    def load_defaults(self) -> dict:
+        """load default config from env file or environment variables
+
+        file format:
+        RECURVE_HOST = https://recurve.test.recurvedata.com
+        RECURVE_USERNAME = foo
+        RECURVE_PASSWORD = pwd123
+        """
+        from dotenv import dotenv_values
+
+        # failed to install dotenv in 3.11  todo
+        return {
+            **dotenv_values(self._config_file),
+            **os.environ,
+        }
+
+    def get_or_default(self, value: str, env_key: str) -> str:
+        if value:
+            return value
+        if env_key not in self.__defaults:
+            raise ValueError("value is required")
+        return self.__defaults[env_key]

recurvedata/core/configurable.py

@@ -0,0 +1,27 @@
+from typing import Any, ClassVar, Union
+
+from pydantic import BaseModel, ConfigDict
+
+
+class BaseConfigModel(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+
+class Configurable:
+    config_model: ClassVar[BaseConfigModel]
+
+    @classmethod
+    def config_schema(cls) -> dict[str, Any]:
+        return cls.config_model.model_json_schema()
+
+    @classmethod
+    def name(cls) -> str:
+        return cls.__name__
+
+    @classmethod
+    def to_dict(cls) -> dict[str, Any]:
+        return {"name": cls.name(), "config_schema": cls.config_schema()}
+
+    @classmethod
+    def validate(cls, configuration: Union[dict, BaseConfigModel]) -> BaseConfigModel:
+        return cls.config_model.model_validate(configuration)

recurvedata/core/consts.py

@@ -0,0 +1,2 @@
+TRACE_KEY = "_otlp_trace"
+TRACING_CONTEXT_KEY = "_otlp_context"