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

recurvedata/connectors/connectors/generic.py

@@ -0,0 +1,49 @@
+import json
+
+from recurvedata.connectors._register import register_connector_class
+from recurvedata.connectors.base import RecurveConnectorBase
+from recurvedata.consts import ConnectorGroup
+from recurvedata.core.translation import _l
+
+CONNECTION_TYPE = "generic"
+UI_CONNECTION_TYPE = "Generic"
+
+
+@register_connector_class([CONNECTION_TYPE, UI_CONNECTION_TYPE])
+class GenericConnector(RecurveConnectorBase):
+    connection_type = CONNECTION_TYPE
+    ui_connection_type = UI_CONNECTION_TYPE
+    group = [ConnectorGroup.DESTINATION]
+    test_required = False
+
+    config_schema = {
+        "type": "object",
+        "properties": {
+            "host": {"type": "string", "title": _l("Host")},
+            "port": {"type": "integer", "title": _l("Port")},
+            "user": {"type": "string", "title": _l("Username")},
+            "password": {"type": "string", "title": _l("Password")},
+            "timeout": {"type": "integer", "title": _l("Timeout (seconds)"), "default": 30},
+            "custom": {
+                "type": "string",
+                "title": _l("Custom Configuration"),
+                "description": _l("Custom configuration parameters in JSON format"),
+                "ui:field": "CodeEditorWithReferencesField",
+                "ui:options": {"type": "code", "lang": "json"},
+            },
+        },
+        "order": ["host", "port", "user", "password", "timeout", "custom"],
+        "required": ["host", "port", "user", "password"],
+        "secret": ["password"],
+    }
+
+    def test_connection(self):
+        pass
+
+    @staticmethod
+    def preprocess_conf(data):
+        data = RecurveConnectorBase.preprocess_conf(data)
+        json_data = data.get("custom")
+        if json_data and isinstance(json_data, str):
+            data["custom"] = json.loads(json_data)
+        return data

recurvedata/connectors/connectors/google_cloud_storage.py

@@ -0,0 +1,115 @@
+from recurvedata.connectors._register import register_connector_class
+from recurvedata.connectors.object_store import ObjectStoreMixin
+from recurvedata.connectors.proxy import HTTP_PROXY_CONFIG_SCHEMA, HttpProxyMixin
+from recurvedata.consts import ConnectorGroup
+from recurvedata.core.translation import _l
+
+CONNECTION_TYPE = "google_cloud_storage"
+UI_CONNECTION_TYPE = "Google Cloud Storage"
+
+
+@register_connector_class([CONNECTION_TYPE, UI_CONNECTION_TYPE])
+class GoogleCloudStorage(HttpProxyMixin, ObjectStoreMixin):
+    connection_type = CONNECTION_TYPE
+    ui_connection_type = UI_CONNECTION_TYPE
+    setup_extras_require = []
+    group = [ConnectorGroup.DESTINATION]
+    test_required = False
+
+    config_schema = {
+        "type": "object",
+        "properties": {
+            "key_dict": {
+                "type": "object",
+                "title": _l("Service Account Key"),
+                "description": _l("Google Cloud service account key credentials"),
+                "properties": {
+                    "type": {
+                        "type": "string",
+                        "title": _l("Account Type"),
+                        "default": "service_account",
+                    },
+                    "project_id": {"type": "string", "title": _l("Google Cloud Project ID")},
+                    "private_key_id": {"type": "string", "title": _l("Google Auth Private Key ID")},
+                    "private_key": {
+                        "type": "string",
+                        "title": _l("Google Auth Private Key"),
+                        "ui:options": {"type": "textarea"},
+                    },
+                    "client_email": {"type": "string", "title": _l("Service Account Email")},
+                    "client_id": {"type": "string", "title": _l("Google OAuth Client ID")},
+                    "auth_uri": {
+                        "type": "string",
+                        "title": _l("Google OAuth Auth URI"),
+                        "default": "https://accounts.google.com/o/oauth2/auth",
+                    },
+                    "token_uri": {
+                        "type": "string",
+                        "title": _l("Google OAuth Token URI"),
+                        "default": "https://oauth2.googleapis.com/token",
+                    },
+                    "auth_provider_x509_cert_url": {
+                        "type": "string",
+                        "title": _l("Google OAuth Certificate URL (Auth Provider)"),
+                        "default": "https://www.googleapis.com/oauth2/v1/certs",
+                    },
+                    "client_x509_cert_url": {
+                        "type": "string",
+                        "title": _l("Google OAuth Certificate URL (Client)"),
+                        "default": "https://www.googleapis.com/robot/v1/metadata/x509/recurvedata-gcs%40brand-portal-prod.iam.gserviceaccount.com",
+                    },
+                },
+                "order": [
+                    "type",
+                    "project_id",
+                    "private_key_id",
+                    "private_key",
+                    "client_email",
+                    "client_id",
+                    "auth_uri",
+                    "token_uri",
+                    "auth_provider_x509_cert_url",
+                    "client_x509_cert_url",
+                ],
+                "required": [
+                    "type",
+                    "project_id",
+                    "private_key_id",
+                    "private_key",
+                    "client_id",
+                ],
+                "secret": [
+                    "private_key",
+                ],
+            },
+            "bucket": {
+                "type": "string",
+                "title": _l("Bucket Name"),
+                "description": _l("Name of the Google Cloud Storage bucket"),
+            },
+            "proxies": HTTP_PROXY_CONFIG_SCHEMA["proxies"],
+        },
+        "order": [
+            "key_dict",
+            "bucket",
+            "proxies",
+        ],
+        "required": [
+            "key_dict",
+        ],
+        "secret": [
+            "key_dict.private_key",
+        ],
+    }
+
+    def init_connection(self, conf):
+        self.connector = None  # todo
+
+    def test_connection(self):
+        # todo
+        pass
+
+    juice_sync_able = True
+
+    def juice_sync_path(self, path: str) -> str:
+        return f"gcs://{path}"  # todo

recurvedata/connectors/connectors/google_service_account.py

@@ -0,0 +1,225 @@
+import re
+
+from recurvedata.consts import ConnectionCategory, ConnectorGroup
+
+try:
+    from pandas import DataFrame
+except ImportError:
+    pass
+
+from recurvedata.connectors._register import register_connector_class
+from recurvedata.connectors.base import RecurveConnectorBase
+from recurvedata.connectors.const import LoadMode
+from recurvedata.connectors.proxy import HTTP_PROXY_CONFIG_SCHEMA, HttpProxyMixin
+from recurvedata.core.translation import _l
+
+CONNECTION_TYPE = "google_service_account"
+UI_CONNECTION_TYPE = "Google Service Account"
+
+try:
+    import gspread  # noqa
+    import pandas as pd
+    from google.oauth2 import service_account
+    from gspread.worksheet import Worksheet  # noqa
+except ImportError:
+    pass
+
+
+@register_connector_class([CONNECTION_TYPE, UI_CONNECTION_TYPE])
+class GoogleServiceAccount(HttpProxyMixin, RecurveConnectorBase):
+    connection_type = CONNECTION_TYPE
+    ui_connection_type = UI_CONNECTION_TYPE
+    group = [ConnectorGroup.DESTINATION]
+    category = [ConnectionCategory.SERVICE]
+    setup_extras_require = [
+        "google-auth",
+    ]
+    # gspread 暂时不加到 setup_extras_require 里，而是加到 operator 的 setup 里
+    default_timeout = 120
+
+    config_schema = {
+        "type": "object",
+        "properties": {
+            "project_id": {"type": "string", "title": _l("Google Cloud Project ID")},
+            "private_key_id": {"type": "string", "title": _l("Google Auth Private Key ID")},
+            "private_key": {
+                "type": "string",
+                "title": _l("Google Auth Private Key"),
+                "ui:options": {"type": "textarea"},
+            },
+            "client_email": {"type": "string", "title": _l("Service Account Email")},
+            "client_id": {"type": "string", "title": _l("Google OAuth Client ID")},
+            "auth_uri": {
+                "type": "string",
+                "title": _l("Google OAuth Auth URI"),
+                "default": "https://accounts.google.com/o/oauth2/auth",
+            },
+            "token_uri": {
+                "type": "string",
+                "title": _l("Google OAuth Token URI"),
+                "default": "https://oauth2.googleapis.com/token",
+            },
+            "auth_provider_x509_cert_url": {
+                "type": "string",
+                "title": _l("Google OAuth Certificate URL (Auth Provider)"),
+                "default": "https://www.googleapis.com/oauth2/v1/certs",
+            },
+            "client_x509_cert_url": {
+                "type": "string",
+                "title": _l("Google OAuth Certificate URL (Client)"),
+                "default": "https://www.googleapis.com/robot/v1/metadata/x509/recurvedata-gcs%40brand-portal-prod.iam.gserviceaccount.com",
+            },
+            "universe_domain": {
+                "type": "string",
+                "title": _l("Universe Domain"),
+                "default": "googleapis.com",
+            },
+            "proxies": HTTP_PROXY_CONFIG_SCHEMA["proxies"],
+        },
+        "order": [
+            "project_id",
+            "private_key_id",
+            "private_key",
+            "client_email",
+            "client_id",
+            "auth_uri",
+            "token_uri",
+            "auth_provider_x509_cert_url",
+            "client_x509_cert_url",
+            "universe_domain",
+            "proxies",
+        ],
+        "required": [
+            "project_id",
+            "private_key_id",
+            "private_key",
+            "client_email",
+        ],
+        "secret": [
+            "private_key",
+        ],
+    }
+
+    def init_credential_key_dict(self):
+        self.private_key = self._convert_private_key(self.private_key)
+        _key_dict = {
+            "type": "service_account",
+            "project_id": self.project_id,
+            "private_key_id": self.private_key_id,
+            "private_key": self.private_key,
+            "client_email": self.client_email,
+            "auth_uri": self.auth_uri,
+            "token_uri": self.token_uri,
+            "auth_provider_x509_cert_url": self.auth_provider_x509_cert_url,
+            "client_x509_cert_url": self.client_x509_cert_url,
+            "universe_domain": self.universe_domain,
+        }
+        if self.client_id:
+            _key_dict["client_id"] = self.client_id
+        return _key_dict
+
+    def init_credential(self):
+        credentials = service_account.Credentials.from_service_account_info(info=self.init_credential_key_dict())
+        return credentials
+
+    @staticmethod
+    def _convert_private_key(private_key: str):
+        # Depending on how the JSON was formatted, it may contain
+        # escaped newlines. Convert those to actual newlines.
+        private_key = private_key.replace("\\\n", "\n")
+        return private_key.replace("\\n", "\n")
+
+    def test_connection(self):
+        # 暂时不校验。如果私钥有问题这里好像会报错
+        with self._init_proxy_manager():
+            _ = self.init_credential()
+
+    def get_sheet(self, url: str, sheet_gid: int = None):
+        """
+        不传 sheet_gid，默认返回第一个 sheet
+        :param url:
+        :param sheet_gid:
+        :return:
+        """
+        with self._init_proxy_manager():
+            gc = gspread.service_account_from_dict(self.init_credential_key_dict())
+            gc.set_timeout(self.default_timeout)
+            spread_sheet = gc.open_by_url(url)
+            sheets = spread_sheet.worksheets()
+            if sheet_gid is not None:
+                for sheet in sheets:
+                    if sheet.id == sheet_gid:
+                        return sheet
+            if sheets:
+                return sheets[0]
+
+    def read_sheet_to_df(
+        self, sheet: "Worksheet", cell_range: str = None, columns: list[str] = None, dataframe_kwargs: dict = None
+    ) -> "pd.DataFrame":
+        """
+        :param sheet:
+        :param cell_range:
+            不传的话，默认读取整个 sheet
+            传的话，例如 'A1:B5'
+        :param columns:
+            不传的话，默认取 sheet 第一行作为 columns
+            传的话，例如 ['col1', 'col2', 'col3']
+        :param dataframe_kwargs:
+            pandas.DataFrame 传入的参数
+        :return:
+        """
+        with self._init_proxy_manager():  # todo: use wrapper
+            if not cell_range:
+                data = sheet.get_all_values()
+            else:
+                data = sheet.get(cell_range)
+            dataframe_kwargs = dataframe_kwargs or {}
+            if not columns:
+                df = pd.DataFrame(data[1:], columns=data[0], **dataframe_kwargs)
+            else:
+                df = pd.DataFrame(data, columns=columns, **dataframe_kwargs)
+            return df
+
+    @staticmethod
+    def parse_sheet_url(url: str) -> (str, int):
+        """输入 URL，返回 token 和 sheet id
+        :param url: https://docs.google.com/spreadsheets/d/118WyiPGFQ3ni7Gp6oNhZtkc9wEmAPfqAWynvP2ufgPk/edit#gid=1996978628
+        :return: ("118WyiPGFQ3ni7Gp6oNhZtkc9wEmAPfqAWynvP2ufgPk", "1996978628")
+        """
+        from gspread.utils import extract_id_from_url
+
+        spread_sheet_id = extract_id_from_url(url)
+        gid_pat = re.compile(r"gid=(?P<gid>\d+)")
+        gid_mobj = gid_pat.search(url)
+        sheet_id = int(gid_mobj.group("gid")) if gid_mobj else None
+        return spread_sheet_id, sheet_id
+
+    def load_df_to_sheet(self, df: DataFrame, sheet: "Worksheet", mode: str, **kwargs):
+        """write data to google sheet
+
+        Args:
+            sheet (Worksheet):
+            df (DataFrame):
+            mode (str): OVERWRITE/APPEND
+        """
+        # Determine the mode and write the data
+        headers = df.columns.values.tolist()
+        values = df.values.tolist()
+        with self._init_proxy_manager():
+            self.load_values_to_sheet(headers, values, sheet, mode, **kwargs)
+
+    @staticmethod
+    def load_values_to_sheet(headers: list[str], values: list[list], sheet: "Worksheet", mode: str, **kwargs):
+        if mode == LoadMode.OVERWRITE:
+            sheet.clear()
+            sheet.update(
+                [
+                    headers,
+                ]
+                + values,
+                **kwargs,
+            )
+        elif mode == LoadMode.APPEND:
+            existing_rows = sheet.get_all_values()
+            next_row = len(existing_rows) + 1
+            sheet.insert_rows(values, row=next_row, **kwargs)

recurvedata/connectors/connectors/hive.py

@@ -0,0 +1,207 @@
+from dataclasses import dataclass
+from typing import Optional, Type
+
+from sqlalchemy.engine.url import URL
+
+from recurvedata.connectors import get_connection_class
+from recurvedata.connectors._register import register_connector_class
+from recurvedata.connectors.connectors.mysql import CONNECTION_TYPE as MYSQL_CONNECTION_TYPE
+from recurvedata.connectors.connectors.postgres import CONNECTION_TYPE as POSTGRES_CONNECTION_TYPE
+from recurvedata.connectors.dbapi import DBAPIBase
+from recurvedata.consts import ConnectionCategory, ConnectorGroup
+from recurvedata.core.translation import _l
+
+HIVE_FIELD_DELIMITER = chr(1)
+HIVE_ARRAY_DELIMITER = chr(2)
+HIVE_MAP_ITEM_DELIMITER = chr(2)
+HIVE_MAP_KV_DELIMITER = chr(3)
+HIVE_NULL = r"\N"
+
+CONNECTION_TYPE = "hive"
+UI_CONNECTION_TYPE = "Apache Hive"
+
+
+@register_connector_class([CONNECTION_TYPE, UI_CONNECTION_TYPE])
+class HiveConnector(DBAPIBase):
+    SYSTEM_DATABASES = [
+        "information_schema",
+    ]
+    connection_type = CONNECTION_TYPE
+    ui_connection_type = UI_CONNECTION_TYPE
+    setup_extras_require = ["PyHive", "thrift-sasl"]
+    driver = "hive"
+    valid_metastore_types = [
+        MYSQL_CONNECTION_TYPE,
+        POSTGRES_CONNECTION_TYPE,
+    ]
+    category = [ConnectionCategory.WAREHOUSE]
+    group = [ConnectorGroup.DESTINATION]
+
+    config_schema = {
+        "type": "object",
+        "properties": {
+            "host": {
+                "type": "string",
+                "title": _l("Host Address"),
+                "default": "127.0.0.1",
+            },
+            "port": {
+                "type": "number",
+                "title": _l("Port Number"),
+                "default": 10000,
+            },
+            "user": {"type": "string", "title": _l("Username")},
+            "password": {"type": "string", "title": _l("Password")},
+            "database": {
+                "type": "string",
+                "title": _l("Database Name"),
+                "description": _l("The name of the database to connect to"),
+                "default": "default",
+            },
+            "hdfs_options": {
+                "type": "object",
+                "title": _l("HDFS Options"),
+                "description": _l("Configuration options for HDFS connection"),
+                "properties": {
+                    "host": {
+                        "type": "string",
+                        "title": _l("Host Address"),
+                        "description": _l("HDFS namenode hostname or IP address"),
+                    },
+                    "port": {
+                        "type": "number",
+                        "title": _l("Port Number"),
+                        "description": _l("HDFS namenode port number"),
+                        "default": 50070,
+                    },
+                    "user": {"type": "string", "title": _l("Username")},
+                    "staging_folder": {
+                        "type": "string",
+                        "title": _l("Transfer Staging Folder"),
+                        "description": _l("Temporary HDFS directory path for data transfer staging"),
+                        "default": "/tmp/recurve",
+                    },
+                },
+                "order": ["host", "port", "user", "staging_folder"],
+            },
+            "auth": {
+                "type": "string",
+                "title": _l("Authentication Type"),
+                "default": "LDAP",
+            },
+            "hive_conf": {
+                "type": "object",
+                "title": _l("Hive Execute Configurations"),
+                "description": _l("Additional Hive execution parameters"),
+                "properties": {
+                    "spark.yarn.queue": {
+                        "type": "string",
+                        "title": _l("Hive On Spark Queue"),
+                        "description": _l("YARN queue name for Spark execution"),
+                    },
+                    "tez.queue.name": {
+                        "type": "string",
+                        "title": _l("Hive On Tez Queue"),
+                        "description": _l("YARN queue name for Tez execution"),
+                    },
+                },
+                "order": ["spark.yarn.queue", "tez.queue.name"],
+            },
+            # 'metastore': {
+            #     'title': 'Hive Metastore Config',
+            #     'type': 'object',
+            #     'properties': {
+            #         'type': {
+            #             'type': 'string',
+            #             'title': 'Metastore Type',
+            #             'default': MYSQL_CONNECTION_TYPE,
+            #         },
+            #         'host': {
+            #             'type': 'string',
+            #             'title': 'Metastore Host Address',
+            #         },
+            #         'user': {
+            #             'type': 'string',
+            #             'title': 'Metastore User Name',
+            #         },
+            #         'password': {
+            #             'type': 'string',
+            #             'title': 'Metastore Password',
+            #         },
+            #         'database': {
+            #             'type': 'string',
+            #             'title': 'Metastore Database Name',
+            #         },
+            #         'port': {
+            #             'type': 'number',
+            #             'title': 'Metastore Port Number',
+            #         },
+            #     },
+            #     "order": ['host', 'port', 'user', 'password', 'database'],
+            #     'secret': ['password'],
+            # },
+            # 'ssh_tunnel': SSH_TUNNEL_CONFIG_SCHEMA,
+        },
+        "order": [
+            "host",
+            "port",
+            "user",
+            "password",
+            "database",
+            "hdfs_options",
+            "auth",
+            "hive_conf",
+        ],
+        "required": ["host", "port"],
+        "secret": ["password"],
+    }
+
+    @property
+    def connect_args(self):
+        return {"auth": "LDAP"}  # todo
+
+    # generate_ddl todo: stored as parquet
+
+    def _extract_column_name(self, column_type):
+        visit_type = column_type.__visit_name__
+        if visit_type == "type_decorator":
+            return column_type.impl.__visit_name__
+        return visit_type
+
+    @property
+    def metastore_connector(self) -> Optional[DBAPIBase]:
+        if not self.metastore:
+            return None
+        metastore_config = MetastoreConfig(**self.metastore)
+        return metastore_config.get_connector(self.conf.get("ssh_tunnel"))
+
+    @property
+    def sqlalchemy_url(self):
+        host, port = self.host, self.port
+        if self.ssh_tunnel and self.ssh_tunnel.is_active:
+            host, port = self.ssh_tunnel.local_bind_host, self.ssh_tunnel.local_bind_port
+
+        return URL(self.driver, self.user, self.password, host, port, self.database, query={"auth": self.auth})
+
+
+@dataclass
+class MetastoreConfig:
+    type: str
+    host: str
+    user: str
+    password: str
+    database: str
+    port: int
+
+    def get_connector(self, ssh_tunnel_config: Optional[dict]) -> DBAPIBase:
+        con_cls: Type[DBAPIBase] = get_connection_class(self.type)
+        return con_cls(
+            conf={
+                "host": self.host,
+                "user": self.user,
+                "password": self.password,
+                "database": self.database,
+                "port": self.port,
+                "ssh_tunnel": ssh_tunnel_config,
+            }
+        )