[MISC] Hold database created hook for pg_hba changes (#1001)

* Hold database created hook for pg_hba changes

* Restore the old blocking hook

* 16/edge lib changes

* Set to expected patch version

Author: dragomirp

lib/charms/postgresql_k8s/v0/postgresql.py

@@ -35,7 +35,7 @@
 
 # Increment this PATCH version before using `charmcraft publish-lib` or reset
 # to 0 if you are raising the major API version
-LIBPATCH = 55
+LIBPATCH = 53
 
 # Groups to distinguish HBA access
 ACCESS_GROUP_IDENTITY = "identity_access"
@@ -780,6 +780,7 @@ def list_valid_privileges_and_roles(self) -> Tuple[Set[str], Set[str]]:
     def set_up_database(self) -> None:
         """Set up postgres database with the right permissions."""
         connection = None
+        cursor = None
         try:
             with self._connect_to_database(
                 database="template1"
@@ -879,6 +880,8 @@ def set_up_database(self) -> None:
             logger.error(f"Failed to set up databases: {e}")
             raise PostgreSQLDatabasesSetupError() from e
         finally:
+            if cursor is not None:
+                cursor.close()
             if connection is not None:
                 connection.close()
 
@@ -1078,3 +1081,21 @@ def validate_group_map(self, group_map: Optional[str]) -> bool:
                     return False
 
         return True
+
+    def is_user_in_hba(self, username: str) -> bool:
+        """Check if user was added in pg_hba."""
+        connection = None
+        try:
+            with self._connect_to_database() as connection, connection.cursor() as cursor:
+                cursor.execute(
+                    SQL(
+                        "SELECT COUNT(*) FROM pg_hba_file_rules WHERE {} = ANY(user_name);"
+                    ).format(Literal(username))
+                )
+                return cursor.fetchone()[0] > 0
+        except psycopg2.Error as e:
+            logger.debug(f"Failed to check pg_hba: {e}")
+            return False
+        finally:
+            if connection:
+                connection.close()

src/relations/postgresql_provider.py

@@ -4,6 +4,7 @@
 """Postgres client relation hooks & helpers."""
 
 import logging
+from datetime import datetime
 
 from charms.data_platform_libs.v0.data_interfaces import (
     DatabaseProvides,
@@ -21,6 +22,7 @@
 from ops.charm import CharmBase, RelationBrokenEvent, RelationChangedEvent, RelationDepartedEvent
 from ops.framework import Object
 from ops.model import ActiveStatus, BlockedStatus, Relation
+from tenacity import RetryError, Retrying, stop_after_attempt, wait_fixed
 
 from constants import (
     DATABASE_PORT,
@@ -66,9 +68,6 @@ def __init__(self, charm: CharmBase, relation_name: str = "database") -> None:
         self.framework.observe(
             self.database_provides.on.database_requested, self._on_database_requested
         )
-        self.framework.observe(
-            charm.on[self.relation_name].relation_changed, self._on_relation_changed
-        )
 
     @staticmethod
     def _sanitize_extra_roles(extra_roles: str | None) -> list[str]:
@@ -163,26 +162,19 @@ def _on_database_requested(self, event: DatabaseRequestedEvent) -> None:
                 if issubclass(type(e), PostgreSQLCreateUserError) and e.message is not None
                 else f"Failed to initialize {self.relation_name} relation"
             )
-
-    def _on_relation_changed(self, event: RelationChangedEvent) -> None:
-        # Check for some conditions before trying to access the PostgreSQL instance.
-        if not self.charm.is_cluster_initialised:
-            logger.debug(
-                "Deferring on_relation_changed: Cluster must be initialized before configuration can be updated with relation users"
-            )
-            event.defer()
             return
 
-        if (
-            not self.charm._patroni.member_started
-            or f"relation_id_{event.relation.id}"
-            not in self.charm.postgresql.list_users(current_host=True)
-        ):
-            logger.debug("Deferring on_relation_changed: user was not created yet")
-            event.defer()
-            return
-
-        self.charm.update_config()
+        # Try to wait for pg_hba trigger
+        try:
+            for attempt in Retrying(stop=stop_after_attempt(3), wait=wait_fixed(1)):
+                with attempt:
+                    if not self.charm.postgresql.is_user_in_hba(user):
+                        raise Exception("pg_hba not ready")
+            self.charm.unit_peer_data.update({
+                "pg_hba_needs_update_timestamp": str(datetime.now())
+            })
+        except RetryError:
+            logger.warning("database requested: Unable to check pg_hba rule update")
 
     def _on_relation_departed(self, event: RelationDepartedEvent) -> None:
         """Set a flag to avoid deleting database users when not wanted."""

tests/unit/test_postgresql_provider.py

@@ -76,9 +76,6 @@ def request_database(_harness):
 def test_on_database_requested(harness):
     with (
         patch("charm.PostgresqlOperatorCharm.update_config"),
-        patch(
-            "relations.postgresql_provider.PostgreSQLProvider._on_relation_changed"
-        ) as _on_relation_changed,
         patch.object(PostgresqlOperatorCharm, "postgresql", Mock()) as postgresql_mock,
         patch.object(EventBase, "defer") as _defer,
         patch("charm.Patroni.member_started", new_callable=PropertyMock) as _member_started,