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

recurvedata/pigeon/connector/owncloud.py

@@ -0,0 +1,92 @@
+import logging
+import os
+import traceback
+
+import pandas as pd
+from owncloud import Client, HTTPResponseError
+
+from recurvedata.pigeon.connector._registry import register_connector_class
+from recurvedata.pigeon.utils.fs import new_tempfile, remove_files_safely
+
+
+class OwncloudException(Exception):
+    pass
+
+
+class OwncloudDownloadException(OwncloudException):
+    pass
+
+
+class OwncloudUploadException(OwncloudException):
+    pass
+
+
+class NewOwncloudClient(Client):
+    def get_webdav_url(self):
+        return self._webdav_url
+
+
+@register_connector_class('owncloud')
+class OwncloudConnector(object):
+    def __init__(self, url: str = None, user: str = None, password: str = None, **kwargs):
+        self.url = url
+        self.user = user
+        self.password = password
+        self.oc = NewOwncloudClient(url, **kwargs)
+        self.oc.login(user, password)
+
+    def download_file(self, remote_path: str, local_path: str):
+        logging.info(f'Downloading remote file {remote_path} to {local_path}.')
+        try:
+            status = self.oc.get_file(remote_path, local_path)
+            if status:
+                logging.info(f'Successfully download remote file {remote_path} to {local_path}.')
+            else:
+                raise OwncloudDownloadException(f'Failed to download remote file {remote_path}, unknown error.')
+        except HTTPResponseError as e:
+            logging.error(traceback.format_exc())
+            raise OwncloudDownloadException(
+                f'Failed to download remote file {remote_path}, HTTPResponseError {e.res}.'
+            )
+        return status
+
+    def upload_file(self, remote_path: str, local_source_file: str, **kwargs):
+        logging.info(f'Uploading local file {local_source_file} to {remote_path}.')
+        try:
+            res = self.oc.put_file(remote_path, local_source_file, **kwargs)
+            logging.info(f'Successfully upload local file {local_source_file} to remote {remote_path}.')
+        except Exception as e:
+            logging.error(traceback.format_exc())
+            raise OwncloudUploadException(
+                f'Failed to upload local file {local_source_file} to remote {remote_path}, {e.args}.'
+            )
+        return res
+
+    def get_pandas_df(self, remote_path: str) -> pd.DataFrame:
+        temp_file_path = new_tempfile()
+        if self.download_file(remote_path, temp_file_path):
+            file_type = os.path.splitext(remote_path)[-1]
+            try:
+                if file_type and file_type.lower() in ('.xlsx', '.xls'):
+                    df = pd.read_excel(temp_file_path)
+                elif file_type and file_type.lower() in ('.parquet', '.parq'):
+                    df = pd.read_parquet(temp_file_path)
+                elif file_type and file_type.lower() == '.json':
+                    df = pd.read_json(temp_file_path)
+                else:
+                    df = pd.read_csv(temp_file_path)
+            except Exception as e:
+                logging.error(traceback.format_exc())
+                raise ValueError(f'Failed to load remote file {remote_path} to pandas df, {e.args}.')
+            finally:
+                remove_files_safely(temp_file_path)
+            logging.info(f'Successfully load remote file {remote_path} to pandas df, {len(df)} rows.')
+            return df
+
+    @property
+    def webdav_url(self):
+        return self.oc.get_webdav_url()
+
+    @property
+    def http_auth_conf(self):
+        return {'username': f'{self.user}', 'password': f'{self.password}'}

recurvedata/pigeon/connector/postgresql.py

@@ -0,0 +1,267 @@
+import psycopg2
+
+from recurvedata.pigeon.connector._registry import register_connector_class
+from recurvedata.pigeon.connector.dbapi import ClosingCursor, DBAPIConnector, NullCursor
+from recurvedata.pigeon.schema import types
+
+_pg_type_to_canonical_type = {
+    16: types.BOOLEAN,
+    21: types.INT16,
+    23: types.INT32,
+    20: types.INT64,
+    114: types.JSON,
+    700: types.FLOAT32,
+    701: types.FLOAT64,
+    1700: types.FLOAT64,
+    1114: types.DATETIME,
+    1184: types.DATETIME,
+    1082: types.DATE,
+    1043: types.STRING,
+    1014: types.STRING,
+    1015: types.STRING,
+    1008: types.STRING,
+    1009: types.STRING,
+    2951: types.STRING,
+}
+
+canonical_type_to_pg_type = {
+    types.BOOLEAN: "BOOLEAN",
+    types.INT8: "INT2",
+    types.INT16: "INT2",
+    types.INT32: "INT4",
+    types.INT64: "INT8",
+    types.FLOAT32: "FLOAT4",
+    types.FLOAT64: "FLOAT8",
+    types.DATETIME: "TIMESTAMP",
+    types.DATE: "DATE",
+    types.STRING: "TEXT",
+    types.JSON: "JSON",
+}
+
+
+class NamedCursor(ClosingCursor):
+    """NamedCursor is a server side cursor, using DECLARE and FETCH internally
+    http://initd.org/psycopg/docs/usage.html#server-side-cursors
+    """
+
+    def __init__(self, connection, commit_on_close=True, name=None):
+        self.connection = connection
+        self._commit_on_close = commit_on_close
+        if name is not None:
+            self._cursor = connection.cursor(name, withhold=True)
+            self._cursor.itersize = 1000
+        else:
+            self._cursor = connection.cursor()
+
+
+@register_connector_class(["postgres", "postgresql"])
+class PostgresConnector(DBAPIConnector):
+    _sqla_driver = "postgresql+psycopg2"
+    _identifier_start_quote = '"'
+    _identifier_end_quote = '"'
+    _default_port = 5432
+
+    def connect_impl(self, autocommit=False, *args, **kwargs):
+        conn = psycopg2.connect(
+            host=self.host,
+            port=self.port,
+            user=self.user,
+            password=self.password,
+            database=self.database,
+            *args,
+            **kwargs,
+        )
+        conn.autocommit = autocommit
+
+        if self.schema:
+            with conn.cursor() as cursor:
+                cursor.execute(f"SET search_path TO {self.schema}, public")
+
+        return conn
+
+    def cursor(self, autocommit=False, dryrun=False, commit_on_close=True, **kwargs):
+        """Returns a DBAPI cursor"""
+        if dryrun:
+            return NullCursor()
+        cursor_name = kwargs.pop("cursor_name", None)
+        conn = self.connect(autocommit, **kwargs)
+        return NamedCursor(conn, commit_on_close=commit_on_close, name=cursor_name)
+
+    def has_table(self, table, database=None, schema="public", **kwargs):
+        schema, table = self._get_schema_table(table, schema)
+
+        if database is not None and database != self.database:
+            conn = self.clone()
+            conn.database = database
+        else:
+            conn = self
+        with conn.cursor() as cursor:
+            cursor.execute(
+                """
+                SELECT EXISTS (
+                  SELECT 1 FROM information_schema.tables
+                  WHERE table_name = %s AND table_schema = %s
+                )
+            """,
+                (table, schema),
+            )
+            return bool(cursor.fetchone()[0])
+
+    def get_columns(self, table, schema="public", database=None):
+        schema, table = self._get_schema_table(table, schema)
+        if database is None:
+            database = self.database
+        if not self.has_table(table, database, schema=schema):
+            raise ValueError("Table {!r}.{!r} not exists in {!r}".format(schema, table, database))
+        with self.cursor() as cursor:
+            cursor.execute('SELECT * FROM "{}"."{}" LIMIT 0'.format(schema, table))
+            cursor.fetchall()
+            return [x[0] for x in cursor.description]
+
+    def generate_ddl(self, table, schema="public", database=None, field_filter=(), if_exists=True):
+        schema, table = self._get_schema_table(table, schema)
+        if database is None:
+            database = self.database
+        if not self.has_table(table, database, schema=schema):
+            raise ValueError(f"Table {schema!r}.{table!r} not exists in {database!r}")
+
+        with self.cursor() as cursor:
+            # get table comment
+            tbl_comment_sql = f"""
+            SELECT pgd.description AS table_comment
+            FROM pg_catalog.pg_description pgd
+            WHERE pgd.objsubid = 0 AND pgd.objoid = (SELECT c.oid
+                                                     FROM pg_catalog.pg_class c
+                                                         LEFT JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
+                                                     WHERE n.nspname = {schema!r}
+                                                           AND c.relname = {table!r} AND
+                                                           c.relkind IN ('r', 'v', 'm', 'f'));
+            """
+            cursor.execute(tbl_comment_sql)
+            t_comment = cursor.fetchall()
+            # get columns
+            col_comment_sql = f"""
+            SELECT
+                a.attname                                       AS "field",
+                pg_catalog.format_type(a.atttypid, a.atttypmod) AS "type",
+                (SELECT pg_catalog.pg_get_expr(d.adbin, d.adrelid)
+                 FROM pg_catalog.pg_attrdef d
+                 WHERE d.adrelid = a.attrelid AND d.adnum = a.attnum
+                       AND a.atthasdef)
+                                                                AS "default",
+                a.attnotnull                                    AS "isnull",
+                pgd.description                                 AS "comment"
+            FROM pg_catalog.pg_attribute a
+                LEFT JOIN pg_catalog.pg_description pgd ON (
+                    pgd.objoid = a.attrelid AND pgd.objsubid = a.attnum)
+            WHERE a.attrelid = (SELECT c.oid
+                                FROM pg_catalog.pg_class c
+                                    LEFT JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
+                                WHERE n.nspname = {schema!r}
+                                      AND c.relname = {table!r} AND c.relkind IN ('r', 'v', 'm', 'f'))
+                  AND a.attnum > 0 AND NOT a.attisdropped
+            ORDER BY a.attnum;
+            """
+            cursor.execute(col_comment_sql)
+            col_info = cursor.fetchall()
+
+        if t_comment:
+            comments = [f"COMMENT ON TABLE {self.quote_identifier(table)} IS {t_comment[0][0]!r};"]
+        else:
+            comments = []
+
+        cols = []
+        for col in col_info:
+            if col in field_filter:
+                continue
+            default = " DEFAULT {col[2]}" if col[2] else ""
+            isnull = " NOT NULL " if col[3] else ""
+            if "character varying" in col[1]:
+                ctype = col[1].replace("character varying", "varchar")
+                cols.append(f"{self.quote_identifier(col[0])} {ctype}{isnull}{default}")
+            else:
+                cols.append(f"{self.quote_identifier(col[0])} {col[1]}{isnull}{default}")
+            if col[4]:
+                comments.append(
+                    f"COMMENT ON COLUMN {self.quote_identifier(table)}.{self.quote_identifier(col[0])} IS {col[4]!r};"
+                )
+        if_exists_stmt = " IF NOT EXISTS " if if_exists else " "
+        cols_stmt = ", ".join(cols)
+        comments_stmt = " ".join(comments)
+        return f"CREATE TABLE{if_exists_stmt}{self.quote_identifier(table)} ({cols_stmt}); {comments_stmt}"
+
+    def is_postgres(self):
+        return True
+
+    @staticmethod
+    def to_canonical_type(type_code, size):
+        return _pg_type_to_canonical_type.get(type_code, types.STRING)
+
+    @staticmethod
+    def from_canonical_type(canonical_type, size):
+        return canonical_type_to_pg_type.get(canonical_type, "TEXT")
+
+    def load_csv(
+        self,
+        table,
+        filename,
+        schema="public",
+        columns=None,
+        delimiter=",",
+        quotechar='"',
+        lineterminator="\r\n",
+        escapechar=None,
+        skiprows=0,
+        using_insert=True,
+        **kwargs,
+    ):
+        # if using_insert:
+        #     method = self.load_csv_by_inserting
+        # else:
+        #     method = self._copy_csv
+        if not using_insert:
+            self.logger.warning("load file directly is not implemented yet, fallback to using bulk INSERT")
+
+        method = self.load_csv_by_inserting
+        schema, table = self._get_schema_table(table, schema)
+        table = self._format_table_name(table, schema)
+
+        return method(
+            table, filename, columns, delimiter, quotechar, lineterminator, escapechar, skiprows=skiprows, **kwargs
+        )
+
+    def _copy_csv(
+        self,
+        table,
+        filename,
+        columns=None,
+        delimiter=",",
+        quotechar='"',
+        lineterminator="\r\n",
+        escapechar=None,
+        skiprows=0,
+        **kwargs,
+    ):
+        conn = self.connect()
+        cursor = conn.cursor()
+        self.logger.info("copy file %s into %s", filename, table)
+        with open(filename, "r") as f:
+            if skiprows:
+                for _ in range(skiprows):
+                    f.readline()
+            # the copy_from method does support standard CSV
+            cursor.copy_from(f, table, sep=delimiter, columns=columns)
+        conn.commit()
+        conn.close()
+
+    def _get_schema_table(self, table, schema):
+        if "." in table:
+            schema, table = table.split(".")
+        if not schema:
+            schema = "public"
+        return schema, table
+
+    def _format_table_name(self, table, schema):
+        if schema and "." not in table:
+            table = self.quote_identifier(f"{schema}.{table}")
+        return table

recurvedata/pigeon/connector/power_bi.py

@@ -0,0 +1,179 @@
+import copy
+import logging
+import time
+from typing import Dict, List, Union
+
+import msal
+import pandas as pd
+import requests
+
+config = dict(
+    # Can be set to 'MasterUser' or 'ServicePrincipal'
+    AUTHENTICATION_MODE='ServicePrincipal',
+    POWER_BI_TENANT_ID='',
+    POWER_BI_CLIENT_ID='',
+    # Client Secret (App Secret) of the AAD app. Required only for ServicePrincipal authentication mode.
+    POWER_BI_CLIENT_SECRET='',
+    # Scope of AAD app. Use the below configuration to use all the permissions provided in the AAD(Azure Active Directory) app through Azure portal.
+    POWER_BI_SCOPE=['https://analysis.windows.net/powerbi/api/.default'],  # 公有云
+    POWER_BI_SCOPE_CN=['https://analysis.chinacloudapi.cn/powerbi/api/.default'],  # 中国区
+    # URL used for initiating authorization request
+    POWER_BI_AUTHORITY='https://login.microsoftonline.com/tenant_id',
+    POWER_BI_AUTHORITY_CN='https://login.chinacloudapi.cn/tenant_id',
+    POWER_BI_API_URL_PREFIX='https://api.powerbi.com/v1.0/myorg',
+    POWER_BI_API_URL_PREFIX_CN='https://api.powerbi.cn/v1.0/myorg'
+)
+
+
+class PBIRefreshFailedException(Exception):
+    pass
+
+
+class PBIRefreshTimeoutException(Exception):
+    pass
+
+
+class PowerBI:
+
+    def __init__(self, tenant_id: str, client_id: str, client_secret: str, **kwargs):
+        self.config = copy.deepcopy(config)
+        self.config["POWER_BI_TENANT_ID"] = tenant_id
+        self.config["POWER_BI_CLIENT_ID"] = client_id
+        self.config["POWER_BI_CLIENT_SECRET"] = client_secret
+        for k, v in kwargs.items():
+            if k in self.config:
+                self.config[k] = v
+        self.access_token = None
+        self.token_abort_time = None
+
+    def get_access_token(self):
+        if self.access_token is not None and time.time() < self.token_abort_time:
+            return self.access_token
+        try:
+            # Service Principal auth is the recommended by Microsoft to achieve App Owns Data Power BI embedding
+            authority = self.config['POWER_BI_AUTHORITY_CN'].replace('tenant_id', self.config['POWER_BI_TENANT_ID'])
+            client_app = msal.ConfidentialClientApplication(
+                client_id=self.config['POWER_BI_CLIENT_ID'],
+                client_credential=self.config['POWER_BI_CLIENT_SECRET'],
+                authority=authority
+            )
+            # Make a client call if Access token is not available in cache
+            response = client_app.acquire_token_for_client(scopes=self.config['POWER_BI_SCOPE_CN'])
+            self.access_token = response
+            self.token_abort_time = time.time() + (response["expires_in"] - 60)
+
+            return response
+
+        except Exception as ex:
+            raise Exception('Error retrieving Access token\n' + str(ex))
+
+    @property
+    def request_header(self):
+        """
+        Get Power BI API request header
+        """
+        access_token = self.get_access_token()
+        return {
+            'Content-Type': 'application/json',
+            'Authorization': f"{access_token['token_type']} {access_token['access_token']}"
+        }
+
+    def get_refresh_job_info(self, group_id: str, dataset_id: str, request_id: str = None, limit: int = 10) -> Union[Dict, List[Dict]]:
+        """
+        Parameters:
+            group_id: The workspace ID
+            dataset_id: The dataset ID
+            request_id: 如果指定 request_id，返回对应 request Dict 否则返回 List[Dict]
+            limit: numbers of recently requests (Descending), default 10
+        """
+        url = f"{self.config['POWER_BI_API_URL_PREFIX_CN']}/groups/{group_id}/datasets/{dataset_id}/refreshes/?$top={limit}"
+        r = requests.get(url, headers=self.request_header)
+        r.raise_for_status()
+        ret = r.json()["value"]
+        if request_id:
+            return next(filter(lambda x: x["requestId"] == request_id, ret), None)
+        return ret
+
+    def refresh_dataset_in_group(self, group_id: str, dataset_id: str, is_wait: bool = True, timeout: int = 300, check_interval: int = 20, limit: int = 10):
+        """
+        推荐使用 refresh_datasets()
+
+        https://learn.microsoft.com/en-us/rest/api/power-bi/datasets/refresh-dataset-in-group \n
+        Limitation:
+            For Shared capacities, a maximum of 8 requests per day, including refreshes executed by using scheduled refresh, can be initiated.
+
+        Parameters:
+            group_id: The workspace ID
+            dataset_id: The dataset ID
+            is_wait: Wait until refresh finish, default wait for refreshing
+            timeout: Default timeout is 5 minutes if is_wait is True
+            check_interval: Default 20 seconds
+            limit: numbers of recently requests (Descending)
+
+        Returns:
+            The refreshing job information
+        """
+
+        logging.info(f"Start refreshing dataset {dataset_id} in group {group_id}")
+        url = f"{self.config['POWER_BI_API_URL_PREFIX_CN']}/groups/{group_id}/datasets/{dataset_id}/refreshes"
+        r = requests.post(url, headers=self.request_header)
+        r.raise_for_status()
+        request_id = r.headers.get("RequestId")
+        # get refresh job information
+        job_info = self.get_refresh_job_info(group_id, dataset_id, request_id, limit)
+        if not job_info:  # if don't receive specific job, wait a second
+            time.sleep(5)
+            job_info = self.get_refresh_job_info(group_id, dataset_id, request_id, limit)
+        logging.info(f"Refresh detail: request_id -> {request_id}, job_info -> {job_info}")
+        if not is_wait:
+            if job_info["status"] == "Failed":
+                logging.info(f"Refresh failed: {dataset_id}")
+                raise PBIRefreshFailedException(job_info)
+            return job_info
+        abort_time = time.time() + timeout
+        while job_info["status"] != "Completed":
+            if job_info["status"] == "Failed":
+                logging.info(f"Refresh failed: {dataset_id}")
+                raise PBIRefreshFailedException(job_info)
+            time.sleep(check_interval)
+            if time.time() > abort_time:
+                logging.info(f"Refresh timeout: {dataset_id}")
+                raise PBIRefreshTimeoutException(job_info)
+            job_info = self.get_refresh_job_info(group_id, dataset_id, request_id, limit)
+            logging.info(f"Retry: {job_info}")
+        logging.info(f"Refresh completed: {dataset_id}")
+        return job_info
+
+    def refresh_datasets(self, refresh_list: pd.DataFrame, is_wait: bool = True, timeout: int = 300, check_interval: int = 20, limit: int = 10):
+        """
+        传入包含 group_id，dataset_id 列的 dataframe，刷新完毕返回刷新的情况
+        """
+        summary = {
+            "Completed": [], "Failed": [], "Timeout": [], "Error": []
+        }
+        if not {"group_id", "dataset_id"}.issubset(refresh_list.columns):
+            raise Exception(f"Contain wrong columns, input must include group_id and dataset_id, while target dataframe has {refresh_list.columns.to_list()} columns.")
+        for _, row in refresh_list.iterrows():
+            group_id, dataset_id = row["group_id"], row["dataset_id"]
+            try:
+                job_info = self.refresh_dataset_in_group(group_id, dataset_id, is_wait, timeout, check_interval, limit)
+                summary["Completed"].append({"group_id": group_id, "dataset_id": dataset_id, "job_info": job_info})
+            except PBIRefreshFailedException as e:
+                summary["Failed"].append({"group_id": group_id, "dataset_id": dataset_id, "job_info": e.args[0]})
+            except PBIRefreshTimeoutException as e:
+                summary["Timeout"].append({"group_id": group_id, "dataset_id": dataset_id, "job_info": e.args[0]})
+            except Exception as e:
+                summary["Error"].append({"group_id": group_id, "dataset_id": dataset_id, "reason": repr(e)})
+        return summary
+
+    def get_datasets_in_group(self, group_id: str) -> pd.DataFrame:
+        """
+        https://learn.microsoft.com/en-us/rest/api/power-bi/datasets/get-datasets-in-group \n
+        Returns a list of datasets from the specified workspace.
+        """
+        url = f"{self.config['POWER_BI_API_URL_PREFIX_CN']}/groups/{group_id}/datasets"
+        r = requests.get(url, headers=self.request_header)
+        r.raise_for_status()
+        df = pd.DataFrame(r.json()["value"]).rename(columns={"id": "dataset_id"})
+        df["group_id"] = group_id
+        return df

recurvedata/pigeon/connector/qcloud_cos.py

@@ -0,0 +1,79 @@
+import os
+
+import qcloud_cos
+
+from recurvedata.pigeon.connector._registry import register_connector_class
+from recurvedata.pigeon.utils.timing import ProgressCallback
+
+
+@register_connector_class("cos")
+class COSConnector(object):
+    def __init__(self, secret_id, secret_key, region, proxies=None, endpoint=None, **kwargs):
+        self.secret_id = secret_id
+        self.secret_key = secret_key
+        self.region = region
+
+        self.config = qcloud_cos.CosConfig(
+            Region=region, SecretId=secret_id, SecretKey=secret_key, Endpoint=endpoint, Proxies=proxies
+        )
+        self.cos = qcloud_cos.CosS3Client(self.config)
+
+    def has_bucket(self, bucket_name):
+        return self.cos.bucket_exists(bucket_name)
+
+    def create_bucket(self, bucket_name):
+        if not self.has_bucket(bucket_name):
+            self.cos.create_bucket(bucket_name)
+
+    def delete_bucket(self, bucket_name):
+        if self.has_bucket(bucket_name):
+            self.cos.delete_bucket(bucket_name)
+
+    def has_object(self, bucket_name, key):
+        return self.cos.object_exists(bucket_name, key)
+
+    def delete_object(self, bucket_name, key):
+        self.cos.delete_object(bucket_name, key)
+
+    def list_objects(self, bucket_name, prefix=""):
+        res = self.cos.list_objects(Bucket=bucket_name, Prefix=prefix)
+        return [x["Key"] for x in res.get("Contents", [])]
+
+    def delete_keys_by_prefix(self, bucket_name, prefix):
+        keys = self.list_objects(bucket_name, prefix)
+        for key in keys:
+            self.delete_object(bucket_name, key)
+
+    def upload(self, bucket_name, 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(bucket_name=bucket_name, key=key):
+                return key
+
+        self.cos.upload_file(
+            Bucket=bucket_name,
+            LocalFilePath=filename,
+            Key=key,
+            MAXThread=num_threads,
+            progress_callback=ProgressCallback(),
+        )
+        return key
+
+    def download(self, bucket_name, key, folder=None, filename=None, overwrite=True, num_threads=4, **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
+
+        self.cos.download_file(Bucket=bucket_name, Key=key, DestFilePath=filename, MAXThread=num_threads)
+        return filename