ry-pg-utils 1.0.9__py3-none-any.whl → 1.0.11__py3-none-any.whl

ry_pg_utils/connect.py

@@ -5,11 +5,10 @@ import threading
 import typing as T
 
 from ryutils import log
-from sqlalchemy import Column, String, create_engine, event
+from sqlalchemy import 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 import DeclarativeBase, Mapped, mapped_column, scoped_session, sessionmaker
 from sqlalchemy.orm.scoping import ScopedSession
 from sqlalchemy_utils import database_exists
 from tenacity import retry, stop_after_attempt, wait_exponential
@@ -22,20 +21,19 @@ 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
+
+# Modern SQLAlchemy 2.0 declarative base
 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)
+    class Base(DeclarativeBase):
+        """Base class with automatic backend_id field."""
+
+        backend_id: Mapped[str] = mapped_column(String(256), nullable=False)
+
 else:
-    Base = declarative_base(name="Base")
 
-# Add type annotation for Base
-Base: DeclarativeMeta  # type: ignore
+    class Base(DeclarativeBase):  # type: ignore[no-redef]
+        """Base class for all SQLAlchemy models."""
 
 
 def get_table_name(

ry_pg_utils/ipc/channels.py

@@ -1,11 +1,11 @@
 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,
+    PostgresMessagePb,
 )
 
 # Channels
-DATABASE_CHANNEL = Channel("DATABASE_CHANNEL", DatabaseConfigPb)
-DATABASE_CONFIG_CHANNEL = Channel("DATABASE_CONFIG_CHANNEL", DatabaseConfigPb)
+DATABASE_CHANNEL = Channel("DATABASE_CHANNEL", PostgresMessagePb)
+DATABASE_CONFIG_CHANNEL = Channel("DATABASE_CONFIG_CHANNEL", PostgresMessagePb)
 DATABASE_NOTIFY_CHANNEL = Channel("DATABASE_NOTIFY_CHANNEL", DatabaseNotificationPb)

ry_pg_utils/pb_types/database_pb2.py

@@ -14,25 +14,17 @@ _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')
+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\"\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
+  _DATABASENOTIFICATIONPB._serialized_start=277
+  _DATABASENOTIFICATIONPB._serialized_end=419
 # @@protoc_insertion_point(module_scope)

ry_pg_utils/pb_types/database_pb2.pyi

@@ -4,9 +4,7 @@ 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
@@ -66,66 +64,6 @@ class PostgresMessagePb(google.protobuf.message.Message):
 
 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

ry_pg_utils/updater.py

@@ -11,7 +11,6 @@ from ryutils.verbose import Verbose
 from ry_pg_utils.connect import close_engine, init_database, is_database_initialized
 from ry_pg_utils.ipc import channels
 from ry_pg_utils.pb_types.database_pb2 import (  # pylint: disable=no-name-in-module
-    DatabaseConfigPb,
     PostgresMessagePb,
     PostgresPb,
 )
@@ -21,28 +20,15 @@ DB_RETRY_TIME = 5.0
 
 
 def get_database_settings(
-    message_pb: DatabaseConfigPb, backend_id: str, verbose: bool = False
+    raw_db_config_pb: PostgresMessagePb | None, verbose: bool = False
 ) -> T.Optional[PostgresInfo]:
-    # Access the nested databaseSettings field
-    if not message_pb.HasField("databaseSettings"):
-        return None
-
-    database_settings = message_pb.databaseSettings
-
-    # Check if this is the primary database
-    if not database_settings.primaryDatabase:
-        return None
-
-    # Get the postgres config for this backend
-    raw_db_config_pb = database_settings.postgres.get(backend_id, None)
-
     if raw_db_config_pb is None:
         return None
 
-    db_config_pb: PostgresPb = raw_db_config_pb
+    db_config_pb: PostgresPb = raw_db_config_pb.postgres
 
     if verbose:
-        log.print_normal(f"Received database settings {backend_id}:\n{db_config_pb}")
+        log.print_normal(f"Received database settings:\n{db_config_pb}")
 
     database_settings_msg = PostgresInfo(
         db_name=db_config_pb.database,
@@ -59,14 +45,12 @@ def get_database_settings(
 class DbUpdater(RedisClientBase):
     do_publish_db: bool
     postgres_info: PostgresInfo
-    backend_id: str
     logging_error_db_callback: T.Callable[[str, str], None] | None
 
     def __init__(
         self,
         redis_info: RedisInfo,
         args: argparse.Namespace,
-        backend_id: str,
         verbose: Verbose,
         logging_error_db_callback: T.Callable[[str, str], None] | None = None,
     ):
@@ -88,12 +72,11 @@ class DbUpdater(RedisClientBase):
         self.last_db_init_retry_time = 0.0
         self.use_local_db_only = args.use_local_db_only if args else True
         self.database_settings_msg: T.Optional[PostgresInfo] = None
-        self.backend_id = backend_id
         self.logging_error_db_callback = logging_error_db_callback
 
     @message_handler
-    def handle_database_config_message(self, message_pb: DatabaseConfigPb) -> None:
-        database_settings = get_database_settings(message_pb, backend_id=self.backend_id)
+    def handle_database_config_message(self, message_pb: PostgresMessagePb) -> None:
+        database_settings = get_database_settings(message_pb)
 
         if database_settings is not None:
             self.database_settings_msg = database_settings

{ry_pg_utils-1.0.9.dist-info → ry_pg_utils-1.0.11.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ry-pg-utils
-Version: 1.0.9
+Version: 1.0.11
 Summary: Utility functions for PostgreSQL
 Author: Ross Yeager
 Author-email: ryeager12@email.com

{ry_pg_utils-1.0.9.dist-info → ry_pg_utils-1.0.11.dist-info}/RECORD

@@ -1,22 +1,22 @@
 ry_pg_utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 ry_pg_utils/config.py,sha256=7W-JLNxZDoursPArtoKyOiEqUABGyO1ZEkTdlVBzdVM,1522
-ry_pg_utils/connect.py,sha256=jF-Sj9oySPDfTlOh4_C6c03t4KW5Sef_MLHx_sh52a4,10108
+ry_pg_utils/connect.py,sha256=ECasbbwXCsj1pJwz_er2aKW613X73X30EQyS_TgX5i0,9942
 ry_pg_utils/dynamic_table.py,sha256=VBkfbATRk_pW3EtAFHvMkGG0fLZMebw_LdWkFaDDugw,6866
 ry_pg_utils/notify_trigger.py,sha256=_dsuPkoqjnSQqWrsdZK3FqI6zG5UuygVXHwtERwN99Y,12069
 ry_pg_utils/parse_args.py,sha256=swFYq2-nBggul67C9ibNzsf9goEXDKZkLNSISbJZgOw,751
 ry_pg_utils/postgres_info.py,sha256=mj9er830jvrJXUFuFKx1EmZMjuEpMDcd901kgYb_Cck,1495
 ry_pg_utils/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-ry_pg_utils/updater.py,sha256=nfOR2hpf83buaJc0oyiHf0xfE1pl51CYOz1PuRid198,6103
+ry_pg_utils/updater.py,sha256=qEj1Wc4f-z5ixw_twqvG4vXdlTICSi7fBPZzzbGeafY,5561
 ry_pg_utils/ipc/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-ry_pg_utils/ipc/channels.py,sha256=3w88R_BW5wLW7NX4iNAS3akgXo0JFQufdR7y-l_dNGE,499
+ry_pg_utils/ipc/channels.py,sha256=U8MKblm8I-LfokmkC-WuCJmUjr-1KGQJB5Fsdnc1R8M,424
 ry_pg_utils/pb_types/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-ry_pg_utils/pb_types/database_pb2.py,sha256=yf__HycGrVoU8pzhYdnOhXhfAkgjeGl80ls_RsLemy4,2780
-ry_pg_utils/pb_types/database_pb2.pyi,sha256=svpmlEY_99dA_u3CJqn7CO5D7CH2pAGTOqlhVMZLXC4,5689
+ry_pg_utils/pb_types/database_pb2.py,sha256=j56AaBNQOFHp19O-8egcil3e0x8rujhlnAFuXIyMvx0,1886
+ry_pg_utils/pb_types/database_pb2.pyi,sha256=PtBakRFPx_o9PkpvYgn0doNoaoz6X03eI9TrnpfamkI,3267
 ry_pg_utils/pb_types/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 ry_pg_utils/tools/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 ry_pg_utils/tools/db_query.py,sha256=sQESUB3zxv3h2bKcodfwJHKHGXdSn6BNaBiyZvg6eNY,19557
-ry_pg_utils-1.0.9.dist-info/licenses/LICENSE,sha256=PYnig94SABlde939TWrqvOqQkkfjuHttf8KpWKlPFlA,1068
-ry_pg_utils-1.0.9.dist-info/METADATA,sha256=eHg0wS2vfpifjX9vEntJEPSLJOucwhylSfnMgrdpN-s,18039
-ry_pg_utils-1.0.9.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-ry_pg_utils-1.0.9.dist-info/top_level.txt,sha256=OruzbmsQHYyPnAw8RichQ5GK1Sxj-MQwWfJ93PEHXIM,12
-ry_pg_utils-1.0.9.dist-info/RECORD,,
+ry_pg_utils-1.0.11.dist-info/licenses/LICENSE,sha256=PYnig94SABlde939TWrqvOqQkkfjuHttf8KpWKlPFlA,1068
+ry_pg_utils-1.0.11.dist-info/METADATA,sha256=g4WXkVP-LJbEMceqkOslKguOEoiA8fZaFZV_gRblhWA,18040
+ry_pg_utils-1.0.11.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+ry_pg_utils-1.0.11.dist-info/top_level.txt,sha256=OruzbmsQHYyPnAw8RichQ5GK1Sxj-MQwWfJ93PEHXIM,12
+ry_pg_utils-1.0.11.dist-info/RECORD,,