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

recurvedata/pigeon/connector/redshift.py

@@ -0,0 +1,123 @@
+import hashlib
+import os
+
+import cytoolz as toolz
+
+from recurvedata.pigeon.connector._registry import register_connector_class
+from recurvedata.pigeon.connector.awss3 import S3Connector
+from recurvedata.pigeon.connector.postgresql import PostgresConnector, canonical_type_to_pg_type
+from recurvedata.pigeon.utils import fs
+
+
+@register_connector_class("redshift")
+class RedshiftConnector(PostgresConnector):
+    _max_text = "VARCHAR(MAX)"
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        self.s3_bucket_name = self.kwargs.get("s3_options", {}).get("bucket")
+
+    def is_redshift(self):
+        return True
+
+    @toolz.memoize
+    def create_s3_connector(self):
+        s3_options = self.kwargs.get("s3_options")
+        if not s3_options:
+            return None
+        return S3Connector(**s3_options)
+
+    def load_csv(
+        self,
+        table,
+        filename,
+        schema="public",
+        columns=None,
+        delimiter=",",
+        quotechar='"',
+        lineterminator="\r\n",
+        escapechar=None,
+        skiprows=0,
+        using_insert=False,
+        **kwargs,
+    ):
+        table = self._format_table_name(table, schema)
+        s3 = self.create_s3_connector()
+        if using_insert or s3 is None:
+            self.load_csv_by_inserting(
+                table, filename, columns, delimiter, quotechar, lineterminator, escapechar, skiprows=skiprows, **kwargs
+            )
+        else:
+            self.load_csv_by_s3(table, filename, columns, skiprows, **kwargs)
+
+    def load_csv_by_s3(self, table, filename, columns=None, skiprows=0, **kwargs):
+        s3 = self.create_s3_connector()
+        bucket = self.generate_s3_bucket_name()
+        if filename.endswith(".gz"):
+            file_to_upload = filename
+        else:
+            self.logger.info("compressing %s", filename)
+            file_to_upload = fs.gzip_compress(filename, using_cmd=True)
+        s, t = self._get_schema_table(table, schema=None)
+        key_name = f"{self.database}/{s}/{t}/{os.path.basename(file_to_upload)}"
+        key_uri = self.format_s3_key_uri(bucket, key_name)
+        self.logger.info("upload %s to %s", file_to_upload, key_uri)
+        s3.upload(bucket, file_to_upload, key_name)
+
+        if columns:
+            field_names = "({})".format(", ".join([self.quote_identifier(x) for x in columns]))
+        else:
+            field_names = ""
+
+        # TODO: null
+        if skiprows:
+            ignore_header = f"IGNOREHEADER AS {int(skiprows)}"
+        else:
+            ignore_header = ""
+        stmt = f"""
+            COPY {table} {field_names} FROM '{key_uri}'
+            credentials 'aws_access_key_id={s3.aws_access_key_id};aws_secret_access_key={s3.aws_secret_access_key}'
+            region '{s3.region}'
+            CSV GZIP ACCEPTINVCHARS EMPTYASNULL {ignore_header}
+        """
+
+        try:
+            self.logger.info("running COPY command")
+            self.execute(stmt, autocommit=False, commit_on_close=True)
+            self.logger.info("COPY finished")
+        except Exception as e:
+            self.logger.exception("failed to copy data to Redshift")
+            raise e
+        finally:
+            if file_to_upload != filename:
+                self.logger.info("delete %s", file_to_upload)
+                fs.remove_files_safely(file_to_upload)
+
+            self.logger.info("delete S3 file: %s", key_uri)
+            try:
+                s3.delete_key(key_name, bucket)
+            except Exception as e:
+                self.logger.error(f"operation on s3 bucket fails: {e}")
+
+    @staticmethod
+    def from_canonical_type(canonical_type, size):
+        rv = canonical_type_to_pg_type.get(canonical_type, "VARCHAR(MAX)")
+        if rv == "TEXT":
+            rv = "VARCHAR(MAX)"
+        return rv
+
+    @staticmethod
+    def get_key_name(filename):
+        return os.path.basename(filename)
+
+    @staticmethod
+    def format_s3_key_uri(bucket, key_name):
+        return f"s3://{bucket}/{key_name}"
+
+    def generate_s3_bucket_name(self):
+        if self.s3_bucket_name:
+            return self.s3_bucket_name
+        cluster_name = self.host.split(".", 1)[0]
+        digest = hashlib.md5(self.host.encode()).hexdigest()
+        return f"pigeon-{cluster_name}-{digest[:15]}"

recurvedata/pigeon/connector/sftp.py

@@ -0,0 +1,73 @@
+import os
+import shutil
+
+import paramiko
+
+from recurvedata.pigeon.connector._registry import register_connector_class
+from recurvedata.pigeon.utils import LoggingMixin
+from recurvedata.pigeon.utils.timing import DisplayProgress
+
+
+@register_connector_class('sftp')
+class SFtpConnector(LoggingMixin):
+    def __init__(self, host, port, username, password, rsa_private_key_file: str = None):
+        self.host = host
+        self.port = port
+        self.username = username
+        self.password = password
+        client = paramiko.Transport((self.host, self.port))
+        if rsa_private_key_file and password:
+            private_key = paramiko.RSAKey.from_private_key_file(rsa_private_key_file)
+            client.start_client(event=None, timeout=15)
+            client.get_remote_server_key()
+            client.auth_publickey(self.username, private_key, event=None)
+            client.auth_password(self.username, self.password, event=None)
+        elif rsa_private_key_file:
+            private_key = paramiko.RSAKey.from_private_key_file(rsa_private_key_file)
+            client.connect(username=self.username, pkey=private_key)
+        else:
+            client.connect(username=self.username, password=self.password)
+        self.sftp = paramiko.SFTPClient.from_transport(client)
+
+    def close(self):
+        self.sftp.close()
+
+    def rename(self, from_name, to_name):
+        self.sftp.rename(from_name, to_name)
+
+    def makedir(self, path):
+        self.sftp.mkdir(path)
+
+    def rmdir(self, path):
+        self.sftp.rmdir(path)
+
+    def rm(self, name):
+        self.sftp.remove(name)
+
+    def pwd(self):
+        return self.sftp.getcwd()
+
+    def size(self, name):
+        return self.sftp.stat(name).st_size
+
+    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:
+            size = self.size(src_file)
+            self.sftp.get(src_file, dst_file, callback=DisplayProgress(size, stream=False))
+            self.logger.info(f'successfully downloaded {src_file} to {dst_file}')
+        except Exception as e:
+            os.unlink(dst_file)
+            if not exists:
+                shutil.rmtree(local_dir)
+
+            self.logger.exception(f'failed to download {src_file}, reason:{e}')
+            raise e
+
+    def upload_file(self, src_file, dst_file):
+        self.sftp.put(src_file, dst_file, callback=DisplayProgress(stream=False))
+        self.logger.info(f'successfully uploaded {src_file} to {dst_file}')

recurvedata/pigeon/connector/sqlite.py

@@ -0,0 +1,42 @@
+import sqlite3
+
+import pandas as pd
+
+
+class SQLiteMemoryDbConnector:
+    """
+    SQLite 内存数据库连接器，用于在内存中的pandas Dataframe执行SQL查询, 支持标准SQL语法
+    暂时只支持内存数据库, 不连接实体表和库
+    python进程关闭后，所有表和数据都会丢失
+    e.g.
+    sqlite_conn = SQLiteMemoryDbConnector()
+    df_1 = pd.DataFrame({'a': [1, 2, 3], 'b': [4, 5, 6]})
+    df_2 = pd.DataFrame({'a': [4, 5, 6], 'b': [7, 8, 9]})
+    sqlite_conn.create_temp_table(df_1, 'df_1')
+    sqlite_conn.create_temp_table(df_2, 'df_2')
+    result = sqlite_conn.get_pandas_df('SELECT *,row_number() over(partition by a order by b) as rn FROM df_1')
+    """
+
+    def __init__(self, max_memory_gb, **kwargs):
+        self.conn = sqlite3.connect(database=':memory:', **kwargs)  # 创建内存中的 SQLite 数据库
+        self.conn.execute(f"PRAGMA max_memory = {max_memory_gb * 1024 * 1024}")
+        self.cursor = self.conn.cursor()
+        self.loaded_tables = set()
+
+    def create_temp_table(self, df, table_name):
+        """ Write a table in the memory database. """
+        df.to_sql(table_name, self.conn, index=False, if_exists='replace')
+
+    def drop_temp_table(self, table_name):
+        """ Drop a table in the memory database. """
+        self.conn.execute(f"DROP TABLE IF EXISTS {table_name}")
+
+    def get_pandas_df(self, query: str) -> pd.DataFrame:
+        """
+        :param query: SQL 查询语句
+        :return: pandas DataFrame
+        """
+        return pd.read_sql_query(query, self.conn)
+
+    def close(self):
+        self.conn.close()

recurvedata/pigeon/connector/starrocks.py

@@ -0,0 +1,144 @@
+import json
+import subprocess
+
+from recurvedata.pigeon.connector._registry import register_connector_class
+from recurvedata.pigeon.connector.mysql import MySQLConnector
+from recurvedata.pigeon.schema import types
+
+_canonical_type_to_starrocks_type = {
+    types.BOOLEAN: "TINYINT",
+    types.INT8: "TINYINT",
+    types.INT16: "SMALLINT",
+    types.INT32: "INT",
+    types.INT64: "BIGINT",
+    types.FLOAT32: "FLOAT",
+    types.FLOAT64: "DOUBLE",
+    types.DATE: "DATE",
+    types.DATETIME: "DATETIME",
+    types.STRING: "STRING",
+    types.JSON: "STRING",
+}
+
+
+@register_connector_class(["starrocks"])
+class StarRocksConnector(MySQLConnector):
+    _default_port = 9030
+    _default_fe_http_port = 8030
+
+    def __init__(self, host, port=None, http_port=None, database=None, user=None, password=None, *args, **kwargs):
+        self.http_port = http_port or self._default_fe_http_port
+        self.user = user
+        self.password = password
+        super().__init__(host=host, port=port, database=database, user=user, password=password, *args, **kwargs)
+
+    @property
+    def load_strict_mode(self) -> bool:
+        if not hasattr(self, "_load_strict_mode"):
+            return False
+        return self._load_strict_mode
+
+    @load_strict_mode.setter
+    def load_strict_mode(self, mode: bool):
+        self._load_strict_mode = mode
+
+    def _load_csv_mysql(
+        self,
+        table,
+        filename,
+        columns=None,
+        delimiter=",",
+        quotechar='"',
+        lineterminator="\r\n",
+        escapechar=None,
+        skiprows=0,
+        **kwargs,
+    ):
+        """
+        stream load data from csv file into table
+        """
+
+        def _split_database_table(table_name: str):
+            tmp_lst = table_name.split(".")
+            if len(tmp_lst) == 1:
+                return self.database, table_name
+            return tmp_lst
+
+        db_name, table_name = _split_database_table(table)
+        shell_cmd = self._format_load_shell(filename, db_name, table_name)
+
+        # Set only authentication environment variables
+        _env = {}
+        if self.user is not None:
+            _env["STARROCKS_USER"] = self.user
+        if self.password is not None:
+            _env["STARROCKS_PASSWORD"] = self.password
+
+        output = subprocess.check_output(shell_cmd, env=_env, shell=True)
+        res_txt = output.decode()
+        if res_txt:
+            res = json.loads(res_txt)
+            self._log(res_txt)
+
+            if res["Status"] != "Success":
+                if "ErrorURL" not in res:
+                    err_output = res["Message"]
+                else:
+                    err_url = res["ErrorURL"]
+                    err_output = subprocess.check_output(["curl", err_url])
+                self._log(f"error: {err_output}")
+                raise Exception("load csv failed")
+
+    def _format_load_shell(self, filename: str, db_name: str, table_name: str) -> str:
+        def __format_column_mapping(db_name: str, table_name: str) -> str:
+            """
+            stream load 有个 bug，最后一列的右引号无法去除
+            在 column_mapping 里处理
+            """
+            columns = self.get_columns(table_name, db_name)
+            # Extract column names from the column metadata
+            column_names = [col["name"] for col in columns]
+            # Escape backticks to prevent shell interpretation
+            cols = [f"\\`{col}\\`" for col in column_names]
+            cols_txt = ",".join(cols)
+            cols2 = [
+                f"\\`{col}\\`=trim(\\`{col}\\`,'\\\"')" if col == column_names[-1] else f"\\`{col}\\`=\\`{col}\\`"
+                for col in column_names
+            ]
+            cols_txt2 = ", ".join(cols2)
+            return f"columns: {cols_txt}, {cols_txt2}"
+
+        def __format_stream_load_url(db_name: str, table_name: str) -> str:
+            return f"http://{self.host}:{self.http_port}/api/{db_name}/{table_name}/_stream_load"
+
+        db_name, table_name = db_name.strip("`"), table_name.strip("`")
+        url = __format_stream_load_url(db_name, table_name)
+        strict_mode = "true" if self.load_strict_mode else "false"
+        column_mapping = __format_column_mapping(db_name, table_name)
+
+        # Handle authentication based on whether password is provided
+        if self.password is not None:
+            auth_part = "-u $STARROCKS_USER:$STARROCKS_PASSWORD"
+        else:
+            auth_part = "-u $STARROCKS_USER:"
+
+        # Construct the full curl command with properly escaped quotes
+        return (
+            f"curl --location-trusted {auth_part} "
+            f'-H "Expect:100-continue" '
+            f'-H "column_separator:," '
+            f'-H "enclose:\\"" '
+            f'-H "trim_double_quotes:true" '
+            f'-H "strict_mode:{strict_mode}" '
+            f'-H "escape:\'" '
+            f'-H "{column_mapping}" '
+            f"-T {filename} -XPUT "
+            f"{url}"
+        ).strip()
+
+    @staticmethod
+    def from_canonical_type(canonical_type, size):
+        if canonical_type == types.STRING:
+            starrocks_type = "STRING"
+        else:
+            starrocks_type = _canonical_type_to_starrocks_type.get(canonical_type, "STRING")
+        return starrocks_type

recurvedata/pigeon/connector/tableau.py

@@ -0,0 +1,162 @@
+"""
+official API doc: https://tableau.github.io/server-client-python/docs/api-ref#views
+tableau-api-lib：https://github.com/divinorum-webb/tableau-api-lib
+"""
+import logging
+import time
+from typing import Optional
+
+import pandas as pd
+import tableauserverclient as TSC
+from tableau_api_lib import TableauServerConnection
+from tableau_api_lib.utils.querying import (
+    get_datasource_connections_dataframe,
+    get_datasources_dataframe,
+    get_embedded_datasources_dataframe,
+    get_projects_dataframe,
+    get_sites_dataframe,
+    get_views_dataframe,
+    get_workbooks_dataframe,
+)
+from tableauserverclient import Server
+
+
+class TableauConnector:
+    def __init__(self, user: str = None, password: str = None, server_url: str = None, site: str = None):
+        self.user = user
+        self.password = password
+        self.server_url = server_url
+        self.site = site
+        self.server: Optional[Server] = None
+        self.library: Optional[TableauServerConnection] = None
+        self._sign_in()
+
+    def _sign_in(self):
+        tableau_auth = TSC.TableauAuth(self.user, self.password, self.site)
+        self.server = TSC.Server(
+            self.server_url, use_server_version=True
+        )  # https://community.tableau.com/s/question/0D54T00000ti0eOSAQ/api-version-upgrade
+        self.server.auth.sign_in(tableau_auth)
+
+        self.library = TableauServerConnection(
+            {
+                "my_env": {
+                    "api_version": self.server.version,
+                    "server": self.server_url,
+                    "username": self.user,
+                    "password": self.password,
+                    "site_name": self.site,
+                    "site_url": self.site,
+                }
+            },
+            env="my_env",
+        )
+        self.library.sign_in()
+
+    @staticmethod
+    def check_columns(df: pd.DataFrame, cols: list) -> pd.DataFrame:
+        if not cols:
+            return df
+        if not set(cols).issubset(df.columns):
+            raise Exception(
+                f"Contain wrong columns, target dataframe has {df.columns.to_list()} columns, while input has {cols}"
+            )
+        return df[cols]
+
+    def get_sites_df(self, cols: list = None):
+        df = get_sites_dataframe(self.library).rename(columns={"id": "site_id", "name": "site_name"})
+        return self.check_columns(df, cols)
+
+    def get_projects_df(self, cols: list = None):
+        df = get_projects_dataframe(self.library).rename(columns={"id": "project_id", "name": "project_name"})
+        return self.check_columns(df, cols)
+
+    def get_workbooks_df(self, cols: list = None):
+        df = get_workbooks_dataframe(self.library).rename(columns={"id": "workbook_id", "name": "workbook_name"})
+        df["project_id"], df["project_name"] = zip(*df["project"].apply(lambda x: (x["id"], x["name"])))
+        # df.drop(columns=["project"], inplace=True)
+        return self.check_columns(df, cols)
+
+    def get_views_df(self, cols: list = None):
+        df = get_views_dataframe(self.library).rename(columns={"id": "view_id", "name": "view_name"})
+        df["workbook_id"], df["workbook_name"] = zip(*df["workbook"].apply(lambda x: (x["id"], x["name"])))
+        project_df = self.get_projects_df(cols=["project_id", "project_name"])
+        df["project_id"] = df["project"].apply(lambda x: x["id"])
+        df = df.merge(project_df, on="project_id", how="left")
+        return self.check_columns(df, cols)
+
+    def get_datasources_df(self, cols: list = None):
+        df = get_datasources_dataframe(self.library).rename(columns={"id": "datasource_id", "name": "datasource_name"})
+        df["project_id"], df["project_name"] = zip(*df["project"].apply(lambda x: (x["id"], x["name"])))
+        return self.check_columns(df, cols)
+
+    def get_embedded_datasources_df(self, workbook_df: pd.DataFrame, cols: list = None):
+        """
+        先筛选指定的 workbook，不然很慢
+        """
+        df = get_embedded_datasources_dataframe(
+            self.library, workbook_df, id_col="workbook_id", name_col="workbook_name"
+        )
+        return self.check_columns(df, cols)
+
+    def get_datasource_connections_df(self, datasources_df: pd.DataFrame = None, cols: list = None):
+        if datasources_df is None:
+            datasources_df = self.get_datasources_df(cols=["datasource_id", "datasource_name"])
+        else:
+            datasources_df = self.check_columns(datasources_df, cols=["datasource_id", "datasource_name"])
+        connections = []
+        for index, row in datasources_df.iterrows():
+            ds_conn = get_datasource_connections_dataframe(self.library, row["datasource_id"])
+            ds_conn["datasource_id"] = row["datasource_id"]
+            connections.append(ds_conn)
+        connections_df = pd.concat(connections, ignore_index=True)
+        connections_df = connections_df.merge(datasources_df, on="datasource_id", how="left")
+        return self.check_columns(connections_df, cols)
+
+    def get_job_status(self, job_id: str):
+        return self.server.jobs.get_by_id(job_id)
+
+    def wait_to_finish(self, job_id, timeout, retry_interval):
+        abort_time = time.time() + timeout
+        job_info = self.get_job_status(job_id)
+        while job_info.completed_at is None:
+            logging.info(
+                f"finish_code: {job_info.finish_code}, progress: {job_info.progress} %. Sleep for {retry_interval} s."
+            )
+            time.sleep(retry_interval)
+            if time.time() > abort_time:
+                logging.warning(f"Timeout {timeout} s. Job_info: {job_info}")
+                break
+            job_info = self.get_job_status(job_id)
+        if job_info.finish_code != 0:
+            logging.warning(f"Job {job_id} is not success")
+        return job_info
+
+    def refresh_workbook(self, workbook_id: str, timeout=600, retry_interval=5):
+        logging.info(f"Start refreshing workbook: {workbook_id}")
+        res = self.server.workbooks.refresh(workbook_id)
+        job_info = self.wait_to_finish(res.id, timeout, retry_interval)
+        logging.info(f"Finish refreshing: {job_info}")
+
+    def refresh_datasource(self, datasource_id: str, timeout=600, retry_interval=5):
+        logging.info(f"Start refreshing datasource: {datasource_id}")
+        res = self.server.datasources.refresh(datasource_id)
+        job_info = self.wait_to_finish(res.id, timeout, retry_interval)
+        logging.info(f"Finish refreshing: {job_info}")
+
+    def screenshot(self, workbook_id: str, view_id: str, save_path: str, maxage: int = 1):
+        """
+        截图可能有延迟
+        """
+        logging.info(f"Start taking screenshot with workbook_id {workbook_id}, view_id {view_id}")
+        workbook = self.server.workbooks.get_by_id(workbook_id)
+        self.server.workbooks.populate_views(workbook)
+
+        view = self.server.views.get_by_id(view_id)
+        image_req_option = TSC.ImageRequestOptions(
+            imageresolution=TSC.ImageRequestOptions.Resolution.High, maxage=maxage
+        )
+        self.server.views.populate_image(view, image_req_option)
+        with open(save_path, "wb") as f:
+            f.write(view.image)
+        logging.info(f"Finish saving screenshot to {save_path}")

recurvedata/pigeon/const.py

@@ -0,0 +1,21 @@
+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"
+
+LOAD_RENAME_OVERWRITE = "RENAME_OVERWRITE"
+LOAD_OVERWRITE = "OVERWRITE"
+LOAD_MERGE = "MERGE"
+LOAD_APPEND = "APPEND"
+
+HIVE_FILE_FORMATS = {
+    "text": "TEXTFILE",
+    "sequence": "SEQUENCEFILE",
+    "parquet": "PARQUET",  # http://parquet.apache.org/documentation/latest/
+    "orc": "ORC",  # optimized row columnar file
+    "rc": "RCFILE",  # record columnar file
+    "avro": "AVRO",  # Apache Avro™ (http://avro.apache.org/docs/current/)
+}
+
+CLICKHOUSE_MAX_ROW_BUFFER = 10000