ry-pg-utils 1.0.2__py3-none-any.whl

ry_pg_utils/config.py

@@ -0,0 +1,44 @@
+import os
+import socket
+from dataclasses import dataclass
+
+import dotenv
+
+
+@dataclass
+class Config:
+    postgres_host: str | None
+    postgres_port: int | None
+    postgres_db: str | None
+    postgres_user: str | None
+    postgres_password: str | None
+    do_publish_db: bool
+    use_local_db_only: bool
+    backend_id: str
+    add_backend_to_all: bool
+    add_backend_to_tables: bool
+    raise_on_use_before_init: bool
+
+
+dotenv.load_dotenv()
+
+# Parse POSTGRES_PORT with proper None handling for mypy
+_postgres_port_str = os.getenv("POSTGRES_PORT")
+_postgres_port = int(_postgres_port_str) if _postgres_port_str is not None else None
+
+pg_config = Config(
+    postgres_host=os.getenv("POSTGRES_HOST"),
+    postgres_port=_postgres_port,
+    postgres_db=os.getenv("POSTGRES_DB"),
+    postgres_user=os.getenv("POSTGRES_USER"),
+    postgres_password=os.getenv("POSTGRES_PASSWORD"),
+    do_publish_db=False,
+    use_local_db_only=True,
+    backend_id=(
+        os.getenv("POSTGRES_USER")
+        or f"{socket.gethostname()}_{socket.gethostbyname(socket.gethostname())}"
+    ),
+    add_backend_to_all=True,
+    add_backend_to_tables=True,
+    raise_on_use_before_init=True,
+)

ry_pg_utils/connect.py

@@ -0,0 +1,288 @@
+import contextlib
+import importlib
+import pkgutil
+import threading
+import typing as T
+
+from ryutils import log
+from sqlalchemy import Column, String, create_engine, event
+from sqlalchemy.engine.base import Engine
+from sqlalchemy.exc import OperationalError
+from sqlalchemy.orm import declarative_base, declared_attr, scoped_session, sessionmaker
+from sqlalchemy.orm.decl_api import DeclarativeMeta
+from sqlalchemy.orm.scoping import ScopedSession
+from sqlalchemy_utils import database_exists
+from tenacity import retry, stop_after_attempt, wait_exponential
+
+from ry_pg_utils import config
+
+_thread_local = threading.local()
+BACKEND_ID_VARIABLE = "backend_id"
+
+ENGINE: T.Dict[str, Engine] = {}
+THREAD_SAFE_SESSION_FACTORY: T.Dict[str, ScopedSession] = {}
+
+# Base class with optional backend_id field
+if config.pg_config.add_backend_to_all:
+    # Add any common fields here
+    class CommonBaseModel:
+        @declared_attr
+        def backend_id(cls: T.Any) -> Column:  # pylint: disable=no-self-argument
+            return Column(String(256), nullable=False)
+
+    Base = declarative_base(name="Base", cls=CommonBaseModel)
+else:
+    Base = declarative_base(name="Base")
+
+# Add type annotation for Base
+Base: DeclarativeMeta  # type: ignore
+
+
+def get_table_name(
+    base_name: str, verbose: bool = False, backend_id: str = config.pg_config.backend_id
+) -> str:
+    if verbose:
+        print(
+            f"{base_name}_{backend_id}"
+            if config.pg_config.add_backend_to_tables
+            else f"{base_name}"
+        )
+    return f"{base_name}_{backend_id}" if config.pg_config.add_backend_to_tables else base_name
+
+
+def init_engine(uri: str, db: str, **kwargs: T.Any) -> Engine:
+    global ENGINE  # pylint: disable=global-variable-not-assigned
+    if db not in ENGINE:
+        # Add pool settings to automatically recycle connections
+        default_pool_settings = {
+            "pool_recycle": 3600,  # Recycle connections after 1 hour
+            "pool_pre_ping": True,  # Enable connection health checks
+            "pool_size": 5,  # Maintain a pool of connections
+            "max_overflow": 10,  # Allow up to 10 additional connections
+        }
+        # Update kwargs with defaults if not already set
+        for key, value in default_pool_settings.items():
+            kwargs.setdefault(key, value)
+        ENGINE[db] = create_engine(uri, **kwargs)
+    return ENGINE[db]
+
+
+def get_engine(db: str) -> Engine:
+    global ENGINE  # pylint: disable=global-variable-not-assigned
+    return ENGINE[db]
+
+
+def clear_db() -> None:
+    global ENGINE  # pylint: disable=global-statement
+    global THREAD_SAFE_SESSION_FACTORY  # pylint: disable=global-statement
+    ENGINE = {}
+    THREAD_SAFE_SESSION_FACTORY = {}
+
+
+def close_engine(db: str) -> None:
+    global ENGINE  # pylint: disable=global-statement, global-variable-not-assigned
+    global THREAD_SAFE_SESSION_FACTORY  # pylint: disable=global-statement, global-variable-not-assigned
+    if db in ENGINE:
+        ENGINE[db].dispose()
+        del ENGINE[db]
+    if db in THREAD_SAFE_SESSION_FACTORY:
+        del THREAD_SAFE_SESSION_FACTORY[db]
+
+
+@retry(
+    stop=stop_after_attempt(3),
+    wait=wait_exponential(multiplier=1, min=4, max=10),
+    retry=lambda e: isinstance(e, (OperationalError, TimeoutError)),
+)
+def _init_session_factory(db: str) -> ScopedSession:
+    """Initialize the THREAD_SAFE_SESSION_FACTORY."""
+    global ENGINE, THREAD_SAFE_SESSION_FACTORY  # pylint: disable=global-variable-not-assigned
+    if db not in ENGINE:
+        raise ValueError(
+            "Initialize ENGINE by calling init_engine before calling _init_session_factory!"
+        )
+    if db not in THREAD_SAFE_SESSION_FACTORY:
+        session_factory = sessionmaker(bind=ENGINE[db])
+        THREAD_SAFE_SESSION_FACTORY[db] = scoped_session(session_factory)
+    return THREAD_SAFE_SESSION_FACTORY[db]
+
+
+def set_backend_id(backend_id: str) -> None:
+    setattr(_thread_local, BACKEND_ID_VARIABLE, backend_id)
+
+
+def get_backend_id() -> T.Optional[str]:
+    return getattr(_thread_local, BACKEND_ID_VARIABLE, None)
+
+
+@event.listens_for(scoped_session, "before_flush")
+def receive_before_flush(session: ScopedSession, _flush_context: T.Any, _instances: T.Any) -> None:
+    backend_id = get_backend_id()
+    if not backend_id:
+        return
+
+    # Automatically add backend_id to instances that have it as a field
+    for instance in session.dirty:
+        if hasattr(instance, "backend_id") and instance.backend_id is None:
+            instance.backend_id = backend_id
+
+    for instance in session.new:
+        if hasattr(instance, "backend_id") and instance.backend_id is None:
+            instance.backend_id = backend_id
+
+
+def is_session_factory_initialized() -> bool:
+    return bool(THREAD_SAFE_SESSION_FACTORY)
+
+
+@contextlib.contextmanager
+def ManagedSession(  # pylint: disable=invalid-name
+    db: T.Optional[str] = None, backend_id: T.Optional[str] = config.pg_config.backend_id
+) -> T.Iterator[T.Optional[ScopedSession]]:
+    """Get a session object whose lifecycle, commits and flush are managed for you.
+    The session will automatically retry operations on connection errors.
+
+    Expected to be used as follows:
+    ```
+    # multiple db_operations are done within one session.
+    with ManagedSession() as session:
+        # db_operations is expected not to worry about session handling.
+        db_operations.select(session, **kwargs)
+        # after the with statement, the session commits to the database.
+        db_operations.insert(session, **kwargs)
+    ```
+    """
+    global THREAD_SAFE_SESSION_FACTORY  # pylint: disable=global-variable-not-assigned
+    if db is None:
+        # assume we're just using the default db
+        db = list(THREAD_SAFE_SESSION_FACTORY.keys())[0]
+
+    if db not in THREAD_SAFE_SESSION_FACTORY:
+        if config.pg_config.raise_on_use_before_init:
+            raise ValueError(f"Call _init_session_factory for {db} before using ManagedSession!")
+        log.print_fail(f"Call _init_session_factory for {db} before using ManagedSession!")
+        yield None
+        return
+
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, min=4, max=10),
+        retry=lambda e: isinstance(e, OperationalError),
+    )
+    def execute_with_retry(session: ScopedSession) -> T.Iterator[ScopedSession]:
+        try:
+            yield session
+            session.commit()
+            session.flush()
+        except Exception:
+            session.rollback()
+            raise
+
+    session = THREAD_SAFE_SESSION_FACTORY[db]()
+
+    if backend_id:
+        set_backend_id(backend_id)
+
+    try:
+        yield from execute_with_retry(session)
+    finally:
+        # source:
+        # https://stackoverflow.com/questions/
+        # 21078696/why-is-my-scoped-session-raising-an-attributeerror-session-object-has-no-attr
+        THREAD_SAFE_SESSION_FACTORY[db].remove()
+
+
+def is_database_initialized(db: str) -> bool:
+    """Check if the database is initialized."""
+    global THREAD_SAFE_SESSION_FACTORY  # pylint: disable=global-variable-not-assigned
+    return db in THREAD_SAFE_SESSION_FACTORY
+
+
+def _import_models_from_module(module_path: str) -> None:
+    """
+    Dynamically import all model classes from a given module path.
+
+    This ensures all SQLAlchemy models are registered with Base.metadata
+    before creating tables.
+
+    Args:
+        module_path: Dot-separated module path (e.g., 'database.models')
+    """
+    try:
+        # Import the base module
+        base_module = importlib.import_module(module_path)
+
+        # Get the package path
+        if hasattr(base_module, "__path__"):
+            package_path = base_module.__path__
+        else:
+            # It's a single module, not a package
+            return
+
+        # Walk through all submodules
+        for _, modname, _ in pkgutil.walk_packages(
+            path=package_path,
+            prefix=f"{module_path}.",
+        ):
+            try:
+                importlib.import_module(modname)
+            except Exception as e:  # pylint: disable=broad-except
+                log.print_warn(f"Failed to import {modname}: {e}")
+                continue
+
+        log.print_ok_blue(f"Imported models from {module_path}")
+    except ModuleNotFoundError:
+        log.print_warn(f"Models module not found: {module_path}")
+    except Exception as e:  # pylint: disable=broad-except
+        log.print_warn(f"Error importing models from {module_path}: {e}")
+
+
+def init_database(
+    db_name: str,
+    db_user: str = "",
+    db_password: str = "",
+    db_host: str = "localhost",
+    db_port: int = 5432,
+    models_module: T.Optional[str] = None,
+) -> None:
+    """
+    Initialize a database connection and create tables.
+
+    Args:
+        db_name: Name of the database
+        db_user: Database username
+        db_password: Database password
+        db_host: Database host
+        db_port: Database port
+        models_module: Optional dot-separated module path to your models
+                      (e.g., 'database.models' or 'myapp.db.models').
+                      All model classes in this module will be automatically
+                      imported to register them with SQLAlchemy.
+    """
+    log.print_normal(f"Initializing database {db_name} at {db_host}:{db_port}")
+
+    if db_user and db_password:
+        uri = f"postgresql://{db_user}:{db_password}@{db_host}:{db_port}/{db_name}"
+    elif db_user:
+        uri = f"postgresql://{db_user}@{db_host}:{db_port}/{db_name}"
+    else:
+        uri = f"postgresql://{db_host}:{db_port}/{db_name}"
+
+    engine = init_engine(uri, db_name)
+
+    if database_exists(engine.url):
+        log.print_normal("Found existing database")
+    else:
+        log.print_ok_blue("Creating new database!")
+
+    # Import models if module path provided
+    if models_module:
+        _import_models_from_module(models_module)
+
+    try:
+        Base.metadata.create_all(bind=engine)
+
+        _init_session_factory(db_name)
+    except OperationalError as exc:
+        log.print_fail(f"Failed to initialize database: {exc}")
+        log.print_normal("Continuing without db connection...")

ry_pg_utils/dynamic_table.py

@@ -0,0 +1,199 @@
+import typing as T
+from datetime import datetime, timezone
+
+from google.protobuf.descriptor import FieldDescriptor
+from google.protobuf.message import Message
+from ryutils import log
+from sqlalchemy import (
+    Column,
+    DateTime,
+    Integer,
+    LargeBinary,
+    MetaData,
+    String,
+    Table,
+    func,
+    insert,
+    inspect,
+    types,
+)
+from sqlalchemy.engine import Engine
+from sqlalchemy.exc import OperationalError
+
+from ry_pg_utils.config import pg_config
+from ry_pg_utils.connect import (
+    BACKEND_ID_VARIABLE,
+    ENGINE,
+    ManagedSession,
+    get_table_name,
+    init_engine,
+)
+
+FIELD_TYPE_MAP: T.Dict[int, str] = {
+    FieldDescriptor.TYPE_DOUBLE: "float",
+    FieldDescriptor.TYPE_FLOAT: "float",
+    FieldDescriptor.TYPE_INT64: "int",
+    FieldDescriptor.TYPE_UINT64: "int",
+    FieldDescriptor.TYPE_INT32: "int",
+    FieldDescriptor.TYPE_FIXED64: "int",
+    FieldDescriptor.TYPE_FIXED32: "int",
+    FieldDescriptor.TYPE_BOOL: "bool",
+    FieldDescriptor.TYPE_STRING: "string",
+    FieldDescriptor.TYPE_GROUP: "string",
+    FieldDescriptor.TYPE_MESSAGE: "message",
+    FieldDescriptor.TYPE_BYTES: "binary",
+    FieldDescriptor.TYPE_UINT32: "int",
+    FieldDescriptor.TYPE_ENUM: "int",
+    FieldDescriptor.TYPE_SFIXED32: "int",
+    FieldDescriptor.TYPE_SFIXED64: "int",
+    FieldDescriptor.TYPE_SINT32: "int",
+    FieldDescriptor.TYPE_SINT64: "int",
+}
+
+ADD_BACKEND_TO_ALL = pg_config.add_backend_to_all
+
+
+def _get_field_types(message_class: T.Type) -> T.Dict[str, str]:
+    """Retrieve the types of each field in a Protobuf message class."""
+    return {field.name: FIELD_TYPE_MAP[field.type] for field in message_class.DESCRIPTOR.fields}
+
+
+def _combine_pb_timestamp(seconds: int, nanos: int) -> datetime:
+    """Combines Protobuf seconds and nanos into a single datetime object."""
+    total_seconds = seconds + nanos / 1e9
+    return datetime.fromtimestamp(total_seconds, tz=timezone.utc)
+
+
+def _create_dynamic_table(channel_name: str, pb_message: T.Type, db_name: str) -> Table:
+    engine: T.Optional[Engine] = ENGINE.get(db_name)
+    if not engine:
+        raise ValueError(f"Database {db_name!r} not initialized")
+    if not pb_message:
+        raise ValueError(f"Protobuf message {pb_message!r} invalid")
+
+    # attempt inspect
+    try:
+        with engine.connect() as conn:
+            inspector = inspect(conn)
+            existing = inspector.get_table_names()
+    except OperationalError:
+        engine.dispose()
+        engine = init_engine(str(engine.url), db_name)
+        with engine.connect() as conn:
+            inspector = inspect(conn)
+            existing = inspector.get_table_names()
+
+    tbl_name = get_table_name(channel_name)
+    metadata = MetaData()
+    if tbl_name in existing:
+        return Table(tbl_name, metadata, autoload_with=engine)
+
+    cols: T.List[Column] = [
+        Column("key", Integer, primary_key=True, autoincrement=True),
+        Column("created_at", DateTime, server_default=func.now()),  # pylint: disable=not-callable
+    ]
+    for name, ftype in _get_field_types(pb_message).items():
+        if ftype == "string":
+            cols.append(Column(name, String))
+        elif ftype == "int":
+            cols.append(Column(name, Integer))
+        elif ftype == "binary":
+            cols.append(Column(name, LargeBinary))
+        elif ftype == "bool":
+            cols.append(Column(name, types.Boolean))
+        elif ftype == "float":
+            cols.append(Column(name, types.Float))
+        elif ftype == "message":
+            cols.append(
+                Column(
+                    name,
+                    types.DateTime,
+                    server_default=func.now(),  # pylint: disable=not-callable
+                )
+            )
+        else:
+            raise ValueError(f"Unsupported field type: {ftype!r}")
+    if ADD_BACKEND_TO_ALL:
+        cols.append(Column(BACKEND_ID_VARIABLE, String(256), nullable=False))
+
+    tbl = Table(tbl_name, metadata, *cols, extend_existing=True)
+    metadata.create_all(engine)
+    return tbl
+
+
+class DynamicTableDb:
+    def __init__(self, db_name: str) -> None:
+        self.db_name = db_name
+
+    @staticmethod
+    def is_in_db(msg: Message, db_name: str, channel: str, attr: str, value: T.Any) -> bool:
+        """Static entry-point for existence check."""
+        return DynamicTableDb(db_name).inst_is_in_db(msg, channel, attr, value)
+
+    def inst_is_in_db(
+        self,
+        message_pb: Message,
+        channel_name: str,
+        attr: str,
+        value: T.Any,
+    ) -> bool:
+        try:
+            tbl = _create_dynamic_table(channel_name, type(message_pb), self.db_name)
+        except ValueError as e:
+            print(f"is_in_db error: {e}")  # explicitly use print() here to avoid circular call
+            return False
+        if not hasattr(tbl.c, attr):
+            raise ValueError(
+                f"Column '{attr}' missing in '{channel_name}', cols={list(tbl.c.keys())}"
+            )
+        with ManagedSession(db=self.db_name) as sess:
+            if sess is None:
+                return False
+            stmt = tbl.select().where(getattr(tbl.c, attr) == value)
+            return bool(sess.execute(stmt).fetchone())
+
+    @staticmethod
+    def log_data_to_db(
+        msg: Message,
+        db_name: str,
+        channel: str,
+        log_print_failure: bool = True,
+    ) -> None:
+        """Static entry-point for logging messages."""
+        assert db_name.strip(), "db_name is required"
+        DynamicTableDb(db_name).add_message(channel, msg, log_print_failure)
+
+    def add_message(
+        self,
+        channel_name: str,
+        message_pb: Message,
+        log_print_failure: bool = True,
+        verbose: bool = False,
+    ) -> None:
+        printer = log.print_fail if log_print_failure else print
+        try:
+            tbl = _create_dynamic_table(channel_name, type(message_pb), self.db_name)
+        except ValueError as e:
+            if verbose:
+                printer(f"add_message table error: {e}")
+            return
+        data = self.protobuf_to_dict(message_pb)
+        with ManagedSession(db=self.db_name) as sess:
+            if sess is None:
+                return
+            stmt = insert(tbl).values(**data)
+            try:
+                sess.execute(stmt)
+            except Exception as e:  # pylint: disable=broad-exception-caught
+                if verbose:
+                    printer(f"insert failed: {e}")
+
+    def protobuf_to_dict(self, message_pb: Message) -> T.Dict[str, T.Any]:
+        out: T.Dict[str, T.Any] = {}
+        for fld in message_pb.DESCRIPTOR.fields:
+            val = getattr(message_pb, fld.name)
+            if fld.type == FieldDescriptor.TYPE_MESSAGE and fld.message_type.name == "Timestamp":
+                out[fld.name] = _combine_pb_timestamp(val.seconds, val.nanos)
+            else:
+                out[fld.name] = val
+        return out

ry_pg_utils/ipc/channels.py

@@ -0,0 +1,14 @@
+from ry_redis_bus.channels import Channel
+
+from ry_pg_utils.pb_types.database_pb2 import DatabaseConfigPb  # pylint: disable=no-name-in-module
+from ry_pg_utils.pb_types.database_pb2 import (  # pylint: disable=no-name-in-module
+    DatabaseNotificationPb,
+)
+from ry_pg_utils.pb_types.database_pb2 import (  # pylint: disable=no-name-in-module
+    DatabaseSettingsPb,
+)
+
+# Channels
+DATABASE_CHANNEL = Channel("DATABASE_CHANNEL", DatabaseConfigPb)
+DATABASE_CONFIG_CHANNEL = Channel("DATABASE_CONFIG_CHANNEL", DatabaseSettingsPb)
+DATABASE_NOTIFY_CHANNEL = Channel("DATABASE_NOTIFY_CHANNEL", DatabaseNotificationPb)