qoptimizer-utils 0.1.0__tar.gz

qoptimizer_utils-0.1.0/PKG-INFO

@@ -0,0 +1,21 @@
+Metadata-Version: 2.4
+Name: qoptimizer-utils
+Version: 0.1.0
+Summary: Shared utilities for Query Optimizer (shared/ + db_utils/)
+License: Proprietary
+Project-URL: Repository, https://github.com/dhkim77000/qoptimize_utils
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Requires-Python: >=3.10
+Requires-Dist: pydantic>=2.10.0
+Requires-Dist: pydantic-settings>=2.7.0
+Requires-Dist: supabase>=2.11.0
+Requires-Dist: cryptography>=44.0.0
+Requires-Dist: structlog>=24.4.0
+Requires-Dist: snowflake-connector-python>=3.12.0
+Requires-Dist: clickhouse-driver>=0.2.9
+Requires-Dist: google-cloud-bigquery>=3.27.0
+Requires-Dist: asyncpg>=0.30.0
+Requires-Dist: neo4j>=5.27.0

qoptimizer_utils-0.1.0/db_utils/__init__.py

@@ -0,0 +1 @@
+"""db_utils package."""

qoptimizer_utils-0.1.0/db_utils/clickhouse_access.py

@@ -0,0 +1,178 @@
+"""ClickHouseCloud connection code."""
+
+import os
+import sys
+
+from clickhouse_driver import Client
+
+_here = os.path.dirname(os.path.abspath(__file__))
+if _here not in sys.path:
+    sys.path.insert(0, _here)
+from db_utils.config import get_env_var
+from db_utils.logger import _init_logger
+
+logger = _init_logger(__file__)
+
+DEFAULT_SERVICE = "rw-standard"
+SERVICES = {
+    DEFAULT_SERVICE: "grkyl47mbo.us-west-2.aws.clickhouse.cloud",
+    "rw-burst": "el5bpab3my.us-west-2.aws.clickhouse.cloud",
+}
+
+
+def get_ch_config(service: str = DEFAULT_SERVICE, **kwargs: dict) -> dict:
+    """Get ClickHouse configuration.
+
+    Parameters
+    ----------
+    service : str, optional
+        The ClickHouse service to connect to, by default DEFAULT_SERVICE.
+    **kwargs : dict
+        Additional keyword arguments to override default configuration.
+
+    Returns
+    -------
+    dict
+        ClickHouse configuration
+    """
+    if service not in SERVICES:
+        msg = f"Invalid service name: {service}. Available services: {list(SERVICES.keys())}"
+        raise ValueError(msg)
+
+    logger.info(f"Connecting to ClickHouse service: {service}")
+
+    config = {
+        "host": SERVICES.get(service, SERVICES[DEFAULT_SERVICE]),
+        "username": get_env_var("ch_user", secret=True),
+        "password": get_env_var("ch_password", secret=True),
+    }
+    config.update(kwargs)
+    return config
+
+
+def clickhouse_connect(**kwargs: dict) -> Client:
+    """Get a connection to ClickHouseCloud.
+
+    Returns
+    -------
+    clickhouse_driver.Client
+        ClickHouse connection client
+    """
+    config = get_ch_config(**kwargs)
+
+    return Client(
+        host=config["host"],
+        user=config["username"],
+        password=config["password"],
+        send_receive_timeout=3300,
+        secure=True,
+    )
+
+
+def run_ch_sql(
+    query: str,
+    *,
+    with_column_types: bool = False,
+    return_dataframe: bool = False,
+    return_dict: bool = False,
+    **kwargs: dict,
+) -> list | tuple:
+    """Run a SQL query on ClickHouseCloud.
+
+    Parameters
+    ----------
+    query : str
+        The SQL query to be executed on ClickHouse.
+    with_column_types : bool, optional
+        If True, include column types in the results, by default False.
+    return_dataframe : bool, optional
+        If True, return results as a pandas DataFrame and column names, by default False.
+    return_dict : bool, optional
+        If True, return results as a list of dictionaries, by default False.
+    **kwargs
+        Additional keyword arguments to pass to the ClickHouse connection.
+
+    Returns
+    -------
+    list or tuple
+        list of tuples with query results
+
+        If return_dataframe is True:
+            tuple containing (DataFrame, list of column names)
+
+    Examples
+    --------
+        >>> run_ch_sql('query', return_dataframe=True)
+        '[ {name: 'John', age: 25}, {name: 'Jane', age: 22} ]'
+    """
+    con = clickhouse_connect(**kwargs)
+
+    if return_dataframe:
+        result = _run_ch_sql_as_dataframe(con, query)
+    elif return_dict:
+        result = _run_ch_sql_as_dict(con, query)
+    else:
+        result = con.execute(query, with_column_types=with_column_types)
+
+    con.disconnect_connection()
+    return result
+
+
+def _run_ch_sql_as_dataframe(con: Client, query: str) -> tuple:
+    """Run a SQL query on ClickHouseCloud and return results as a DataFrame.
+
+    Parameters
+    ----------
+    con : clickhouse_driver.Client
+        The ClickHouse connection client.
+    query : str
+        The SQL query to be executed on ClickHouse.
+
+    Returns
+    -------
+    DataFrame, list
+        A tuple containing a pandas DataFrame and a list of column names.
+    """
+    from pandas import DataFrame  # inline import since Pandas can be heavy
+
+    result, columns = con.execute(query, with_column_types=True)
+
+    column_names = [c[0] for c in columns]
+    dict_result = DataFrame(result, columns=column_names).to_dict(orient="records")
+
+    return dict_result, column_names
+
+
+def _run_ch_sql_as_dict(con: Client, query: str) -> list[dict]:
+    """Run a SQL query on ClickHouseCloud and return results as a list of dictionaries.
+
+    Parameters
+    ----------
+    con : clickhouse_driver.Client
+        The ClickHouse connection client.
+    query : str
+        The SQL query to be executed on ClickHouse.
+
+    Returns
+    -------
+    list
+        A list of dictionaries representing the query results.
+    """
+    from clickhouse_driver.dbapi.extras import DictCursor
+
+    cur = DictCursor(con, connection=None)
+    cur.execute(query)
+    return cur.fetchall()
+
+
+if __name__ == "__main__":
+    con = clickhouse_connect()
+    result = con.execute("select current_date();")
+    print(result)
+    con.disconnect_connection()
+
+    with clickhouse_connect() as ch_con:
+        result = ch_con.execute("select current_date();")
+        print(result)
+
+    print(run_ch_sql("select current_date();", return_results=True, service="rw-burst"))

qoptimizer_utils-0.1.0/db_utils/config.py

@@ -0,0 +1,97 @@
+"""Configuration."""
+
+import json
+import os
+from typing import Optional
+from .logger import _init_logger
+
+logger = _init_logger(__file__)
+
+# Vendor-related env keys: always from AWS Secrets Manager (SecretId = user_id or AWS_VENDOR_SECRET_ID).
+
+
+def _is_vendor_key(env_var: str) -> bool:
+    if env_var.startswith("PSQL_") or env_var.startswith("SF_") or env_var.startswith("BQ_"):
+        return True
+    if env_var in ("CH_HOST", "CH_USERNAME", "CH_PASSWORD", "CH_USER"):
+        return True
+    if env_var.startswith("CH_") and not env_var.startswith("CH_ANALYTICS_"):
+        return True
+    return False
+
+
+def _get_secret_from_aws(secret_id: str) -> dict:
+    """Fetch secret JSON from AWS Secrets Manager."""
+    import boto3
+    from botocore.exceptions import ClientError
+    region = os.environ.get("AWS_REGION", "us-east-1")
+    client = boto3.client("secretsmanager", region_name=region)
+    try:
+        resp = client.get_secret_value(SecretId=secret_id)
+        raw = resp.get("SecretString") or "{}"
+        return json.loads(raw)
+    except ClientError as e:
+        logger.debug("AWS SM get_secret_value %s: %s", secret_id, e)
+        return {}
+    except json.JSONDecodeError:
+        return {}
+
+
+def get_chartmetric_data_script_path() -> str:
+    """Get local directory path for chartmetric_data_script.
+
+    Returns
+    -------
+    str
+        local directory path
+    """
+    return os.environ.get("CHARTMETRIC_DATA_SCRIPT_PATH", "/home/ec2-user/chartmetric_data_script")
+
+
+def get_env_var(name: str, secret: Optional[bool] = None, user_id: Optional[str] = None) -> Optional[str]:
+    """Get value from AWS Secrets Manager (vendor keys) or environment.
+
+    Vendor-related keys (PSQL_*, SF_*, CH_* except CH_ANALYTICS_*, BQ_*):
+      - Fetched from AWS SM. SecretId = user_id or env AWS_VENDOR_SECRET_ID.
+    Other keys:
+      - From os.environ only.
+
+    Parameters
+    ----------
+    name : str
+        variable name (e.g. psql_user, ch_password)
+    secret : bool, optional
+        unused; kept for API compatibility
+    user_id : str, optional
+        user UUID for per-user secret; if None, uses AWS_VENDOR_SECRET_ID (system secret)
+
+    Returns
+    -------
+    str or None
+        value if found, else None
+    """
+    assert name, "name cannot be empty."
+    env_var = name.upper()
+
+    if _is_vendor_key(env_var):
+        secret_id = user_id or os.environ.get("AWS_VENDOR_SECRET_ID")
+        if secret_id:
+            secrets = _get_secret_from_aws(secret_id)
+            if secrets:
+                val = secrets.get(env_var) or (secrets.get("CH_USERNAME") if env_var == "CH_USER" else None)
+                if val is not None and str(val).strip():
+                    logger.debug("Using %s from AWS SM (secret_id=%s)", env_var, secret_id)
+                    return str(val).strip()
+
+    if value := os.environ.get(env_var):
+        logger.info("Found %s in environment variables.", env_var)
+        return value
+
+    logger.error("Could not find %s in AWS SM or environment.", name)
+    return None
+
+
+if __name__ == "__main__":
+    # test with python3 -m 'src.db_utils.config'
+    assert get_env_var("psql_writer_host") == "prod2.cluster-cni52ceaa2ty.us-west-2.rds.amazonaws.com"
+    assert get_env_var("chartmetric_api_username") == "data-eng@chartmetric.com"

qoptimizer_utils-0.1.0/db_utils/db_access.py

@@ -0,0 +1,9 @@
+"""DB access."""
+
+from .postgres.connect import db_connect as db_reader_connect
+from .postgres.connect import get_connection as get_reader_db_connection
+from .postgres.connect import get_cursor as get_reader_db_cursor
+from .postgres.connect import run_pg_sql
+from .postgres.helpers import  insert_data
+from .postgres.helpers import disconnect as db_disconnect
+from .postgres.writer import db_connect, get_db_connection, get_db_cursor

qoptimizer_utils-0.1.0/db_utils/logger.py

@@ -0,0 +1,89 @@
+"""Implements functions for logging."""
+
+import datetime
+import logging
+
+
+class Logger:
+    """Logger with info, skip, select, insert, update, delete, error, exception, stack_trace."""
+
+    def __init__(self, name: str) -> None:
+        self._log = logging.getLogger(name)
+
+    def info(self, msg: str, *args) -> None:
+        self._log.info(msg, *args)
+
+    def skip(self, msg: str) -> None:
+        self._log.info("skip: %s", msg)
+
+    def select(self, msg: str) -> None:
+        self._log.info("select: %s", msg)
+
+    def insert(self, msg: str) -> None:
+        self._log.info("insert: %s", msg)
+
+    def update(self, msg: str) -> None:
+        self._log.info("update: %s", msg)
+
+    def delete(self, msg: str) -> None:
+        self._log.info("delete: %s", msg)
+        
+    def debug(self, msg: str, *args) -> None:
+        self._log.debug(msg, *args)
+
+    def error(self, msg: str, *args) -> None:
+        self._log.error(msg, *args)
+
+    def exception(self, msg: str) -> None:
+        self._log.exception(msg)
+
+    def stack_trace(self) -> None:
+        self._log.info("stack_trace", exc_info=True)
+
+
+
+def _init_logger(name: str) -> Logger:
+    """Initialize the logger.
+
+    Parameters
+    ----------
+    name : str
+        Name of the logger
+
+    Returns
+    -------
+    Logger
+        Initialized logger instance
+    """
+    return Logger(name)
+
+
+def log_with_print(message: str) -> None:
+    """Print the message with a timestamp prepended to it.
+
+    Parameters
+    ----------
+    message : str
+        message to log
+    """
+    print(f"{datetime.datetime.now(tz=datetime.timezone.utc)} {message}")
+
+
+if __name__ == "__main__":
+    logger = _init_logger(__file__)
+
+    logger.info("start processing")
+    logger.skip("update=False - spotify_track_id=56Kjskx12ksQss")
+    logger.select("spotify_artist - tier=1 - artists_count=30")
+    logger.insert("spotify_artist_insights - id=3963, timestp=2022-08-30")
+    logger.update("spotify_artist_insights - id=3963, timestp=2022-08-30")
+    logger.delete("spotify_artist_insights - id=3963, timestp=2022-08-30")
+    logger.error("spotify_artist - tier=1 - artists_count=30")
+
+    try:
+        x = [1, 2, 3]
+        print(x[5])  # This will raise an IndexError
+    except IndexError:
+        logger.exception("As expected, an error occurred while accessing the list.")
+
+    logger.stack_trace()

qoptimizer_utils-0.1.0/db_utils/neo4j_access.py

@@ -0,0 +1,146 @@
+"""Neo4j connection code."""
+
+from contextlib import contextmanager
+
+from neo4j import GraphDatabase
+
+from .config import get_env_var
+from .logger import _init_logger
+
+logger = _init_logger(__file__)
+
+
+def user_id_to_neo4j_db(user_id: int | str) -> str:
+    """Convert a user id to a Neo4j database name.
+
+    Neo4j database names must be lowercase alphanumeric + dots.
+    Format: ``user.<id>`` for int ids, ``user.<uuid_hex>`` for legacy UUID strings.
+
+    Example: ``42`` → ``"user.42"``
+    """
+    if isinstance(user_id, int):
+        return f"user.{user_id}"
+    # Legacy UUID string support
+    hex_id = str(user_id).replace("-", "").lower()
+    return f"user.{hex_id}"
+
+
+def ensure_user_database(driver, user_id: str) -> str:
+    """Create a per-user Neo4j database if it does not exist.
+
+    Must run against the ``system`` database.
+    Returns the database name, or ``None`` if the instance does not support
+    multi-database (e.g. Neo4j Aura Free/Pro).
+    """
+    db_name = user_id_to_neo4j_db(user_id)
+    try:
+        with driver.session(database="system") as session:
+            session.run(f"CREATE DATABASE `{db_name}` IF NOT EXISTS")
+        logger.info("ensure_user_database: created %s", db_name)
+        return db_name
+    except Exception as e:
+        if "not supported" in str(e).lower() or "UnsupportedAdministration" in str(e):
+            logger.info(
+                "ensure_user_database: CREATE DATABASE not supported "
+                "(Neo4j Aura?). Using default database.",
+            )
+            return None
+        raise
+
+
+def get_neo4j_config(**kwargs) -> dict:
+    """Get Neo4j configuration.
+
+    Parameters
+    ----------
+    kwargs
+        Customizable Neo4j configuration keys:
+        - uri
+        - user
+        - password
+
+    Returns
+    -------
+    dict
+        Neo4j configuration
+    """
+    config = {
+        "uri": get_env_var("neo4j_uri"),
+        "user": get_env_var("neo4j_user") or "neo4j",
+        "password": get_env_var("neo4j_password", secret=True),
+    }
+    config.update(kwargs)
+    return config
+
+
+def neo4j_connect(**kwargs):
+    """Get a Neo4j driver.
+
+    Returns
+    -------
+    neo4j.BoltDriver
+        Neo4j driver
+    """
+    config = get_neo4j_config(**kwargs)
+    return GraphDatabase.driver(
+        config["uri"],
+        auth=(config["user"], config["password"]),
+    )
+
+
+@contextmanager
+def get_neo4j_session(database: str | None = None, **kwargs):
+    """Get Neo4j session context (auto-closes driver).
+
+    Parameters
+    ----------
+    database : str, optional
+        Target Neo4j database name. ``None`` uses the server default.
+
+    Yields
+    ------
+    neo4j.Session
+        Neo4j session
+    """
+    driver = neo4j_connect(**kwargs)
+    try:
+        session_kwargs = {}
+        if database:
+            session_kwargs["database"] = database
+        with driver.session(**session_kwargs) as session:
+            yield session
+    finally:
+        driver.close()
+
+
+def run_neo4j_query(query: str, params: dict = None, database: str | None = None, **kwargs):
+    """Run a Cypher query and return results.
+
+    Parameters
+    ----------
+    query : str
+        Cypher query
+    params : dict, optional
+        Query parameters
+    database : str, optional
+        Target Neo4j database name.
+    **kwargs
+        Additional keyword arguments for connection
+
+    Returns
+    -------
+    list
+        List of record dicts
+    """
+    with get_neo4j_session(database=database, **kwargs) as session:
+        result = session.run(query, params or {})
+        return [dict(record) for record in result]
+
+
+if __name__ == "__main__":
+    driver = neo4j_connect()
+    with driver.session() as s:
+        result = s.run("MATCH (n) RETURN labels(n)[0] AS label, count(n) AS cnt ORDER BY label")
+        for rec in result:
+            print(f"  {rec['label']}: {rec['cnt']}")
+    driver.close()

qoptimizer_utils-0.1.0/db_utils/postgres/__init__.py

@@ -0,0 +1 @@
+"""Postgres."""

qoptimizer_utils-0.1.0/db_utils/postgres/connect.py

@@ -0,0 +1,138 @@
+"""Postgres connection code."""
+
+from collections.abc import Generator
+from contextlib import contextmanager
+from typing import Any, Optional
+
+import psycopg2
+from psycopg2 import connect as _connect
+from psycopg2.extensions import UNICODE, connection
+from psycopg2.extras import RealDictCursor
+from psycopg2.pool import ThreadedConnectionPool
+
+from .helpers import MAX_CLIENTS, get_config
+
+psycopg2.extensions.register_type(UNICODE)
+
+
+def db_connect(**kwargs: dict) -> connection:
+    """Connect Postgres (Reader).
+
+    Parameters
+    ----------
+    user : str, optional
+        Postgres user, by default DB_USER
+
+    Returns
+    -------
+    psycopg2.extensions.connection
+        Postgres connection
+    """
+    config = get_config(**kwargs)
+    return _connect(
+        database=config["database"],
+        user=config["user"],
+        password=config["password"],
+        host=config["host"],
+        port=config["port"],
+    )
+
+
+@contextmanager
+def get_connection(**kwargs: dict) -> Generator[connection, None, None]:
+    """Get Postgres connection (Reader).
+
+    Parameters
+    ----------
+    user : str, optional
+        Postgres user, by default DB_USER
+
+    Yields
+    ------
+    psycopg2.extensions.connection
+        Postgres connection
+    """
+    config = get_config(**kwargs)
+    pool = ThreadedConnectionPool(
+        1,
+        MAX_CLIENTS,
+        database=config["database"],
+        user=config["user"],
+        password=config["password"],
+        host=config["host"],
+        port=config["port"],
+    )
+    try:
+        conn = pool.getconn()
+        yield conn
+    finally:
+        pool.putconn(conn, close=True)
+
+
+@contextmanager
+def get_cursor(commit: Optional[bool] = None, **kwargs: dict) -> Generator[RealDictCursor, None, None]:
+    """Get Postgres cursor (Reader).
+
+    Parameters
+    ----------
+    commit : boolean, optional
+        whether to auto-commit at the end, by default False
+
+    Yields
+    ------
+    psycopg2.extras.RealDictCursor
+        Postgres cursor
+    """
+    with get_connection(**kwargs) as conn:
+        cursor = conn.cursor(cursor_factory=RealDictCursor)
+        try:
+            yield cursor
+            if commit:
+                conn.commit()
+        finally:
+            cursor.close()
+
+
+def run_pg_sql(
+    query: str, *, writer: bool = False, return_results: bool = False, **kwargs: dict
+) -> Optional[list[dict[str, Any]]]:
+    """Get data from the database for the specified query.
+
+    Parameters
+    ----------
+    query : str
+        Query
+    user : str, optional
+        Postgres user, by default DB_USER
+    writer : bool, optional
+        whether to connect to the writer node, by default False
+    return_results : bool, optional
+        whether to return the query results, by default False
+
+    Returns
+    -------
+    list
+        list of {column: value} dicts
+    """
+    with get_connection(writer=writer, **kwargs) as con, con.cursor(cursor_factory=RealDictCursor) as cur:
+        cur.execute(query)
+        if writer:
+            con.commit()
+        if return_results:
+            return cur.fetchall()
+        return None
+
+
+if __name__ == "__main__":
+    con = db_connect()
+    cur = con.cursor()
+    cur.execute("select current_date;")
+    print(cur.fetchall())
+    con.close()
+
+    with get_connection() as pg_con:
+        cur = pg_con.cursor()
+        cur.execute("select current_date;")
+        print(cur.fetchall())
+
+    print(run_pg_sql("select current_date;", return_results=True))