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

recurvedata/connectors/const.py

@@ -0,0 +1,315 @@
+import json
+import logging
+import os
+from enum import Enum
+
+from recurvedata.connectors.base import RecurveConnectorBase
+from recurvedata.connectors.proxy import HttpProxyMixin
+from recurvedata.core.translation import _l
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_TIMEOUT = 10  # in seconds
+
+# todo(chenjingmeng): using auto generated by scripts/gen_const.py
+ALL_CONNECTION_SECRET_WORDS = [
+    "account_key",
+    "api_key",
+    "api_secret_key",
+    "app_secret",
+    "blob_options.sas_token",
+    "password",
+    "private_key",
+    "sas_token",
+    "secret_access_key",
+    "secret_key",
+    "key_dict.private_key",
+    "client_secret",
+    "access_key_secret",
+]
+
+
+def _format_connector_module_name(connection_type: str) -> str:
+    """
+    connection_type can be ui_connection_type
+    """
+    connection_type = connection_type.replace(" ", "_").lower()
+    return f"recurvedata.connectors.connectors.{connection_type}"
+
+
+DBAPI_TYPES = [
+    "azure_synapse",
+    "bigquery",
+    "clickhouse",
+    "elasticsearch",
+    "hive",
+    "impala",
+    "mongodb",
+    "mssql",
+    "mysql",
+    "phoenix",
+    "postgres",
+    "redis",
+    "redshift",
+    "starrocks",
+    "tidb",
+    "doris",
+    "microsoft_fabric",
+]
+CONNECTION_TYPE_MODULE_MAPPING = {  # todo(chenjingmeng): auto generated
+    "Tencent COS": "recurvedata.connectors.connectors.tencent_cos",
+    "cos": "recurvedata.connectors.connectors.tencent_cos",
+    "Elastic Search": "recurvedata.connectors.connectors.es",
+    "elasticsearch": "recurvedata.connectors.connectors.es",
+    "Ding Talk": "recurvedata.connectors.connectors.dingtalk",
+    "feishu_bot": "recurvedata.connectors.connectors.feishu",
+    "SelectDB(Doris)": "recurvedata.connectors.connectors.doris",
+    "azure_mssql": "recurvedata.connectors.connectors.mssql",
+    "google_bigquery": "recurvedata.connectors.connectors.bigquery",
+    "Google BigQuery": "recurvedata.connectors.connectors.bigquery",
+    "BigQuery": "recurvedata.connectors.connectors.bigquery",
+    "PostgreSQL": "recurvedata.connectors.connectors.postgres",
+    "MongoDB": "recurvedata.connectors.connectors.mongo",
+    "mongodb": "recurvedata.connectors.connectors.mongo",
+    "Microsoft SQL Server": "recurvedata.connectors.connectors.mssql",
+    "selectdb(doris)": "recurvedata.connectors.connectors.doris",
+    "apache impala": "recurvedata.connectors.connectors.impala",
+    "Apache Impala": "recurvedata.connectors.connectors.impala",
+    "Aliyun OSS": "recurvedata.connectors.connectors.oss",
+}
+
+
+def get_module_name(connection_type: str) -> str:
+    if connection_type in CONNECTION_TYPE_MODULE_MAPPING:
+        return CONNECTION_TYPE_MODULE_MAPPING[connection_type]
+    return _format_connector_module_name(connection_type)
+
+
+SQL_OPERATOR_TYPES = [
+    "azure_mssql",
+    "azure_synapse",
+    "bigquery",
+    "clickhouse",
+    "google_bigquery",
+    "hive",
+    "impala",
+    "mssql",
+    "mysql",
+    "phoenix",
+    "postgres",
+    "redshift",
+    "starrocks",
+    "tidb",
+    "doris",
+    "microsoft_fabric",
+]
+JUICE_SYNC_ABLE_DBAPI_TYPES = ["azure_blob", "cos", "google_cloud_storage", "oss", "s3", "sftp"]
+
+# This Const is manually built, refer to  "https://docs.getdbt.com/docs/supported-data-platforms"
+# is there any web's api available?
+DBT_SUPPORTED_TYPES = [
+    # official trusted, seems that these database/data
+    # warehouse connector are more robust
+    # ------------------------------------------------
+    "spark",
+    "azure_synapse",
+    "bigquery",
+    "postgres",
+    "redshift",
+    # not implemented connectors
+    "alloy_db",
+    "athena",
+    "databricks",
+    "dremio",
+    "glue",
+    "materialize",
+    "microsoft_fabric",
+    "oracle_autonomous_database",
+    "snowflake",
+    "starburst",
+    "teradata",
+    # --------------------------------------------------
+    # community maintained
+    # --------------------------------------------------
+    "mysql",
+    "starrocks",
+    "clickhouse",
+    "doris",
+    "tidb",
+    "hive",
+    "impala",
+    # not implemented connectors
+    "duckdb",
+    "exasol_analytics",
+    "extrica",
+    "ibm_db2",
+    "infer",
+    "iomete",
+    "mindsdb",
+    "risingwave",
+    "rockset",
+    "single_store",
+    "sql_server",
+    "sqlite",
+    "timescaledb",
+    "upsolver",
+    "vertica",
+    "databend_cloud",
+    "yellowbrick",
+]
+
+# This is also manually built :)). Refer to: https://cube.dev/docs/product/configuration/data-sources
+CUBE_SUPPORTED_TYPES = [
+    "doris",
+    "postgres",
+    "bigquery",
+    "starrocks",
+    "mysql",
+]
+
+
+class ProcessDBMixin(object):
+    @staticmethod
+    def auth_preprocess_conf(data):
+        data = RecurveConnectorBase.preprocess_conf(data)
+        json_data = data.get("extra")
+        if json_data and isinstance(json_data, str):
+            data["extra"] = json.loads(json_data)
+        return data
+
+    @classmethod
+    def bigquery_preprocess_conf(cls, data: dict) -> dict:
+        data = RecurveConnectorBase.preprocess_conf(data)
+        proxies = data.get("proxies")
+        if proxies and not HttpProxyMixin.check_proxy(proxies):
+            logger.warning(f"proxies {proxies} is not available, use direct connect")
+            data["proxies"] = None
+        return data
+
+    @classmethod
+    def google_cloud_storage_preprocess_conf(cls, data: dict) -> dict:
+        data = RecurveConnectorBase.preprocess_conf(data)
+        proxies = data.get("proxies")
+        if proxies and not HttpProxyMixin.check_proxy(proxies):
+            logger.warning(f"proxies {proxies} is not available, use direct connect")
+            data["proxies"] = None
+        return data
+
+    @classmethod
+    def google_service_account_preprocess_conf(cls, data: dict) -> dict:
+        data = RecurveConnectorBase.preprocess_conf(data)
+        proxies = data.get("proxies")
+        if proxies and not HttpProxyMixin.check_proxy(proxies):
+            logger.warning(f"proxies {proxies} is not available, use direct connect")
+            data["proxies"] = None
+        return data
+
+    @classmethod
+    def oss_preprocess_conf(cls, data: dict) -> dict:
+        data = RecurveConnectorBase.preprocess_conf(data)
+        proxies = data.get("proxies")
+        if proxies and not HttpProxyMixin.check_proxy(proxies):
+            logger.warning(f"proxies {proxies} is not available, use direct connect")
+            data["proxies"] = None
+        return data
+
+    @staticmethod
+    def other_preprocess_conf(data):
+        data = RecurveConnectorBase.preprocess_conf(data)
+        json_data = data.get("data")
+        if json_data and isinstance(json_data, str):
+            real_data = json.loads(json_data)
+            return real_data
+        return data
+
+    @classmethod
+    def s3_preprocess_conf(cls, data: dict) -> dict:
+        data = RecurveConnectorBase.preprocess_conf(data)
+        proxies = data.get("proxies")
+        if proxies and not HttpProxyMixin.check_proxy(proxies):
+            logger.warning(f"proxies {proxies} is not available, use direct connect")
+            data["proxies"] = None
+        return data
+
+    @staticmethod
+    def spark_preprocess_conf(data):
+        data = RecurveConnectorBase.preprocess_conf(data)
+        execution_config = data.get("execution_config")
+        if execution_config:
+            execution_config_conf = execution_config.get("conf")
+            if execution_config_conf and isinstance(execution_config_conf, str):
+                execution_config_conf = json.loads(execution_config_conf)
+                execution_config["conf"] = execution_config_conf
+        return data
+
+
+CONNECTION_TYPE_PREPROCESS_CONF_MAPPING = {
+    "auth": ProcessDBMixin.auth_preprocess_conf,
+    "Auth": ProcessDBMixin.auth_preprocess_conf,
+    "bigquery": ProcessDBMixin.bigquery_preprocess_conf,
+    "BigQuery": ProcessDBMixin.bigquery_preprocess_conf,
+    "google_cloud_storage": ProcessDBMixin.google_cloud_storage_preprocess_conf,
+    "Google Cloud Storage": ProcessDBMixin.google_cloud_storage_preprocess_conf,
+    "google_service_account": ProcessDBMixin.google_service_account_preprocess_conf,
+    "Google Service Account": ProcessDBMixin.google_service_account_preprocess_conf,
+    "oss": ProcessDBMixin.oss_preprocess_conf,
+    "OSS": ProcessDBMixin.oss_preprocess_conf,
+    "s3": ProcessDBMixin.s3_preprocess_conf,
+    "S3": ProcessDBMixin.s3_preprocess_conf,
+    "spark": ProcessDBMixin.spark_preprocess_conf,
+    "Spark": ProcessDBMixin.spark_preprocess_conf,
+}
+
+
+# auto generated finish
+
+
+def preprocess_conf(connection_type: str, data: dict):
+    func = CONNECTION_TYPE_PREPROCESS_CONF_MAPPING.get(connection_type, RecurveConnectorBase.preprocess_conf)
+    return func(data)
+
+
+SSH_TUNNEL_CONFIG_SCHEMA = {
+    "type": "object",
+    "title": _l("SSH Tunnel Configuration"),
+    "description": _l("Configuration for establishing an SSH tunnel connection"),
+    "properties": {
+        "host": {"type": "string", "title": _l("Host Address")},
+        "user": {"type": "string", "title": _l("Username")},
+        "port": {
+            "type": "number",
+            "title": _l("Port Number"),
+            "default": 22,
+        },
+        "password": {"type": "string", "title": _l("Password")},
+        "private_key_str": {
+            "type": "string",
+            "title": _l("SSH Private Key"),
+            "description": _l("Private key content for SSH key-based authentication"),
+        },
+        "private_key_passphrase": {
+            "type": "string",
+            "title": _l("SSH Private Key Passphrase"),
+            "description": _l("Passphrase to decrypt the SSH private key if encrypted"),
+        },
+    },
+    "order": ["host", "user", "port", "password", "private_key_str", "private_key_passphrase"],
+    "secret": ["password", "private_key_str", "private_key_passphrase"],
+}
+
+
+class LoadMode(str, Enum):
+    OVERWRITE = "OVERWRITE"
+    APPEND = "APPEND"
+
+
+ENV_VAR_DBT_USER = '{{ env_var("DBT_USER") }}'  # after yaml dump, single quote will become '', which cause dbt error
+ENV_VAR_DBT_PASSWORD = '{{ env_var("DBT_PASSWORD") }}'
+
+
+def set_env_dbt_user(user_name: str):
+    os.environ["DBT_USER"] = user_name
+
+
+def set_env_dbt_password(password: str):
+    os.environ["DBT_PASSWORD"] = password

recurvedata/connectors/datasource.py

@@ -0,0 +1,189 @@
+"""
+之前 Pigeon 里叫法是 DataSource
+这里先封装一个类似的，
+之后再整合到 Base 里
+"""
+
+import copy
+from dataclasses import dataclass
+from typing import Optional
+
+from recurvedata.connectors.const import DBAPI_TYPES, preprocess_conf
+
+
+@dataclass
+class DataSourceBase:
+    connection_type: str
+    data: dict
+    name: str = ""
+    extra: dict = None
+
+    def __post_init__(self):
+        self.data = preprocess_conf(self.connection_type, self.data)
+        self.process_pigeon_keyword()
+        self.extra = copy.deepcopy(self.data)
+
+    def process_pigeon_keyword(self):
+        """
+        OneFlow 使用的是 pigeon 的关键词，部分和 Recurve 不一致。
+        历史原因，从 OneFlow 迁移过来时，是直接复制数据库数据，没有转换关键词，
+        导致 Recurve 数据库里同时存在两套。
+        这里把旧的 pigeon 关键词转成新的 Recurve 关键词
+        :return:
+        """
+        from recurvedata.connectors.pigeon import DataSource as PigeonDataSource
+
+        keyword_renames: dict = PigeonDataSource.PIGEON_KEYWORD_MAPPING.get(self.connection_type, {})
+        for recurve_keyword, pigeon_keyword in keyword_renames.items():
+            if not pigeon_keyword or recurve_keyword in self.data or pigeon_keyword not in self.data:
+                continue
+            self.data[recurve_keyword] = self.data[pigeon_keyword]
+
+    @property
+    def recurve_connector_cls(self):
+        from recurvedata.connectors._register import get_connection_class
+
+        return get_connection_class(self.connection_type)
+
+    @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.data)
+        return recurve_con
+
+    def juice_sync_path(self, path: str) -> tuple[str, str]:
+        """
+        Return the paths used in juice sync.
+        The first return value is the path with a secret key,
+        and the second return value is the path without a secret key, intended for display purposes
+        """
+        if not self.recurve_connector_cls:
+            raise ValueError(f"{self.connection_type} is not juice sync able")
+        return self.recurve_connector.juice_sync_path(path)
+
+
+class DataSource(DataSourceBase):
+    """
+    pigeon 里 DataSource 等同于 Connection，这里保留两种叫法
+    """
+
+    @property
+    def connector(self):
+        """
+        暂时保留之前 OneFlow 做法，返回 pigeon Connector 对象
+        """
+        from recurvedata.connectors.pigeon import DataSource as PigeonDataSource
+
+        if PigeonDataSource.is_support_connection_type(self.connection_type):
+            pigeon_ds = PigeonDataSource(connection_type=self.connection_type, name=self.name, data=self.data)
+            return pigeon_ds.connector
+        recurve_cls = self.recurve_connector_cls
+        if not recurve_cls:
+            raise ValueError(f"Unknown connection type: {self.connection_type}")
+        recurve_con = recurve_cls(self.data)
+        return recurve_con
+
+    def create_engine(self):
+        return self.connector.create_engine()
+
+    @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 password(self):
+        # used in email load
+        return self.data.get("password")
+
+    @property
+    def user(self):
+        # used in email load
+        return self.data.get("user")
+
+
+class DataSourceWrapper(object):
+    """封装 DataSource，只保留必要的只读功能"""
+
+    def __init__(self, ds: DataSource):
+        self.__ds = ds
+
+    @property
+    def name(self) -> str:
+        return self.__ds.name
+
+    @property
+    def ds_type(self) -> str:
+        return self.__ds.connection_type
+
+    @property
+    def host(self) -> Optional[str]:
+        return self.__ds.host
+
+    @property
+    def database(self) -> Optional[str]:
+        return self.__ds.database
+
+    @property
+    def user(self) -> Optional[str]:
+        return self.__ds.user
+
+    @property
+    def password(self) -> Optional[str]:
+        # 最初的设计里，ds_wrapper 不提供 password，但是后来使用的地方较多
+        return self.__ds.password
+
+    @property
+    def port(self) -> Optional[int]:
+        return self.__ds.port
+
+    @property
+    def is_dbapi(self) -> bool:
+        return self.ds_type in DBAPI_TYPES
+
+    @property
+    def connector(self):
+        return self.__ds.connector
+
+    @property
+    def recurve_connector(self):
+        return self.__ds.recurve_connector
+
+    @property
+    def data(self) -> dict:
+        if self.ds_type == "other":
+            return self.__ds.data.get("data", self.__ds.data)
+        return self.__ds.data
+
+    @property
+    def extra(self) -> dict:
+        if self.ds_type == "other":
+            return self.__ds.extra.get("data", self.__ds.extra)
+        return self.__ds.extra
+
+    def __getattr__(self, name):
+        if name in self.data:
+            return self.data[name]
+        return super().__getattribute__(name)  # raise