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

recurvedata/pigeon/connector/ftp.py

@@ -0,0 +1,163 @@
+import datetime
+import ftplib
+import logging
+import os
+import shutil
+import time
+
+import humanize
+
+from recurvedata.pigeon.connector._registry import register_connector_class
+from recurvedata.pigeon.utils import LoggingMixin
+
+
+@register_connector_class("ftp")
+class FtpConnector(LoggingMixin):
+    def __init__(self, host, user, password, **kwargs):
+        self.host = host
+        self.username = user
+        self.password = password
+        self.port = kwargs.pop("port", ftplib.FTP_PORT)
+
+        # The ftplib.FTP is stupid, it doesn't support the port parameter in the constructor
+        # so we have to use the connect method to specify the port
+        if self.port == ftplib.FTP_PORT:
+            self.ftp = ftplib.FTP(host=self.host, user=self.username, passwd=self.password, **kwargs)
+        else:
+            self.ftp = ftplib.FTP(**kwargs)
+            self.ftp.connect(self.host, self.port)
+            if user:
+                self.ftp.login(user=user, passwd=password, acct=kwargs.get("acct", ""))
+
+    def close(self):
+        ftp = self.ftp
+        ftp.quit()
+        self.ftp = None
+
+    def list_dir(self, path):
+        try:
+            return self.ftp.nlst(path)
+        except ftplib.all_errors:
+            return []
+
+    def rename(self, from_name, to_name):
+        return self.ftp.rename(from_name, to_name)
+
+    def makedir(self, path):
+        self.ftp.mkd(path)
+
+    def rmdir(self, path):
+        self.ftp.rmd(path)
+
+    def rm(self, name):
+        self.ftp.delete(name)
+
+    def pwd(self):
+        return self.ftp.pwd()
+
+    def size(self, name):
+        return self.ftp.size(name)
+
+    def is_ftp_dir(self, path):
+        original_cwd = self.pwd()
+        try:
+            self.ftp.cwd(path)
+            self.ftp.cwd(original_cwd)
+            return True
+        except ftplib.all_errors:
+            return False
+
+    def download_file(self, src_file, dst_file):
+        exists = True
+        local_dir = os.path.dirname(dst_file)
+        if not os.path.exists(local_dir):
+            exists = False
+            os.makedirs(local_dir)
+        try:
+            total_bytes = self.size(src_file)
+            with open(dst_file, "wb") as f:
+                writer = StatsReaderWriter(f, total_bytes)
+                self.ftp.retrbinary(f"RETR {src_file}", writer.write)
+            writer.show_stat()
+            self.logger.info(f"successfully downloaded {src_file} to {dst_file}")
+        except ftplib.all_errors as e:
+            os.unlink(dst_file)
+            if not exists:
+                shutil.rmtree(local_dir)
+            self.logger.exception(f"failed to download {src_file}")
+            raise e
+
+    def upload_file(self, src_file, dst_file):
+        if not os.path.isfile(src_file):
+            raise ValueError(f"{src_file} is not a file")
+        try:
+            total_bytes = os.stat(src_file).st_size
+            with open(src_file, "rb") as f:
+                reader = StatsReaderWriter(f, total_bytes)
+                self.ftp.storbinary(f"STOR {dst_file}", reader)
+            reader.show_stat()
+            self.logger.info(f"successfully uploaded {src_file} to {dst_file}")
+        except ftplib.all_errors as e:
+            self.logger.exception(f"failed to upload {src_file}")
+            raise e
+
+
+class StatsReaderWriter(object):
+    def __init__(self, fp, total_bytes, show_stats_bytes=1024 * 1024):
+        self.fp = fp
+        self.total_bytes = total_bytes
+        self.show_stats_bytes = show_stats_bytes
+
+        self._transferred_bytes = 0
+        self._start_time = time.time()
+        self._end_time = 0
+
+    def read(self, n):
+        rv = self.fp.read(n)
+        self._incr_transferred_bytes(n)
+        return rv
+
+    def write(self, data):
+        rv = self.fp.write(data)
+        self._incr_transferred_bytes(len(data))
+        return rv
+
+    def close(self):
+        if self.fp.closed:
+            if self._end_time == 0:
+                self._end_time = time.time()
+            return
+        try:
+            self.fp.close()
+        except Exception:
+            pass
+        self._end_time = time.time()
+
+    def _incr_transferred_bytes(self, n):
+        for _ in range(n):
+            self._transferred_bytes += 1
+            if self._transferred_bytes % self.show_stats_bytes == 0:
+                self.show_stat()
+
+    def show_stat(self):
+        if self._end_time == 0:
+            end_time = time.time()
+        else:
+            end_time = self._end_time
+        duration = end_time - self._start_time
+        if duration == 0:
+            speed = 0
+        else:
+            speed = self._transferred_bytes / duration
+
+        if self.total_bytes == 0:
+            progress = 0
+        else:
+            progress = 100 * self._transferred_bytes / self.total_bytes
+        logging.info(
+            "transferred %s in %s, average speed: %s/s, progress: %.2f%%",
+            humanize.naturalsize(self._transferred_bytes, gnu=True),
+            datetime.timedelta(seconds=duration),
+            humanize.naturalsize(speed, gnu=True),
+            progress,
+        )

recurvedata/pigeon/connector/google_bigquery.py

@@ -0,0 +1,283 @@
+import copy
+import os
+from urllib import parse
+
+import cytoolz as toolz
+from google import auth
+from google.cloud.bigquery import Client, LoadJobConfig, SourceFormat, dbapi, enums, job
+from google.cloud.bigquery.dataset import DatasetReference
+from google.cloud.bigquery.table import TableReference
+from google.cloud.exceptions import BadRequest, NotFound
+from google.oauth2 import service_account
+from requests import Session
+
+from recurvedata.pigeon.connector._registry import register_connector_class
+from recurvedata.pigeon.connector.dbapi import ClosingCursor, DBAPIConnector, NullCursor
+from recurvedata.pigeon.schema import types
+
+_bigquery_type_to_canonical_type = {
+    enums.SqlTypeNames.STRING: types.STRING,
+    enums.SqlTypeNames.INT64: types.INT64,
+    enums.SqlTypeNames.INTEGER: types.INT64,
+    enums.SqlTypeNames.FLOAT: types.FLOAT64,
+    enums.SqlTypeNames.FLOAT64: types.FLOAT64,
+    enums.SqlTypeNames.NUMERIC: types.FLOAT64,
+    enums.SqlTypeNames.BOOLEAN: types.BOOLEAN,
+    enums.SqlTypeNames.BOOL: types.BOOLEAN,
+    enums.SqlTypeNames.TIMESTAMP: types.DATETIME,
+    enums.SqlTypeNames.DATETIME: types.DATETIME,
+    enums.SqlTypeNames.DATE: types.DATE,
+}
+
+GOOGLE_DRIVE_API = (
+    "https://www.googleapis.com/auth/drive"  # external table linked with google sheet, need google drive api enabled
+)
+
+
+class IterCursor(ClosingCursor):
+    def __init__(self, connection, commit_on_close=True, pagesize=None):
+        super().__init__(connection, commit_on_close)
+        self._cursor.arraysize = pagesize
+
+    def __iter__(self):
+        self._cursor._try_fetch()
+        return self._cursor._query_data
+
+
+@register_connector_class(["google_bigquery", "gbq"])
+class GoogleBigqueryConnector(DBAPIConnector):
+    _sqla_driver = "bigquery"
+
+    if GOOGLE_DRIVE_API in Client.SCOPE:
+        _scopes = Client.SCOPE
+    else:
+        _scopes = Client.SCOPE + (GOOGLE_DRIVE_API,)
+
+    def __init__(
+        self,
+        key_path: str = None,
+        key_dict: dict = None,
+        project: str = None,
+        http: Session = None,
+        proxies: dict = None,
+        location: str = None,
+        dataset: str = None,
+        pagesize: int = None,
+        *args,
+        **kwargs,
+    ):
+        """
+        instance of gbq
+        :param project: project_id
+        :param key_path: path to json key file
+        :param key_dict: dict of key
+        :param http: requests session
+        :param proxies: proxy
+        :param location: location
+        :param dataset: dataset_id
+        """
+        super().__init__(host=None, database=dataset, *args, **kwargs)
+        self._project_id = project
+        self._key_path = key_path
+        self._key_dict = key_dict
+        self._http = http
+        self._proxies = proxies
+        self._location = location
+        self.dataset = dataset
+        self.pagesize = pagesize
+
+    def is_google_bigquery(self):
+        return True
+
+    def get_credentials(self):
+        if not any([self._key_path, self._key_dict]):
+            credentials, project_id = auth.default(scopes=self._scopes, request=self._http)
+            self._project_id = self._project_id or project_id
+        elif self._key_path:
+            credentials = service_account.Credentials.from_service_account_file(
+                filename=self._key_path, scopes=self._scopes
+            )
+        else:
+            _key_dict = copy.deepcopy(self._key_dict)
+            # Fix private key format with robust conversion
+            _key_dict["private_key"] = self._convert_private_key(_key_dict["private_key"])
+            credentials = service_account.Credentials.from_service_account_info(info=_key_dict, scopes=self._scopes)
+        self._project_id = self._project_id or credentials.project_id
+
+        return credentials
+
+    @staticmethod
+    def _convert_private_key(private_key: str) -> str:
+        """
+        Convert private key from various escape formats to proper PEM format.
+        Handles multiple levels of escaping that can occur during transmission/storage.
+        """
+        if not private_key:
+            return private_key
+            
+        # Remove any leading/trailing whitespace
+        private_key = private_key.strip()
+        
+        # Handle various escape sequence patterns
+        # Multiple replacement passes to handle nested escaping
+        
+        # Replace quadruple-escaped newlines (\\\\n -> \\n)
+        private_key = private_key.replace("\\\\n", "\\n")
+        
+        # Replace double-escaped newlines (\\n -> \n) 
+        private_key = private_key.replace("\\n", "\n")
+        
+        # Handle edge case where literal \n strings need to become actual newlines
+        # This covers cases where the key was stored as a literal string
+        if "-----BEGIN PRIVATE KEY-----" in private_key and "\n" not in private_key:
+            # If we have the BEGIN marker but no actual newlines, it's likely escaped
+            private_key = private_key.replace("-----BEGIN PRIVATE KEY-----", "-----BEGIN PRIVATE KEY-----\n")
+            private_key = private_key.replace("-----END PRIVATE KEY-----", "\n-----END PRIVATE KEY-----")
+            
+            # Split the key content and add newlines every 64 characters (standard PEM format)
+            lines = private_key.split('\n')
+            if len(lines) >= 2:
+                # Extract the key content between BEGIN and END
+                begin_line = lines[0]
+                end_line = lines[-1] 
+                key_content = ''.join(lines[1:-1])
+                
+                # Split key content into 64-character lines
+                formatted_lines = [begin_line]
+                for i in range(0, len(key_content), 64):
+                    formatted_lines.append(key_content[i:i+64])
+                formatted_lines.append(end_line)
+                
+                private_key = '\n'.join(formatted_lines)
+        
+        return private_key
+
+    @toolz.memoize
+    def connect_impl(self, *args, **kwargs):
+        return dbapi.connect(client=self.client())
+
+    def client(self):
+        if self._proxies:
+            for scheme in ["http", "https"]:
+                os.environ[f"{scheme}_proxy"] = self._proxies[scheme]
+
+        client = Client(
+            project=self._project_id, credentials=self.get_credentials(), location=self._location, _http=self._http
+        )
+        return client
+
+    def cursor(self, autocommit=False, dryrun=False, commit_on_close=True, **kwargs):
+        if dryrun:
+            return NullCursor()
+        conn = self.connect(autocommit, **kwargs)
+        return IterCursor(conn, commit_on_close=commit_on_close, pagesize=self.pagesize)
+
+    def _get_sqlalchemy_uri(self):
+        params = {"location": self._location}
+        if self._key_path:
+            params.update({"credentials_path": self._key_path})
+        uri = f"{self._sqla_driver}://{self._project_id}"
+        if self.dataset:
+            uri = os.path.join(uri, self.dataset)
+        return f"{uri}?{parse.urlencode(params)}"
+
+    def get_pandas_df(self, query, parameters=None, **kwargs):
+        """Get pandas dataframe
+        Note: pd.read_gbq 无法正常工作，改用 client to_dataframe()
+        """
+        format_operation = dbapi.cursor._format_operation(query, parameters=parameters)
+        query_parameters = dbapi._helpers.to_query_parameters(parameters)
+        config = job.QueryJobConfig(use_legacy_sql=False)
+        config.query_parameters = query_parameters
+        result = self.client().query(format_operation, job_config=config).result()
+        return result.to_dataframe()
+
+    def table_ref(self, table, dataset):
+        return TableReference(DatasetReference(self._project_id, dataset), table)
+
+    def has_table(self, table, dataset=None, **kwargs):
+        if dataset is None:
+            dataset = self.dataset
+        try:
+            self.client().get_table(self.table_ref(table, dataset))
+            return True
+        except NotFound:
+            return False
+
+    def list_partitions(self, table, dataset):
+        """查询 partition keys"""
+        try:
+            return self.client().list_partitions(self.table_ref(table, dataset))
+        except BadRequest:
+            return []
+
+    @staticmethod
+    def to_canonical_type(type_code, size):
+        return _bigquery_type_to_canonical_type.get(type_code, types.STRING)
+
+    @staticmethod
+    def from_canonical_type(canonical_type, size):
+        _canonical_type_to_bigquery_type = {v: k for k, v in _bigquery_type_to_canonical_type.items()}
+        return _canonical_type_to_bigquery_type.get(canonical_type, "STRING")
+
+    def generate_ddl(self, table, dataset=None, if_exists=True):
+        cols = [f"{col.name} {col.field_type}" for col in self.get_schema(table, dataset)]
+        if_exists_stmt = " IF NOT EXISTS " if if_exists else " "
+        full_table_name = f"{self.quote_identifier(dataset)}.{self.quote_identifier(table)}"
+        return f'CREATE TABLE{if_exists_stmt}{full_table_name} ({", ".join(cols)})'
+
+    def get_columns(self, table, dataset=None, exclude=()):
+        cols = []
+        for col in self.get_schema(table, dataset):
+            if col.name in exclude:
+                continue
+            cols.append(col.name)
+        return cols
+
+    def get_schema(self, table, dataset):
+        if dataset is None:
+            dataset = self.dataset
+        if not self.has_table(table, dataset):
+            raise ValueError(f"Table {table} not exists in {dataset!r}")
+        table = self.client().get_table(self.table_ref(table, dataset))
+        return table.schema
+
+    def load_csv(
+        self,
+        table,
+        filename=None,
+        gcs_uri=None,
+        delimiter=",",
+        quotechar='"',
+        skiprows=0,
+        write_disposition="WRITE_APPEND",
+        schema=None,
+        **kwargs,
+    ):
+        dataset, table = table.split(".")
+        job_config = LoadJobConfig(
+            source_format=SourceFormat.CSV,
+            skip_leading_rows=skiprows,
+            # autodetect=True,
+            field_delimiter=delimiter,
+            quote_character=quotechar,
+            write_disposition=write_disposition,
+            schema=schema,
+            **kwargs,
+        )
+
+        if filename:
+            with open(filename, "rb") as file:
+                load_job = self.client().load_table_from_file(
+                    file_obj=file, destination=self.table_ref(table, dataset), job_config=job_config
+                )
+        elif gcs_uri:
+            load_job = self.client().load_table_from_uri(
+                source_uris=gcs_uri, destination=self.table_ref(table, dataset), job_config=job_config
+            )
+        else:
+            self.logger.error("no file or gcs uri is provided")
+
+        self.logger.info("start loading csv to bigquery")
+        load_job.result()
+        self.logger.info("finish loading csv to bigquery")

recurvedata/pigeon/connector/google_cloud_storage.py

@@ -0,0 +1,130 @@
+import copy
+import logging
+import os
+
+from google import auth
+from google.cloud import storage
+from google.oauth2 import service_account
+
+from recurvedata.pigeon.connector._registry import register_connector_class
+
+
+@register_connector_class(["google_cloud_storage", "gcs"])
+class GoogleCloudStorageConnector(object):
+    def __init__(
+        self,
+        key_path: str = None,
+        key_dict: dict = None,
+        project: str = None,
+        proxies: dict = None,
+        bucket_name: str = None,
+        **kwargs,
+    ):
+        """
+        instance of gcs
+        :param project: project_id
+        :param key_path: path to json key file
+        :param key_dict: dict of key
+        :param proxies: proxy
+        :param bucket_name: bucket's name
+        """
+        self._project_id = project
+        self._key_path = key_path
+        self._key_dict = key_dict
+        self._proxies = proxies
+        self.bucket_name = bucket_name
+
+        if not any([self._key_path, self._key_dict]):
+            # 需要配置好 GOOGLE_APPLICATION_CREDENTIALS 环境变量
+            # export GOOGLE_APPLICATION_CREDENTIALS='{service account key 文件路径}'
+            self._credentials, auth_project_id = auth.default()
+            self._project_id = self._project_id or auth_project_id
+        elif self._key_path:
+            # 传入 service account key 文件路径
+            self._credentials = service_account.Credentials.from_service_account_file(filename=self._key_path)
+        else:
+            # 传入 service account key dict
+            _key_dict = copy.deepcopy(self._key_dict)
+            _key_dict["private_key"] = _key_dict["private_key"].replace("\\n", "\n")
+            self._credentials = service_account.Credentials.from_service_account_info(info=_key_dict)
+        self._project_id = self._project_id or self._credentials.project_id
+
+        if self._proxies:
+            for scheme in ["http", "https"]:
+                os.environ[f"{scheme}_proxy"] = self._proxies[scheme]
+
+        self.client = storage.Client(project=self._project_id, credentials=self._credentials, **kwargs)
+
+    def create_bucket(self, bucket_name, location=None):
+        logging.info(f"Start creating bucket {bucket_name} at location {location}")
+        new_bucket = self.client.create_bucket(bucket_name, location=location)
+        logging.info(f"Successfully created bucket {bucket_name} at location {location}")
+        return new_bucket
+
+    def get_buckets(self):
+        buckets = self.client.list_buckets()
+        return [bucket.name for bucket in buckets]
+
+    def has_key(self, key, bucket_name=None):
+        if not bucket_name:
+            bucket_name = self.bucket_name
+
+        bucket = self.client.bucket(bucket_name)
+        return bucket.blob(key).exists()
+
+    def get_keys(self, bucket_name=None, prefix=""):
+        if not bucket_name:
+            bucket_name = self.bucket_name
+
+        keys = self.client.list_blobs(bucket_name, prefix=prefix)
+        return [key.name for key in keys]
+
+    def delete_key(self, key, bucket_name=None):
+        if not bucket_name:
+            bucket_name = self.bucket_name
+
+        bucket = self.client.bucket(bucket_name)
+        blob = bucket.blob(key)
+        logging.info(f"Start deleting storage object {key}")
+        blob.delete()
+        logging.info(f"Successfully deleted storage object {key}")
+
+    def upload(self, filename, bucket_name=None, key=None, folder=None, overwrite=True, **kwargs):
+        if not bucket_name:
+            bucket_name = self.bucket_name
+        if not key:
+            key = os.path.basename(filename)
+        if folder:
+            key = os.path.join(folder, key)
+
+        bucket = self.client.bucket(bucket_name)
+        blob = bucket.blob(key)
+
+        if not overwrite and blob.exists():
+            return key
+
+        logging.info(f"Start uploading file {filename} to {key}.")
+        blob.upload_from_filename(filename, **kwargs)
+        logging.info(f"Successfully uploaded file {filename} to {key}.")
+        return key
+
+    def download(self, key, bucket_name=None, folder=None, filename=None, overwrite=True, **kwargs):
+        if not bucket_name:
+            bucket_name = self.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
+
+        bucket = self.client.bucket(bucket_name)
+        blob = bucket.get_blob(key)
+        logging.info(f"Start downloading storage object {key} from bucket {bucket_name} to local file {filename}.")
+        logging.info(f"Size: {round(blob.size / 1024 / 1024, 2)} MB")
+        blob.download_to_filename(filename, **kwargs)
+        logging.info(
+            f"Successfully downloaded storage object {key} from bucket {bucket_name} to local file {filename}."
+        )
+        return filename

recurvedata/pigeon/connector/hbase_phoenix.py

@@ -0,0 +1,108 @@
+import phoenixdb
+from phoenixdb.cursor import Cursor
+
+from recurvedata.pigeon.connector._registry import register_connector_class
+from recurvedata.pigeon.connector.dbapi import DBAPIConnector
+from recurvedata.pigeon.schema import types
+
+# Phoenix Data Types: http://phoenix.apache.org/language/datatypes.html
+_phoenix_type_to_canonical_type = {
+    'INTEGER': types.INT32,
+    'UNSIGNED_INT': types.INT32,
+    'BIGINT': types.INT64,
+    'UNSIGNED_LONG': types.INT64,
+    'TINYINT': types.INT8,
+    'UNSIGNED_TINYINT': types.INT8,
+    'SMALLINT': types.INT16,
+    'UNSIGNED_SMALLINT': types.INT16,
+    'FLOAT': types.FLOAT32,
+    'UNSIGNED_FLOAT': types.FLOAT32,
+    'DOUBLE': types.FLOAT64,
+    'UNSIGNED_DOUBLE': types.FLOAT64,
+    'DECIMAL': types.FLOAT64,
+    'BOOLEAN': types.BOOLEAN,
+
+    'TIME': types.STRING,
+    'UNSIGNED_TIME': types.STRING,
+    'DATE': types.DATE,
+    'UNSIGNED_DATE': types.DATE,
+    'TIMESTAMP': types.DATETIME,
+    'UNSIGNED_TIMESTAMP': types.DATETIME,
+
+    'VARCHAR': types.STRING,
+    'CHAR': types.STRING,
+
+    # default: types.STRING
+}
+
+_canonical_type_to_phoenix_type = {
+    types.BOOLEAN: 'BOOLEAN',
+    types.INT8: 'TINYINT',
+    types.INT16: 'SMALLINT',
+    types.INT32: 'INTEGER',
+    types.INT64: 'BIGINT',
+    types.FLOAT32: 'FLOAT',
+    types.FLOAT64: 'DOUBLE',
+
+    types.DATE: 'DATE',
+    types.DATETIME: 'DATETIME',
+
+    types.STRING: 'VARCHAR',
+    types.JSON: 'VARCHAR',
+}
+
+
+class PhoenixCursor(Cursor):
+    itersize = 1000
+
+
+@register_connector_class(['phoenix'])
+class PhoenixConnector(DBAPIConnector):
+    _default_port = 8765
+    _identifier_start_quote = '"'
+    _identifier_end_quote = '"'
+
+    def connect_impl(self, autocommit=False, *args, **kwargs):
+        url = f'http://{self.host}:{self.port}'
+        return phoenixdb.connect(url=url, autocommit=autocommit, cursor_factory=PhoenixCursor, *args, **kwargs)
+
+    def create_engine(self, engine_kwargs=None):
+        raise NotImplementedError
+
+    def has_table(self, table, database=None, **kwargs):
+        if database is None:
+            query = 'SELECT 1 FROM system.catalog WHERE table_name = ? LIMIT 1'
+            params = [table]
+        else:
+            query = 'SELECT 1 FROM system.catalog WHERE table_name = ? AND table_schem = ? LIMIT 1'
+            params = [table, database]
+        return self.fetchone(query, params) is not None
+
+    def get_columns(self, table, database=None, exclude=None):
+        if database:
+            clause = f'table_schem = {database!r}'
+        else:
+            clause = 'table_schem IS NULL'
+        query = f'''
+            SELECT column_name FROM system.catalog
+            WHERE {clause} AND table_name = ? AND ordinal_position IS NOT NULL
+            ORDER BY ordinal_position
+        '''
+        cols = [x[0] for x in self.fetchall(query, [table])]
+        if not cols:
+            raise ValueError('Table {!r} not exists in {!r}'.format(table, database))
+
+        if exclude:
+            cols = [x for x in cols if x not in exclude]
+        return cols
+
+    def is_phoenix(self):
+        return True
+
+    @staticmethod
+    def to_canonical_type(type_code, size):
+        return _phoenix_type_to_canonical_type.get(type_code, types.STRING)
+
+    @staticmethod
+    def from_canonical_type(canonical_type, size):
+        return _canonical_type_to_phoenix_type.get(canonical_type, 'VARCHAR')