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

recurvedata/pigeon/connector/__init__.py

@@ -0,0 +1,294 @@
+from functools import partial
+
+from recurvedata.pigeon.connector._registry import get_connector_class
+
+
+def new_azure_synapse_connector(connection=None, database=None):
+    """
+    only connection string accepted
+    database switching between  azure data warehouses is not allowed.
+    """
+    from .azure_synapse import AzureSynapseConnector
+
+    conf = connection.copy()
+    if database:
+        conf["database"] = database
+    return AzureSynapseConnector(**conf)
+
+
+def new_azure_blob_connector(
+    conn_string: str = None,
+    account_url: str = None,
+    endpoint_suffix: str = "core.chinacloudapi.cn",
+    account_name: str = None,
+    sas_token: str = None,
+    **kwargs,
+):
+    """only connection string accepted"""
+    from .azure_blob import AzureBlobConnector
+
+    return AzureBlobConnector(
+        connection_string=conn_string,
+        account_url=account_url,
+        endpoint_suffix=endpoint_suffix,
+        account_name=account_name,
+        sas_token=sas_token,
+        **kwargs,
+    )
+
+
+def new_mysql_connector(connection=None, database=None, **kwargs):
+    """Factory function to create a new MySQLConnector.
+
+    :param connection: the connection properties,
+    :type connection: dict
+    :param database: the optional database name
+    :type database: str
+    """
+    from .mysql import MySQLConnector
+
+    conf = connection.copy()
+    return MySQLConnector(database=database, **conf)
+
+
+def new_tidb_connector(connection=None, database=None):
+    """Factory function to create a new TiDBConnector (MySQLConnector).
+
+    Similar to new_mysql_connector, but with different default connection parameter.
+
+    :param connection: the connection properties
+    :type connection: dict
+    :param database: the optional database name
+    :type database: str
+    """
+    return new_mysql_connector(connection, database)
+
+
+def new_starrocks_connector(connection=None, database=None):
+    """Factory function to create a new StarRocksConnector.
+
+    :type connection: dict
+    :param database: the optional database name
+    :type database: str
+    """
+    from .starrocks import StarRocksConnector
+
+    return StarRocksConnector(database=database, **connection)
+
+
+def new_hive_connector(connection=None, database=None, **kwargs):
+    """Factory function to create a new HiveConnector.
+
+    :param connection: the connection properties
+    :type connection: dict
+    :param database: the optional database name
+    :type database: str
+    """
+    from .hive_impala import HiveConnector
+
+    conf = connection.copy()
+    return HiveConnector(database=database, **conf)
+
+
+def new_impala_connector(connection=None, database=None, **kwargs):
+    """Factory function to create a new ImpalaConnector.
+
+    :param connection: the connection properties
+    :type connection: dict
+    :param database: the optional database name
+    :type database: str
+    """
+    from .hive_impala import ImpalaConnector
+
+    conf = connection.copy()
+    return ImpalaConnector(database=database, **conf)
+
+
+def new_webhdfs_connector(conf=None, **kwargs):
+    from .hdfs import HDFSConnector
+
+    conf = conf.copy()
+    return HDFSConnector(**conf)
+
+
+def new_redshift_connector(connection=None, database=None):
+    """Factory function to create a new RedshiftConnector.
+
+    :param connection: the connection properties
+    :type connection: dict
+    :param database: the optional database name
+    :type database: str
+    """
+    from .redshift import RedshiftConnector
+
+    conf = connection.copy()
+    return RedshiftConnector(database=database, **conf)
+
+
+def new_postgresql_connector(connection=None, database=None):
+    """Factory function to create a new PostgresConnector.
+
+    :param connection: the connection properties
+    :type connection: dict
+    :param database: the optional database name
+    :type database: str
+    """
+    from .postgresql import PostgresConnector
+
+    conf = connection.copy()
+    if database is not None:
+        conf["database"] = database
+    return PostgresConnector(**conf)
+
+
+def new_cassandra_connector(connection, database=None):
+    """Factory function to create a new CassandraConnector
+
+    :param connection: the connection properties
+    :type connection: dict
+    :param database: the optional database name
+    :type database: str
+    """
+    from .cass import CassandraConnector
+
+    return CassandraConnector(database=database, **connection)
+
+
+def new_s3_connector(conf=None):
+    from .awss3 import S3Connector
+
+    conf = conf.copy()
+    return S3Connector(**conf)
+
+
+def new_elasticsearch_connector(conf=None):
+    from .es import ElasticSearchConnector
+
+    return ElasticSearchConnector(**conf)
+
+
+def new_ftp_connector(conf=None):
+    from .ftp import FtpConnector
+
+    conf = (conf or {}).copy()
+    return FtpConnector(**conf)
+
+
+def new_sftp_connector(conf):
+    from .sftp import SFtpConnector
+
+    return SFtpConnector(**conf)
+
+
+def new_mssql_connector(connection=None, database=None, is_azure=False):
+    from .mssql import AzureSQLServerConnector, MSSQLConnector
+
+    conf = connection.copy()
+    if database:
+        conf["database"] = database
+    if is_azure:
+        connector_cls = AzureSQLServerConnector
+    else:
+        connector_cls = MSSQLConnector
+    return connector_cls(**conf)
+
+
+def new_clickhouse_connector(connection=None, database=None, native=True):
+    conf = connection.copy()
+    if not native:
+        from .clickhouse import ClickHouseConnector
+    else:
+        from .clickhouse_native import ClickHouseConnector
+    return ClickHouseConnector(database=database, **conf)
+
+
+def new_phoenix_connector(connection=None, **kwargs):
+    from .hbase_phoenix import PhoenixConnector
+
+    conf = connection.copy()
+    return PhoenixConnector(**conf)
+
+
+def new_mongodb_connector(connection=None, **kwargs):
+    from .mongodb import MongoDBConnector
+
+    conf = connection.copy()
+    return MongoDBConnector(**conf)
+
+
+def new_google_bigquery_connector(*args, **kwargs):
+    from .google_bigquery import GoogleBigqueryConnector
+
+    return GoogleBigqueryConnector(*args, **kwargs)
+
+
+def new_feishu_connector(app_id=None, app_secret=None):
+    from .feishu import FeishuBot
+
+    conf = {}
+    if app_id:
+        conf["app_id"] = app_id
+        conf["app_secret"] = app_secret
+    return FeishuBot(**conf)
+
+
+def new_owncloud_connector(url: str = None, user: str = None, password: str = None, **kwargs):
+    from .owncloud import OwncloudConnector
+
+    conf = {}
+    if url and user and password:
+        conf["url"] = url
+        conf["user"] = user
+        conf["password"] = password
+        conf.update(kwargs)
+    else:
+        raise ValueError("You must provide owncloud URL, user and password.")
+    return OwncloudConnector(**conf)
+
+
+def new_sqlite_connector(in_memory: bool, max_memory_gb: int = 2, **kwargs):
+    if not in_memory:
+        raise ValueError("Currently only supports in-memory database.")
+    from .sqlite import SQLiteMemoryDbConnector
+
+    conf = {}
+    conf.update(kwargs)
+    return SQLiteMemoryDbConnector(max_memory_gb=max_memory_gb, **conf)
+
+
+def new_doris_connector(connection=None, database=None):
+    from .doris import DorisConnector
+
+    conf = connection.copy()
+    return DorisConnector(database=database, **conf)
+
+
+_factory_registry = {
+    "mysql": new_mysql_connector,
+    "tidb": new_tidb_connector,
+    "hive": new_hive_connector,
+    "impala": new_impala_connector,
+    "redshift": new_redshift_connector,
+    "cassandra": new_cassandra_connector,
+    "s3": new_s3_connector,
+    "elasticsearch": new_elasticsearch_connector,
+    "es": new_elasticsearch_connector,
+    "ftp": new_ftp_connector,
+    "azure_synapse": new_azure_synapse_connector,
+    "azure_blob": new_azure_blob_connector,
+    "mssql": new_mssql_connector,
+    "clickhouse": new_clickhouse_connector,
+    "clickhouse_native": partial(new_clickhouse_connector, native=True),
+    "phoenix": new_phoenix_connector,
+    "mongodb": new_mongodb_connector,
+    "gbq": new_google_bigquery_connector,
+    "google_bigquery": new_google_bigquery_connector,
+    "sqlite": new_sqlite_connector,
+    "postgres": new_postgresql_connector,
+    "doris": new_doris_connector,
+    "starrocks": new_starrocks_connector,
+}
+
+
+def get_connector(db_type, *args, **kwargs):
+    return _factory_registry[db_type](*args, **kwargs)

recurvedata/pigeon/connector/_registry.py

@@ -0,0 +1,17 @@
+from recurvedata.pigeon.utils import ensure_str_list
+
+_registry = {}
+
+
+class register_connector_class(object):
+    def __init__(self, ctype):
+        self.ctype = ensure_str_list(ctype)
+
+    def __call__(self, connector):
+        for t in self.ctype:
+            _registry[t] = connector
+        return connector
+
+
+def get_connector_class(ctype):
+    return _registry[ctype]

recurvedata/pigeon/connector/aliyun_oss.py

@@ -0,0 +1,80 @@
+import logging
+import os
+
+import oss2
+
+from recurvedata.pigeon.connector._registry import register_connector_class
+from recurvedata.pigeon.utils.timing import ProgressCallback
+
+
+@register_connector_class("oss")
+class OSSBucketConnector(object):
+    def __init__(self, access_key_id, access_key_secret, endpoint, bucket_name, **kwargs):
+        self.access_key_id = access_key_id
+        self.access_key_secret = access_key_secret
+        self.endpoint = endpoint
+        self.bucket_name = bucket_name
+
+        if not all((self.access_key_id, self.access_key_secret)):
+            logging.info("access_key_id or access_key_secret is missing, fallback to ")
+            self._auth = oss2.AnonymousAuth()
+        else:
+            self._auth = oss2.make_auth(self.access_key_id, self.access_key_secret)
+
+        self.bucket = oss2.Bucket(self._auth, self.endpoint, self.bucket_name)
+        proxies = kwargs.get("proxies")
+        if proxies:
+            # pass proxies to the underlying requests.Session
+            logging.info("use %s as proxies", proxies)
+            self.bucket.session.session.proxies = proxies
+
+    def has_object(self, key):
+        return self.bucket.object_exists(key)
+
+    def delete_key(self, key):
+        self.bucket.delete_object(key)
+
+    def delete_keys_by_prefix(self, prefix):
+        keys = []
+        batch_size = 100
+        for obj in oss2.ObjectIteratorV2(bucket=self.bucket, prefix=prefix):
+            keys.append(obj.key)
+            if len(keys) >= batch_size:
+                self.bucket.batch_delete_objects(keys)
+                keys = []
+        if keys:
+            self.bucket.batch_delete_objects(keys)
+
+    def get_keys(self, prefix="", delimiter=""):
+        keys = [x.key for x in oss2.ObjectIteratorV2(bucket=self.bucket, prefix=prefix, delimiter=delimiter)]
+        if delimiter:
+            keys = [x for x in keys if not x.endswith(delimiter)]
+
+        return keys
+
+    def upload(self, filename, key=None, folder=None, overwrite=True, num_threads=4, **kwargs):
+        if not key:
+            key = os.path.basename(filename)
+        if folder:
+            key = os.path.join(folder, key)
+
+        if not overwrite:
+            if self.has_object(key=key):
+                return key
+
+        oss2.resumable_upload(self.bucket, key, filename, progress_callback=ProgressCallback(), num_threads=num_threads)
+        return key
+
+    def download(self, key, folder=None, filename=None, overwrite=True, num_threads=4, **kwargs):
+        if not filename:
+            filename = os.path.basename(key)
+        if folder:
+            filename = os.path.join(folder, filename)
+
+        if not overwrite and os.path.exists(filename):
+            return filename
+
+        oss2.resumable_download(
+            self.bucket, key, filename, progress_callback=ProgressCallback(), num_threads=num_threads
+        )
+        return filename

recurvedata/pigeon/connector/awss3.py

@@ -0,0 +1,123 @@
+import os
+import urllib.parse
+
+import boto3
+import botocore.exceptions
+from botocore.config import Config
+
+from recurvedata.pigeon.connector._registry import register_connector_class
+from recurvedata.pigeon.utils.timing import DisplayProgress
+
+
+@register_connector_class("s3")
+class S3Connector(object):
+    def __init__(self, aws_access_key_id, aws_secret_access_key, region="cn-north-1", proxies=None, **kwargs):
+        self.aws_access_key_id = aws_access_key_id
+        self.aws_secret_access_key = aws_secret_access_key
+        self.region = region
+
+        self.s3 = boto3.resource(
+            "s3",
+            region_name=self.region,
+            aws_access_key_id=aws_access_key_id,
+            aws_secret_access_key=aws_secret_access_key,
+            config=Config(proxies=proxies),
+        )
+
+    def create_bucket(self, bucket_name):
+        return self.s3.create_bucket(Bucket=bucket_name, CreateBucketConfiguration={"LocationConstraint": self.region})
+
+    def has_bucket(self, bucket_name):
+        exists = True
+        try:
+            self.s3.meta.client.head_bucket(Bucket=bucket_name)
+        except botocore.exceptions.ClientError as e:
+            error_code = int(e.response["Error"]["Code"])
+            if error_code == 404:
+                exists = False
+        return exists
+
+    def get_bucket(self, bucket_name):
+        if self.has_bucket(bucket_name):
+            return self.s3.Bucket(bucket_name)
+        return self.create_bucket(bucket_name)
+
+    def delete_bucket(self, bucket_name):
+        bucket = self.get_bucket(bucket_name)
+        for key in bucket.objects.all():
+            key.delete()
+        bucket.delete()
+
+    def has_object(self, bucket_name, key):
+        exists = True
+        try:
+            self.s3.meta.client.head_object(Bucket=bucket_name, Key=key)
+        except botocore.exceptions.ClientError as e:
+            error_code = int(e.response["Error"]["Code"])
+            if error_code == 404:
+                exists = False
+        return exists
+
+    @staticmethod
+    def parse_s3_url(s3url):
+        parsed_url = urllib.parse.urlparse(s3url)
+        if not parsed_url.netloc:
+            raise ValueError("Please provide a bucket_name")
+
+        bucket_name = parsed_url.netloc
+        key = parsed_url.path.strip("/")
+        return bucket_name, key
+
+    def delete_key(self, key, bucket_name=None):
+        if bucket_name is None:
+            bucket_name, key = self.parse_s3_url(key)
+        bucket = self.get_bucket(bucket_name)
+        bucket.Object(key).delete()
+
+    def delete_keys_by_prefix(self, bucket_name, prefix):
+        bucket = self.get_bucket(bucket_name)
+        for key in bucket.objects.filter(Prefix=prefix):
+            key.delete()
+
+    def get_keys(self, bucket_name, prefix=None):
+        bucket = self.get_bucket(bucket_name)
+        if prefix is not None:
+            all_keys = bucket.objects.filter(Prefix=prefix)
+        else:
+            all_keys = bucket.objects.all()
+
+        return [x.key for x in all_keys]
+
+    def upload(self, bucket_name, filename, key=None, folder=None, overwrite=True, **kwargs):
+        if not key:
+            key = os.path.basename(filename)
+        if folder:
+            key = os.path.join(folder, key)
+
+        if not overwrite:
+            if self.has_object(bucket_name=bucket_name, key=key):
+                return key
+
+        size = os.path.getsize(filename)
+        bucket = self.get_bucket(bucket_name)
+        with open(filename, "rb") as data:
+            bucket.upload_fileobj(data, key, Callback=DisplayProgress(size), **kwargs)
+        return key
+
+    def download(self, bucket_name, key, folder=None, filename=None, overwrite=True, **kwargs):
+        if not self.has_object(bucket_name, key):
+            raise ValueError(f"{key} not exists in {bucket_name}")
+
+        if not filename:
+            filename = os.path.basename(key)
+        if folder:
+            filename = os.path.join(folder, filename)
+
+        if not overwrite and os.path.exists(filename):
+            return filename
+
+        size = float(self.s3.meta.client.head_object(Bucket=bucket_name, Key=key)["ContentLength"])
+        bucket = self.get_bucket(bucket_name)
+        with open(filename, "wb") as data:
+            bucket.download_fileobj(key, data, Callback=DisplayProgress(size), **kwargs)
+        return filename

recurvedata/pigeon/connector/azure_blob.py

@@ -0,0 +1,176 @@
+import logging
+import os
+from typing import List, Optional
+
+from azure.core.exceptions import ResourceExistsError
+from azure.identity import AzureAuthorityHosts, ClientSecretCredential
+from azure.storage.blob import BlobServiceClient, StorageStreamDownloader
+
+from recurvedata.pigeon.connector._registry import register_connector_class
+from recurvedata.pigeon.utils.timing import DisplayProgress
+
+logger = logging.getLogger(__name__)
+
+
+@register_connector_class("azure_blob")
+class AzureBlobConnector:
+    """Connector for Azure Blob Storage.
+
+    Four ways to config:
+    - using connection_string
+    - using account_url + sas_token
+    - using endpoint_suffix + account_name + sas_token
+    - using endpoint_suffix + account_name + tenant_id + client_id + client_secret
+
+    kwargs:
+        spn_authority_host: authority host for spn, default is AzureAuthorityHosts.AZURE_CHINA
+    """
+
+    def __init__(
+        self,
+        connection_string: str = None,
+        account_url: str = None,
+        endpoint_suffix: str = "core.chinacloudapi.cn",
+        account_name: str = None,
+        sas_token: str = None,
+        tenant_id: str = None,
+        client_id: str = None,
+        client_secret: str = None,
+        **kwargs,
+    ):
+        self.conn_string = connection_string
+        self.account_url = account_url
+        self.endpoint_suffix = endpoint_suffix
+        self.account_name = account_name
+        self.sas_token = sas_token
+        self.kwargs = kwargs
+        self.spn_authority_host = self.kwargs.get("spn_authority_host") or AzureAuthorityHosts.AZURE_CHINA
+
+        authorize_by_conn_string = False
+        authorize_by_sas_token = False
+        authorize_by_spn_secret = False
+        if connection_string:
+            authorize_by_conn_string = True
+        if sas_token and (account_url or all((account_name, endpoint_suffix))):
+            authorize_by_sas_token = True
+        if tenant_id and client_id and client_secret and all((account_name, endpoint_suffix)):
+            authorize_by_spn_secret = True
+
+        if not any((authorize_by_conn_string, authorize_by_sas_token, authorize_by_spn_secret)):
+            raise ValueError(
+                """
+                    invalid authorization info
+                    Four ways to config:
+                    - using connection_string
+                    - using account_url + sas_token
+                    - using endpoint_suffix + account_name + sas_token
+                    - using endpoint_suffix + account_name + tenant_id + client_id + client_secret
+                """
+            )
+
+        if authorize_by_conn_string:
+            self.blob_service = BlobServiceClient.from_connection_string(connection_string, **kwargs)
+        elif authorize_by_sas_token:
+            if not account_url:
+                account_url = f"https://{account_name}.blob.{endpoint_suffix}"
+            self.blob_service = BlobServiceClient(account_url, credential=sas_token, **kwargs)
+        else:
+            credential = ClientSecretCredential(tenant_id, client_id, client_secret, authority=self.spn_authority_host)
+            account_url = f"https://{account_name}.blob.{endpoint_suffix}"
+            self.blob_service = BlobServiceClient(account_url=account_url, credential=credential)
+
+    @property
+    def account_key(self) -> Optional[str]:
+        if not self.conn_string:
+            return None
+        kvs = self.parse_conn_string(self.conn_string)
+        return kvs["accountkey"]
+
+    @staticmethod
+    def parse_conn_string(conn_string: str):
+        parts = conn_string.strip(";").split(";")
+        kvs = {}
+        for p in parts:
+            k, v = p.split("=", 1)
+            kvs[k.lower()] = v
+        return kvs
+
+    def get_url(self, container: str, blob: str) -> str:
+        return f"https://{self.blob_service.primary_hostname}/{container}/{blob}"
+
+    def create_container(self, container_name: str, exist_ok=True):
+        """create container"""
+        try:
+            return self.blob_service.create_container(container_name)
+        except ResourceExistsError as e:
+            if exist_ok:
+                logger.info(f"container {container_name} already exists, skip")
+            else:
+                raise e
+
+    def delete_container(self, container_name: str, **kwargs):
+        """if container not exists, error will be suppressed with the fail_not_exist parameter"""
+        self.blob_service.delete_container(container_name, **kwargs)
+
+    def exists(self, container_name: str, blob_name: str = None, **kwargs) -> bool:
+        """
+        if blob name is none, check whether container exists or not
+        if blob name specified, check blob exists or not in the  container
+        """
+        if blob_name is None:
+            client = self.blob_service.get_container_client(container_name)
+        else:
+            client = self.blob_service.get_blob_client(container_name, blob_name)
+        return client.exists(**kwargs)
+
+    def delete_blob(self, container_name, blob_name, **kwargs):
+        container = self.blob_service.get_container_client(container_name)
+        container.delete_blob(blob_name, **kwargs)
+
+    def list_blobs(self, container_name, name_starts_with=None, include=None, **kwargs) -> List[str]:
+        container = self.blob_service.get_container_client(container_name)
+        generator = container.list_blobs(name_starts_with=name_starts_with, include=include, **kwargs)
+        return [blob.name for blob in generator]
+
+    def upload(self, container_name, local_file_path, blob_name=None, overwrite=True, is_progress_hook=True, **kwargs):
+        """
+        Upload local file to container with specified blob name.
+        The specified container will also be created if not exists.
+        """
+        if not blob_name:
+            blob_name = os.path.basename(local_file_path)
+
+        # container_blob = f'{container_name}/{blob_name}'
+        blob = self.blob_service.get_blob_client(container_name, blob_name)
+        if not overwrite and blob.exists():
+            logger.info("Blob exists, skip!")
+            return blob_name
+
+        size = os.path.getsize(local_file_path)
+        options = {"overwrite": True, "max_concurrency": 4}
+        if is_progress_hook:
+            options["progress_hook"] = DisplayProgress(size, stream=False)
+
+        options.update(kwargs)
+        with open(local_file_path, "rb") as data:
+            blob.upload_blob(data, **options)
+        return blob_name
+
+    def download(self, container_name, blob_name, local_file_path, **kwargs):
+        """download blob to local"""
+        blob = self.blob_service.get_blob_client(container_name, blob_name)
+        size = blob.get_blob_properties().size
+        if size == 0:
+            logging.warning("blob %s has no content, create an empty file and exit", blob_name)
+            with open(local_file_path, "w"):
+                return
+
+        options = {
+            "max_concurrency": 4,
+            "progress_hook": DisplayProgress(size, stream=False),
+        }
+        options.update(kwargs)
+        with open(local_file_path, "wb") as f:
+            data: StorageStreamDownloader = blob.download_blob(**options)
+            data.readinto(f)
+        return local_file_path