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

ry_pg_utils/notify_trigger.py

@@ -0,0 +1,346 @@
+import json
+import threading
+import time
+from contextlib import contextmanager
+from typing import Any, Callable, Dict, Generator, Iterator, List, Optional, Set
+
+import psycopg2
+from psycopg2.extensions import ISOLATION_LEVEL_AUTOCOMMIT
+from sqlalchemy import text
+from sqlalchemy.engine import Engine
+
+from ry_pg_utils.connect import get_engine
+
+
+def create_notify_trigger(
+    engine: Engine,
+    table_name: str,
+    channel_name: Optional[str] = None,
+    trigger_name: Optional[str] = None,
+    events: Optional[List[str]] = None,
+    columns: Optional[List[str]] = None,
+) -> None:
+    """
+    Create a notification trigger on a specified table.
+
+    Args:
+        engine: SQLAlchemy engine instance
+        table_name: Name of the table to add the trigger to
+        channel_name: Name of the notification channel (defaults to table_name)
+        trigger_name: Name of the trigger (defaults to f"{table_name}_notify_trigger")
+        events: List of events to trigger on (defaults to ['INSERT', 'UPDATE', 'DELETE'])
+        columns: List of columns to include in the notification payload. If None, all columns.
+    """
+    if events is None:
+        events = ["INSERT", "UPDATE", "DELETE"]
+
+    # Validate events
+    valid_events = {"INSERT", "UPDATE", "DELETE"}
+    invalid_events = set(events) - valid_events
+    if invalid_events:
+        raise ValueError(f"Invalid events: {invalid_events}. Valid events are: {valid_events}")
+
+    channel_name = channel_name or table_name
+    trigger_name = trigger_name or f"{table_name}_notify_trigger"
+
+    # Validate requested columns
+    if columns is not None:
+        with engine.connect() as conn:
+            result = conn.execute(
+                text(
+                    """
+                    SELECT column_name
+                    FROM information_schema.columns
+                    WHERE table_name = :table_name
+                    AND table_schema = current_schema()
+                """
+                ),
+                {"table_name": table_name},
+            )
+            available_columns = {row[0] for row in result}
+            invalid_columns = set(columns) - available_columns
+            if invalid_columns:
+                raise ValueError(
+                    f"Invalid columns: {invalid_columns}. "
+                    f"Available columns are: {available_columns}"
+                )
+
+    # Build JSON payload expression
+    if columns:
+        data_builder = (
+            "json_build_object("
+            + ", ".join(
+                f"'{col}', CASE WHEN TG_OP = 'DELETE' THEN OLD.{col} ELSE NEW.{col} END"
+                for col in columns
+            )
+            + ")"
+        )
+    else:
+        data_builder = "row_to_json(CASE WHEN TG_OP = 'DELETE' THEN OLD ELSE NEW END)"
+
+    trigger_function = f"""
+    CREATE OR REPLACE FUNCTION "{trigger_name}_function"()
+    RETURNS TRIGGER AS $$
+    BEGIN
+        PERFORM pg_notify(
+            :channel,
+            json_build_object(
+                'table', TG_TABLE_NAME,
+                'action', TG_OP,
+                'data', {data_builder}
+            )::text
+        );
+        RETURN CASE WHEN TG_OP = 'DELETE' THEN OLD ELSE NEW END;
+    END;
+    $$ LANGUAGE plpgsql;
+    """
+
+    drop_commands = f"""
+    DROP TRIGGER IF EXISTS "{trigger_name}_insert" ON "{table_name}";
+    DROP TRIGGER IF EXISTS "{trigger_name}_update" ON "{table_name}";
+    DROP TRIGGER IF EXISTS "{trigger_name}_delete" ON "{table_name}";
+    DROP FUNCTION IF EXISTS "{trigger_name}_function"();
+    """
+
+    with engine.begin() as conn:
+        conn.execute(text(drop_commands))
+        conn.execute(text(trigger_function), {"channel": channel_name})
+        for event in events:
+            trigger_sql = f"""
+            CREATE TRIGGER "{trigger_name}_{event.lower()}"
+            AFTER {event} ON "{table_name}"
+            FOR EACH ROW EXECUTE FUNCTION "{trigger_name}_function"();
+            """
+            conn.execute(text(trigger_sql))
+
+
+def drop_notify_trigger(
+    engine: Engine,
+    table_name: str,
+    trigger_name: Optional[str] = None,
+) -> None:
+    """
+    Drop a notification trigger and its associated function.
+
+    Args:
+        engine: SQLAlchemy engine instance
+        table_name: Name of the table the trigger is on
+        trigger_name: Name of the trigger (defaults to f"{table_name}_notify_trigger")
+    """
+    trigger_name = trigger_name or f"{table_name}_notify_trigger"
+
+    drop_commands = f"""
+    DROP TRIGGER IF EXISTS "{trigger_name}_insert" ON "{table_name}";
+    DROP TRIGGER IF EXISTS "{trigger_name}_update" ON "{table_name}";
+    DROP TRIGGER IF EXISTS "{trigger_name}_delete" ON "{table_name}";
+    DROP FUNCTION IF EXISTS "{trigger_name}_function"();
+    """
+
+    with engine.begin() as conn:
+        conn.execute(text(drop_commands))
+
+
+@contextmanager
+def subscribe_to_notifications(
+    engine: Engine,
+    channel_name: str,
+    callback: Optional[Callable[[Dict[str, Any]], None]] = None,
+    timeout: float = 60.0,
+) -> Iterator[Generator[Dict[str, Any], None, None]]:
+    """
+    Subscribe to PostgreSQL notifications from a channel. Clean shutdown on exit.
+
+    If a callback is given, notifications are processed in a background thread.
+    Otherwise, notifications can be consumed via the yielded generator.
+    """
+    conn_params = engine.url.translate_connect_args()
+    conn = psycopg2.connect(
+        dbname=conn_params.get("database"),
+        user=conn_params.get("username"),
+        password=conn_params.get("password"),
+        host=conn_params.get("host"),
+        port=conn_params.get("port"),
+    )
+    conn.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
+
+    stop_flag = threading.Event()
+
+    def notification_generator() -> Generator[Dict[str, Any], None, None]:
+        start_time = time.time()
+        with conn.cursor() as cur:
+            cur.execute(f'LISTEN "{channel_name}";')
+            while not stop_flag.is_set() and (time.time() - start_time) < timeout:
+                conn.poll()
+                while conn.notifies:
+                    notify = conn.notifies.pop()
+                    try:
+                        notification = json.loads(notify.payload)
+                        if callback:
+                            callback(notification)
+                        else:
+                            yield notification
+                    except json.JSONDecodeError:
+                        continue
+                time.sleep(0.1)
+
+    try:
+        if callback:
+            thread = threading.Thread(target=lambda: list(notification_generator()), daemon=True)
+            thread.start()
+            yield notification_generator()  # yield generator even in callback mode
+            stop_flag.set()
+            thread.join(timeout=1.0)
+        else:
+            yield notification_generator()
+    finally:
+        if not conn.closed:
+            conn.close()
+
+
+class NotificationListener:
+    """
+    A class to handle PostgreSQL notifications in the background.
+
+    Example usage:
+    ```python
+    # Create a listener
+    listener = NotificationListener(engine)
+
+    # Create a listener for a specific table
+    listener.create_listener(table_name="my_table", channel_name="my_channel")
+
+    # Add a callback for a specific channel
+    def handle_changes(notification):
+        print(f"Received notification: {notification}")
+
+    listener.add_callback("my_table", handle_changes)
+
+    # Start listening in the background
+    listener.start()
+
+    # ... your main application code ...
+
+    # When done, stop the listener
+    listener.stop()
+    ```
+    """
+
+    def __init__(self, db_name: str) -> None:
+        self.db_name = db_name
+        self._running = False
+        self._stop_flag = threading.Event()
+        self._thread: Optional[threading.Thread] = None
+        self._callbacks: Dict[str, Set[Callable[[Dict[str, Any]], None]]] = {}
+        self._lock = threading.Lock()
+
+    def add_callback(self, channel_name: str, callback: Callable[[Dict[str, Any]], None]) -> None:
+        with self._lock:
+            self._callbacks.setdefault(channel_name, set()).add(callback)
+
+    def remove_callback(
+        self, channel_name: str, callback: Callable[[Dict[str, Any]], None]
+    ) -> None:
+        with self._lock:
+            callbacks = self._callbacks.get(channel_name)
+            if callbacks:
+                callbacks.discard(callback)
+                if not callbacks:
+                    del self._callbacks[channel_name]
+
+    def _process_notification(self, notify: Any) -> None:
+        """Process a single notification and execute callbacks."""
+        try:
+            notification = json.loads(notify.payload)
+            with self._lock:
+                callbacks = self._callbacks.get(notify.channel, set())
+                for cb in callbacks:
+                    try:
+                        cb(notification)
+                    except Exception as e:  # pylint: disable=broad-exception-caught
+                        print(f"Error in callback for channel {notify.channel}: {e}")
+        except json.JSONDecodeError:
+            pass
+
+    def _handle_connection(self, conn: Any, channels: List[str]) -> None:
+        """Handle the active connection and process notifications."""
+        with conn.cursor() as cur:
+            for channel in channels:
+                cur.execute(f'LISTEN "{channel}";')
+
+            while not self._stop_flag.is_set():
+                conn.poll()
+                while conn.notifies:
+                    notify = conn.notifies.pop()
+                    self._process_notification(notify)
+                time.sleep(0.1)
+
+    def _create_connection(self, conn_params: Dict[str, Any]) -> Any:
+        """Create and configure a new database connection."""
+        conn = psycopg2.connect(
+            dbname=conn_params.get("database"),
+            user=conn_params.get("username"),
+            password=conn_params.get("password"),
+            host=conn_params.get("host"),
+            port=conn_params.get("port"),
+        )
+        conn.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
+        return conn
+
+    def _listen_loop(self) -> None:
+        """Main loop for listening to database notifications."""
+        engine = get_engine(self.db_name)
+        conn_params = engine.url.translate_connect_args()
+
+        while not self._stop_flag.is_set():
+            conn = None
+            try:
+                conn = self._create_connection(conn_params)
+                with self._lock:
+                    channels = list(self._callbacks.keys())
+                self._handle_connection(conn, channels)
+            except (psycopg2.OperationalError, psycopg2.InterfaceError) as e:
+                print(f"Listener connection error: {e}")
+                time.sleep(5)
+            finally:
+                if conn and not conn.closed:
+                    conn.close()
+
+    def start(self) -> None:
+        if not self._running:
+            self._stop_flag.clear()
+            self._running = True
+            self._thread = threading.Thread(target=self._listen_loop, daemon=True)
+            self._thread.start()
+
+    def stop(self) -> None:
+        if self._running:
+            self._stop_flag.set()
+            if self._thread is not None:
+                self._thread.join(timeout=2.0)
+            self._thread = None
+            self._running = False
+
+    def create_listener(
+        self,
+        table_name: str,
+        channel_name: str,
+        columns: Optional[List[str]] = None,
+        events: Optional[List[str]] = None,
+    ) -> None:
+        engine = get_engine(self.db_name)
+        create_notify_trigger(
+            engine=engine,
+            table_name=table_name,
+            channel_name=channel_name,
+            trigger_name=f"{table_name}_notify_trigger",
+            events=events,
+            columns=columns,
+        )
+
+    def remove_listener(self, table_name: str) -> None:
+        engine = get_engine(self.db_name)
+        drop_notify_trigger(
+            engine=engine,
+            table_name=table_name,
+            trigger_name=f"{table_name}_notify_trigger",
+        )

ry_pg_utils/parse_args.py

@@ -0,0 +1,15 @@
+import argparse
+
+from . import config
+
+
+def add_postrgres_db_args(parser: argparse.ArgumentParser) -> None:
+    postgres_parser = parser.add_argument_group("postgres-options")
+    postgres_parser.add_argument("--postgres-host", default=config.pg_config.postgres_host)
+    postgres_parser.add_argument(
+        "--postgres-port", type=int, default=config.pg_config.postgres_port
+    )
+    postgres_parser.add_argument("--postgres-db", default=config.pg_config.postgres_db)
+    postgres_parser.add_argument("--postgres-user", default=config.pg_config.postgres_user)
+    postgres_parser.add_argument("--postgres-password", default=config.pg_config.postgres_password)
+    postgres_parser.add_argument("--do-publish-db", action="store_true", default=False)

ry_pg_utils/pb_types/database_pb2.py

@@ -0,0 +1,38 @@
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: database.proto
+"""Generated protocol buffer code."""
+from google.protobuf.internal import builder as _builder
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import symbol_database as _symbol_database
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+from google.protobuf import timestamp_pb2 as google_dot_protobuf_dot_timestamp__pb2
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x0e\x64\x61tabase.proto\x12\x08\x64\x61tabase\x1a\x1fgoogle/protobuf/timestamp.proto\"m\n\nPostgresPb\x12\x10\n\x08\x64\x61tabase\x18\x01 \x01(\t\x12\x0c\n\x04user\x18\x02 \x01(\t\x12\x10\n\x08password\x18\x03 \x01(\t\x12\x0c\n\x04host\x18\x04 \x01(\t\x12\x0c\n\x04port\x18\x05 \x01(\x05\x12\x11\n\tbackendId\x18\x06 \x01(\t\"f\n\x11PostgresMessagePb\x12)\n\x05utime\x18\x01 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12&\n\x08postgres\x18\x02 \x01(\x0b\x32\x14.database.PostgresPb\"\xb2\x01\n\x12\x44\x61tabaseSettingsPb\x12\x17\n\x0fprimaryDatabase\x18\x01 \x01(\x08\x12<\n\x08postgres\x18\x02 \x03(\x0b\x32*.database.DatabaseSettingsPb.PostgresEntry\x1a\x45\n\rPostgresEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12#\n\x05value\x18\x02 \x01(\x0b\x32\x14.database.PostgresPb:\x02\x38\x01\"u\n\x10\x44\x61tabaseConfigPb\x12)\n\x05utime\x18\x01 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\x36\n\x10\x64\x61tabaseSettings\x18\x02 \x01(\x0b\x32\x1c.database.DatabaseSettingsPb\"\x8e\x01\n\x16\x44\x61tabaseNotificationPb\x12)\n\x05utime\x18\x01 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\x12\n\ntable_name\x18\x02 \x01(\t\x12\x14\n\x0c\x63hannel_name\x18\x03 \x01(\t\x12\x0e\n\x06\x61\x63tion\x18\x04 \x01(\t\x12\x0f\n\x07payload\x18\x05 \x01(\tb\x06proto3')
+
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, globals())
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'database_pb2', globals())
+if _descriptor._USE_C_DESCRIPTORS == False:
+
+  DESCRIPTOR._options = None
+  _DATABASESETTINGSPB_POSTGRESENTRY._options = None
+  _DATABASESETTINGSPB_POSTGRESENTRY._serialized_options = b'8\001'
+  _POSTGRESPB._serialized_start=61
+  _POSTGRESPB._serialized_end=170
+  _POSTGRESMESSAGEPB._serialized_start=172
+  _POSTGRESMESSAGEPB._serialized_end=274
+  _DATABASESETTINGSPB._serialized_start=277
+  _DATABASESETTINGSPB._serialized_end=455
+  _DATABASESETTINGSPB_POSTGRESENTRY._serialized_start=386
+  _DATABASESETTINGSPB_POSTGRESENTRY._serialized_end=455
+  _DATABASECONFIGPB._serialized_start=457
+  _DATABASECONFIGPB._serialized_end=574
+  _DATABASENOTIFICATIONPB._serialized_start=577
+  _DATABASENOTIFICATIONPB._serialized_end=719
+# @@protoc_insertion_point(module_scope)

ry_pg_utils/pb_types/database_pb2.pyi

@@ -0,0 +1,156 @@
+"""
+@generated by mypy-protobuf.  Do not edit manually!
+isort:skip_file
+"""
+
+import builtins
+import collections.abc
+import google.protobuf.descriptor
+import google.protobuf.internal.containers
+import google.protobuf.message
+import google.protobuf.timestamp_pb2
+import typing
+
+DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
+
+@typing.final
+class PostgresPb(google.protobuf.message.Message):
+    """Message for Postgres"""
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    DATABASE_FIELD_NUMBER: builtins.int
+    USER_FIELD_NUMBER: builtins.int
+    PASSWORD_FIELD_NUMBER: builtins.int
+    HOST_FIELD_NUMBER: builtins.int
+    PORT_FIELD_NUMBER: builtins.int
+    BACKENDID_FIELD_NUMBER: builtins.int
+    database: builtins.str
+    user: builtins.str
+    password: builtins.str
+    host: builtins.str
+    port: builtins.int
+    backendId: builtins.str
+    def __init__(
+        self,
+        *,
+        database: builtins.str = ...,
+        user: builtins.str = ...,
+        password: builtins.str = ...,
+        host: builtins.str = ...,
+        port: builtins.int = ...,
+        backendId: builtins.str = ...,
+    ) -> None: ...
+    def ClearField(self, field_name: typing.Literal["backendId", b"backendId", "database", b"database", "host", b"host", "password", b"password", "port", b"port", "user", b"user"]) -> None: ...
+
+global___PostgresPb = PostgresPb
+
+@typing.final
+class PostgresMessagePb(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    UTIME_FIELD_NUMBER: builtins.int
+    POSTGRES_FIELD_NUMBER: builtins.int
+    @property
+    def utime(self) -> google.protobuf.timestamp_pb2.Timestamp: ...
+    @property
+    def postgres(self) -> global___PostgresPb: ...
+    def __init__(
+        self,
+        *,
+        utime: google.protobuf.timestamp_pb2.Timestamp | None = ...,
+        postgres: global___PostgresPb | None = ...,
+    ) -> None: ...
+    def HasField(self, field_name: typing.Literal["postgres", b"postgres", "utime", b"utime"]) -> builtins.bool: ...
+    def ClearField(self, field_name: typing.Literal["postgres", b"postgres", "utime", b"utime"]) -> None: ...
+
+global___PostgresMessagePb = PostgresMessagePb
+
+@typing.final
+class DatabaseSettingsPb(google.protobuf.message.Message):
+    """Message for DatabaseSettings"""
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    @typing.final
+    class PostgresEntry(google.protobuf.message.Message):
+        DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+        KEY_FIELD_NUMBER: builtins.int
+        VALUE_FIELD_NUMBER: builtins.int
+        key: builtins.str
+        @property
+        def value(self) -> global___PostgresPb: ...
+        def __init__(
+            self,
+            *,
+            key: builtins.str = ...,
+            value: global___PostgresPb | None = ...,
+        ) -> None: ...
+        def HasField(self, field_name: typing.Literal["value", b"value"]) -> builtins.bool: ...
+        def ClearField(self, field_name: typing.Literal["key", b"key", "value", b"value"]) -> None: ...
+
+    PRIMARYDATABASE_FIELD_NUMBER: builtins.int
+    POSTGRES_FIELD_NUMBER: builtins.int
+    primaryDatabase: builtins.bool
+    @property
+    def postgres(self) -> google.protobuf.internal.containers.MessageMap[builtins.str, global___PostgresPb]: ...
+    def __init__(
+        self,
+        *,
+        primaryDatabase: builtins.bool = ...,
+        postgres: collections.abc.Mapping[builtins.str, global___PostgresPb] | None = ...,
+    ) -> None: ...
+    def ClearField(self, field_name: typing.Literal["postgres", b"postgres", "primaryDatabase", b"primaryDatabase"]) -> None: ...
+
+global___DatabaseSettingsPb = DatabaseSettingsPb
+
+@typing.final
+class DatabaseConfigPb(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    UTIME_FIELD_NUMBER: builtins.int
+    DATABASESETTINGS_FIELD_NUMBER: builtins.int
+    @property
+    def utime(self) -> google.protobuf.timestamp_pb2.Timestamp: ...
+    @property
+    def databaseSettings(self) -> global___DatabaseSettingsPb: ...
+    def __init__(
+        self,
+        *,
+        utime: google.protobuf.timestamp_pb2.Timestamp | None = ...,
+        databaseSettings: global___DatabaseSettingsPb | None = ...,
+    ) -> None: ...
+    def HasField(self, field_name: typing.Literal["databaseSettings", b"databaseSettings", "utime", b"utime"]) -> builtins.bool: ...
+    def ClearField(self, field_name: typing.Literal["databaseSettings", b"databaseSettings", "utime", b"utime"]) -> None: ...
+
+global___DatabaseConfigPb = DatabaseConfigPb
+
+@typing.final
+class DatabaseNotificationPb(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    UTIME_FIELD_NUMBER: builtins.int
+    TABLE_NAME_FIELD_NUMBER: builtins.int
+    CHANNEL_NAME_FIELD_NUMBER: builtins.int
+    ACTION_FIELD_NUMBER: builtins.int
+    PAYLOAD_FIELD_NUMBER: builtins.int
+    table_name: builtins.str
+    channel_name: builtins.str
+    action: builtins.str
+    payload: builtins.str
+    @property
+    def utime(self) -> google.protobuf.timestamp_pb2.Timestamp: ...
+    def __init__(
+        self,
+        *,
+        utime: google.protobuf.timestamp_pb2.Timestamp | None = ...,
+        table_name: builtins.str = ...,
+        channel_name: builtins.str = ...,
+        action: builtins.str = ...,
+        payload: builtins.str = ...,
+    ) -> None: ...
+    def HasField(self, field_name: typing.Literal["utime", b"utime"]) -> builtins.bool: ...
+    def ClearField(self, field_name: typing.Literal["action", b"action", "channel_name", b"channel_name", "payload", b"payload", "table_name", b"table_name", "utime", b"utime"]) -> None: ...
+
+global___DatabaseNotificationPb = DatabaseNotificationPb

ry_pg_utils/postgres_info.py

@@ -0,0 +1,47 @@
+class PostgresInfo:
+    """Class to store the Postgres database information."""
+
+    def __init__(self, db_name: str, user: str, password: str, host: str, port: int) -> None:
+        self.db_name = db_name
+        self.user = user
+        self.password = password
+        self.host = host
+        self.port = port
+
+    @staticmethod
+    def null() -> "PostgresInfo":
+        """Return a null PostgresInfo object."""
+        return PostgresInfo("", "", "", "", 0)
+
+    def is_null(self) -> bool:
+        """Return if the PostgresInfo object is null."""
+        if bool(self == PostgresInfo.null()):
+            return True
+        if not self.db_name.strip() or not self.host.strip() or not self.user.strip():
+            return True
+        if self.port == 0:
+            return True
+        return False
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, PostgresInfo):
+            return False
+
+        return bool(
+            self.db_name == other.db_name
+            and self.user == other.user
+            and self.password == other.password
+            and self.host == other.host
+            and self.port == other.port
+        )
+
+    def __str__(self) -> str:
+        values = [
+            f"{key}={'*' * 8 if key == 'password' else value}"
+            for key, value in self.__dict__.items()
+        ]
+        values_string = "\n\t".join(values)
+        return f"PostgresInfo({values_string})"
+
+    def __repr__(self) -> str:
+        return self.__str__()