rdxz2-utill 1.0.0__py3-none-any.whl

utill/my_pg.py

@@ -0,0 +1,159 @@
+import csv
+import json
+import os
+import psycopg
+import psycopg.rows
+
+from loguru import logger
+from textwrap import dedent
+
+from .my_env import PG_FILENAME
+from .my_string import generate_random_string
+from .my_tunnel import establish_tunnel
+
+
+class PG:
+    def __init__(
+        self,
+        connection=None,
+        config_source: str | dict = PG_FILENAME,
+        autocommit: bool = True,
+    ) -> None:
+        if type(config_source) == str:
+            if connection is None:
+                raise ValueError('Connection name must be provided when using file source!')
+            conf = json.loads(open(os.path.expanduser(config_source)).read())[connection]
+        elif type(config_source) == dict:
+            conf = config_source
+        (_, host, port) = establish_tunnel(conf)
+        self.db_host = host
+        self.db_port = port
+        self.db_username = conf['username']
+        self.db_password = conf['password']
+        self.db_name = conf['db']
+        self.conf = conf
+
+        self.conn = None
+        self.cursor = None
+        self.establish_connection(autocommit)
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, exc_tb):
+        self.close()
+
+    def establish_connection(self, autocommit: bool, row_factory: psycopg.rows = psycopg.rows.dict_row):
+        self.conn = psycopg.connect(f'postgresql://{self.db_username}:{self.db_password}@{self.db_host}:{self.db_port}/{self.db_name}', autocommit=autocommit)
+        self.cursor = self.conn.cursor(row_factory=row_factory)
+        logger.debug(f'PG client open: {self.db_username}@{self.db_host}:{self.db_port}/{self.db_name}, autocommit={self.conn.autocommit}')
+
+    def change_autocommit(self, autocommit: bool):
+        if autocommit == self.conn.autocommit:
+            return
+
+        self.conn.autocommit = autocommit
+
+    def execute_query(self, query: str, *params):
+        # Make sure connection alive
+        if self.conn.closed:
+            self.establish_connection(self.conn.autocommit)
+
+        query = query.strip()
+        logger.debug(f'🔎 Query:\n{query}')
+
+        return self.cursor.execute(query, params)
+
+    def download_csv(self, query: str, file_path: str) -> None:
+        query = dedent(
+            f'''
+            COPY ({query})
+            TO STDOUT
+            WITH DELIMITER ','
+            CSV HEADER;
+            '''
+        )
+        logger.debug(f'🔎 Query:\n{query}')
+        with open(os.path.expanduser(file_path), 'wb') as f:
+            with self.cursor.copy(query) as copy:
+                for data in copy:
+                    f.write(data)
+
+    def pg_to_pg(self, pg: "PG", source_table: str, target_table: str, cols: list[str] = None) -> None:
+        tmp_filename = generate_random_string() + '.csv'
+        cols_str = ','.join([f'"{x}"' for x in cols]) if (cols is not None and cols != []) else '*'
+        try:
+            self.download_csv(f'SELECT {cols_str} FROM {source_table}', tmp_filename)
+            pg.upload_csv(tmp_filename, target_table)
+        except:
+            raise
+        finally:
+            os.remove(tmp_filename) if os.path.exists(tmp_filename) else None
+
+    def check_table_existence(self, table_name: str) -> bool:
+        if not self.execute_query('''SELECT count(1) AS "cnt" FROM "information_schema"."tables" WHERE "table_schema" || '.' || "table_name" = '%s';''', table_name).fetchone()['cnt']:
+            raise Exception(f'Target table \'{table_name}\' not created, please create it first!')
+
+    def upload_tuples(self, cols: list[str], tuples: list[tuple], table_name: str) -> None:
+        self.check_table_existence(table_name)
+
+        cols_str = ','.join([f'"{x}"' for x in cols])
+        query = f'''COPY {table_name}({cols_str}) FROM STDIN'''
+        logger.debug(f'🔎 Query:\n{query}')
+        with self.cursor.copy(query) as copy:
+            for row in tuples:
+                copy.write_row(row)
+
+    def upload_list_of_dict(self, data: list[dict], table_name: str) -> None:
+        self.check_table_existence(table_name)
+
+        if len(data) == 0:
+            raise ValueError('No data to upload!')
+
+        cols = data[0].keys()
+        cols_str = ','.join([f'"{x}"' for x in cols])
+        query = f'''COPY {table_name}({cols_str}) FROM STDIN'''
+        logger.debug(f'🔎 Query:\n{query}')
+        with self.cursor.copy(query) as copy:
+            for row in data:
+                copy.write_row(tuple(row[col] for col in cols))
+
+    def upload_csv(self, file_path: str, table_name: str) -> None:
+        self.check_table_existence(table_name)
+
+        cols_str = ','.join([f'"{x}"' for x in next(csv.reader(open(file_path, 'r')))])
+        query = dedent(
+            f'''
+            COPY {table_name}({cols_str})
+            FROM STDIN
+            DELIMITER ','
+            CSV HEADER;
+            '''
+        )
+        logger.debug(f'🔎 Query:\n{query}')
+        with open(os.path.expanduser(file_path), 'r') as f:
+            with self.cursor.copy(query) as copy:
+                while data := f.read(1024):
+                    copy.write(data)
+
+    def create_index(self, table_name: str, index: str | list[str], unique: bool = False) -> None:
+        try:
+            index = index if type(index) == list else [index]
+            indexes = ','.join([f'"{x}"' for x in index])
+            self.execute_query(f'CREATE {"UNIQUE " if unique else ""}INDEX ON "{table_name}" ({indexes});', return_df=False)
+        except Exception as e:
+            self.rollback()
+            raise e
+
+    def rollback(self):
+        self.conn.rollback()
+        logger.debug('🚫 Transaction rollback')
+
+    def commit(self):
+        self.conn.commit()
+        logger.debug('✅ Transaction commit')
+
+    def close(self):
+        self.cursor.close()
+        self.conn.close()
+        logger.debug('PG client close')

utill/my_queue.py

@@ -0,0 +1,66 @@
+import queue
+import concurrent.futures
+
+from loguru import logger
+
+
+class ThreadingQ:
+    def __init__(self) -> None:
+        self.q = queue.Queue()
+
+        self.producer_func = None
+        self.producer_args = None
+        self.consumer_func = None
+
+    def add_producer(self, func, *args):
+        self.producer_func = func
+        self.producer_args = args or []
+        return self
+
+    def add_consumer(self, func):
+        self.consumer_func = func
+        # The consume args is based on producer output
+        return self
+
+    def execute(self):
+        if not all([self.producer_func is not None, self.producer_args is not None, self.consumer_func is not None]):
+            raise Exception('Producer and Consumer functions must be defined!')
+
+        def producer():
+            results = []
+
+            for item in self.producer_func(*self.producer_args):
+                self.q.put(item)
+                results.append(item)
+                logger.debug(f'🌾 Produced {item}')
+
+            self.q.put(None)
+            return results
+
+        def consumer():
+            results = []
+
+            while True:
+                item = self.q.get()
+                if item is None:
+                    break
+
+                result = self.consumer_func(*item)
+                results.append(result)
+
+                self.q.task_done()
+                logger.debug(f'🔥 Consumed {item}')
+
+            return results
+
+        with concurrent.futures.ThreadPoolExecutor() as executor:
+            # Schedule the producer and consumer
+            self.future_producer = executor.submit(producer)
+            self.future_consumer = executor.submit(consumer)
+
+            producer_result = self.future_producer.result()
+            logger.debug('✅ Producer done')
+            consumer_result = self.future_consumer.result()
+            logger.debug('✅ Consumer done')
+
+        return producer_result, consumer_result

utill/my_string.py

@@ -0,0 +1,20 @@
+import hashlib
+import random
+import re
+import string
+
+
+def generate_random_string(length: int = 4, alphanum: bool = False): return ''.join(random.choice(string.ascii_letters + string.digits + (r'!@#$%^&*()-=_+[]{};\':",./<>?' if not alphanum else '')) for _ in range(length))
+
+
+def replace_nonnumeric(string: str, replace: str) -> str: return re.sub('[^0-9a-zA-Z]+', replace, string)
+
+
+def mask(string: str, mask_length_min: int = 5, mask_length_max: int = 50, display_length: int = 5) -> str:
+    if not string:
+        mask_length = mask_length_min
+    else:
+        hash_value = int(hashlib.sha256(string.encode()).hexdigest(), 16)
+        mask_length = mask_length_min + (hash_value % (mask_length_max - mask_length_min + 1))
+
+    return ('*' * mask_length) + (string[(-display_length if len(string) > display_length else -1):])

utill/my_style.py

@@ -0,0 +1,39 @@
+class Styles:
+    NONE = '\033[0m'
+    ITALIC = '\033[3m'
+    BOLD = '\033[1m'
+    UNDERLINE = '\033[4m'
+
+
+class Colors:
+    HEADER = '\033[95m'
+    OKBLUE = '\033[94m'
+    OKCYAN = '\033[96m'
+    OKGREEN = '\033[92m'
+    WARNING = '\033[93m'
+    RED = '\033[91m'
+    BOLD = '\033[1m'
+    UNDERLINE = '\033[4m'
+
+
+def make_style(styles_or_colors: list[Styles | Colors] | Styles | Colors, string: str) -> str:
+    if type(styles_or_colors) == list:
+        return ''.join(styles_or_colors) + string + Styles.NONE
+    else:
+        return styles_or_colors + string + Styles.NONE
+
+
+def bold(string: str) -> str:
+    return make_style(Styles.BOLD, string)
+
+
+def italic(string: str) -> str:
+    return make_style(Styles.ITALIC, string)
+
+
+def underline(string: str) -> str:
+    return make_style(Styles.UNDERLINE, string)
+
+
+def color(string: str, color: Colors) -> str:
+    return make_style(color, string)

utill/my_tunnel.py

@@ -0,0 +1,42 @@
+import socket
+
+from loguru import logger
+from sshtunnel import SSHTunnelForwarder
+
+LOCALHOST = '127.0.0.1'
+
+
+def _get_random_port() -> int:
+    s = socket.socket()
+    s.bind((LOCALHOST, 0))
+    return s.getsockname()[1]
+
+
+def start_tunnel(host: str, port: int, user: str, key: str, target_host: str, target_port: int, local_port: int = None) -> int:
+    local_port = local_port or _get_random_port()
+
+    tunnel = SSHTunnelForwarder(
+        (host, port),
+        ssh_username=user,
+        ssh_private_key=key,
+        remote_bind_address=(target_host, target_port),
+        local_bind_address=(LOCALHOST, local_port),
+    )
+
+    tunnel.start()
+
+    return (tunnel, LOCALHOST, local_port)
+
+
+def establish_tunnel(conf: dict, local_port: int = None) -> tuple:
+    using_tunnel = bool(conf.get('tunnel_host'))
+    local_host = LOCALHOST if using_tunnel else conf['host']
+
+    z = start_tunnel(conf['tunnel_host'], conf['tunnel_port'], conf['tunnel_username'], conf['tunnel_key'], conf['host'], conf['port'], local_port=local_port)\
+        if using_tunnel\
+        else (None, local_host, conf['port'])
+
+    if using_tunnel:
+        logger.debug(f'🛣️  Tunnel established: {conf["host"]}:{conf["port"]} --> {conf["tunnel_username"]}@{conf["tunnel_host"]} --> {z[1]}:{z[2]}')
+
+    return z

utill/my_xlsx.py

@@ -0,0 +1,21 @@
+import duckdb
+import multiprocessing
+
+from loguru import logger
+
+
+def xlsx_to_csv(filename: str, sheet: str):
+    con = duckdb.connect()
+    return con.execute('install spatial;')\
+        .execute('load spatial;')\
+        .execute(f'select * from st_read(\'{filename}\', layer=\'{sheet}\');')\
+        .fetchall()
+
+
+def csv_to_xlsx(filename: str, output_file_path: str):
+    logger.info(f'Converting csv \'{filename}\' into xlsx \'{output_file_path}\' ...')
+    con = duckdb.connect()
+    con.execute('install spatial;')\
+        .execute('load spatial;')\
+        .execute(f'set threads to {multiprocessing.cpu_count()};')\
+        .execute(f'copy \'{filename}\' to \'{output_file_path}\' with(format gdal, driver \'xlsx\')')