Backport #959 to 14/edge (#963)

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 = 56
 
 # Groups to distinguish HBA access
 ACCESS_GROUP_IDENTITY = "identity_access"
@@ -1078,3 +1078,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,
@@ -22,6 +23,7 @@
 from ops.charm import CharmBase, RelationBrokenEvent, RelationChangedEvent
 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 (
     ALL_CLIENT_RELATIONS,
@@ -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]:
@@ -145,27 +144,16 @@ def _on_database_requested(self, event: DatabaseRequestedEvent) -> 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
-
-        user = f"relation-{event.relation.id}"
+        # Try to wait for pg_hba trigger
         try:
-            if user not in self.charm.postgresql.list_users():
-                logger.debug("Deferring on_relation_changed: user was not created yet")
-                event.defer()
-                return
-        except PostgreSQLListUsersError:
-            logger.debug("Deferring on_relation_changed: failed to list users")
-            event.defer()
-            return
-
-        self.charm.update_config()
+            for attempt in Retrying(stop=stop_after_attempt(3), wait=wait_fixed(1)):
+                with attempt:
+                    self.charm.postgresql.is_user_in_hba(user)
+            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_broken(self, event: RelationBrokenEvent) -> None:
         """Correctly update the status."""

tests/integration/ha_tests/test_async_replication.py

@@ -125,8 +125,12 @@ async def test_deploy_async_replication_setup(
             num_units=CLUSTER_SIZE,
             config={"profile": "testing"},
         )
-    await ops_test.model.deploy(APPLICATION_NAME, channel="latest/edge", num_units=1)
-    await second_model.deploy(APPLICATION_NAME, channel="latest/edge", num_units=1)
+    await ops_test.model.deploy(
+        APPLICATION_NAME, channel="latest/edge", num_units=1, config={"sleep_interval": 1000}
+    )
+    await second_model.deploy(
+        APPLICATION_NAME, channel="latest/edge", num_units=1, config={"sleep_interval": 1000}
+    )
 
     async with ops_test.fast_forward(), fast_forward(second_model):
         await gather(
@@ -535,36 +539,32 @@ async def test_scaling(
     logger.info("checking whether writes are increasing")
     await are_writes_increasing(ops_test)
 
-    async with ops_test.fast_forward(FAST_INTERVAL), fast_forward(second_model, FAST_INTERVAL):
-        logger.info("scaling out the first cluster")
-        first_cluster_original_size = len(first_model.applications[DATABASE_APP_NAME].units)
-        await scale_application(ops_test, DATABASE_APP_NAME, first_cluster_original_size + 1)
-
-        logger.info("checking whether writes are increasing")
-        await are_writes_increasing(ops_test, extra_model=second_model)
-
-        logger.info("scaling out the second cluster")
-        second_cluster_original_size = len(second_model.applications[DATABASE_APP_NAME].units)
-        await scale_application(
-            ops_test, DATABASE_APP_NAME, second_cluster_original_size + 1, model=second_model
-        )
-
-        logger.info("checking whether writes are increasing")
-        await are_writes_increasing(ops_test, extra_model=second_model)
-
-        logger.info("scaling in the first cluster")
-        await scale_application(ops_test, DATABASE_APP_NAME, first_cluster_original_size)
+    logger.info("scaling out the clusters")
+    first_cluster_original_size = len(first_model.applications[DATABASE_APP_NAME].units)
+    second_cluster_original_size = len(second_model.applications[DATABASE_APP_NAME].units)
+    await gather(
+        scale_application(ops_test, DATABASE_APP_NAME, first_cluster_original_size + 1),
+        scale_application(
+            ops_test,
+            DATABASE_APP_NAME,
+            second_cluster_original_size + 1,
+            model=second_model,
+        ),
+    )
 
-        logger.info("checking whether writes are increasing")
-        await are_writes_increasing(ops_test, extra_model=second_model)
+    logger.info("checking whether writes are increasing")
+    await are_writes_increasing(ops_test, extra_model=second_model)
 
-        logger.info("scaling in the second cluster")
-        await scale_application(
+    logger.info("scaling in the clusters")
+    await gather(
+        scale_application(ops_test, DATABASE_APP_NAME, first_cluster_original_size),
+        scale_application(
             ops_test, DATABASE_APP_NAME, second_cluster_original_size, model=second_model
-        )
+        ),
+    )
 
-        logger.info("checking whether writes are increasing")
-        await are_writes_increasing(ops_test, extra_model=second_model)
+    logger.info("checking whether writes are increasing")
+    await are_writes_increasing(ops_test, extra_model=second_model)
 
     # Verify that no writes to the database were missed after stopping the writes
     # (check that all the units have all the writes).

tests/integration/helpers.py

@@ -950,7 +950,12 @@ async def run_command_on_unit(ops_test: OpsTest, unit_name: str, command: str) -
 
 
 async def scale_application(
-    ops_test: OpsTest, application_name: str, count: int, model: Model = None
+    ops_test: OpsTest,
+    application_name: str,
+    count: int,
+    model: Model = None,
+    timeout=2000,
+    idle_period: int = 30,
 ) -> None:
     """Scale a given application to a specific unit count.
 
@@ -959,6 +964,8 @@ async def scale_application(
         application_name: The name of the application
         count: The desired number of units to scale to
         model: The model to scale the application in
+        timeout: timeout period
+        idle_period: idle period
     """
     if model is None:
         model = ops_test.model
@@ -971,8 +978,8 @@ async def scale_application(
     await model.wait_for_idle(
         apps=[application_name],
         status="active",
-        timeout=2000,
-        idle_period=30,
+        timeout=timeout,
+        idle_period=idle_period,
         wait_for_exact_units=count,
     )
 

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("subprocess.check_output", return_value=b"C"),
         patch("charm.PostgreSQLProvider.update_endpoints") as _update_endpoints,