From c4fdfb9cadd1c1d6476a4151db511aa97b2e01e0 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <6687393+dragomirp@users.noreply.github.com>
Date: Mon, 8 Jul 2024 17:03:04 +0300
Subject: [PATCH] [DPE-4772] Subordinate scale up (#274)

* Subordinate scale up

* Wait for backend rel to mark as ready

* Update ready flag

* Sync lib changes

* Update lib
---
 .../data_platform_libs/v0/data_interfaces.py  | 25 ++++++++++++++++++-
 src/relations/backend_database.py             |  2 ++
 src/relations/pgbouncer_provider.py           | 16 ++++++++++--
 3 files changed, 40 insertions(+), 3 deletions(-)

diff --git a/lib/charms/data_platform_libs/v0/data_interfaces.py b/lib/charms/data_platform_libs/v0/data_interfaces.py
index 59a97226a..a2162aa0b 100644
--- a/lib/charms/data_platform_libs/v0/data_interfaces.py
+++ b/lib/charms/data_platform_libs/v0/data_interfaces.py
@@ -331,7 +331,7 @@ def _on_topic_requested(self, event: TopicRequestedEvent):
 
 # Increment this PATCH version before using `charmcraft publish-lib` or reset
 # to 0 if you are raising the major API version
-LIBPATCH = 37
+LIBPATCH = 38
 
 PYDEPS = ["ops>=2.0.0"]
 
@@ -2606,6 +2606,14 @@ def set_version(self, relation_id: int, version: str) -> None:
         """
         self.update_relation_data(relation_id, {"version": version})
 
+    def set_subordinated(self, relation_id: int) -> None:
+        """Raises the subordinated flag in the application relation databag.
+
+        Args:
+            relation_id: the identifier for a particular relation.
+        """
+        self.update_relation_data(relation_id, {"subordinated": "true"})
+
 
 class DatabaseProviderEventHandlers(EventHandlers):
     """Provider-side of the database relation handlers."""
@@ -2842,6 +2850,21 @@ def _on_relation_created_event(self, event: RelationCreatedEvent) -> None:
 
     def _on_relation_changed_event(self, event: RelationChangedEvent) -> None:
         """Event emitted when the database relation has changed."""
+        is_subordinate = False
+        remote_unit_data = None
+        for key in event.relation.data.keys():
+            if isinstance(key, Unit) and not key.name.startswith(self.charm.app.name):
+                remote_unit_data = event.relation.data[key]
+            elif isinstance(key, Application) and key.name != self.charm.app.name:
+                is_subordinate = event.relation.data[key].get("subordinated") == "true"
+
+        if is_subordinate:
+            if not remote_unit_data:
+                return
+
+            if remote_unit_data.get("state") != "ready":
+                return
+
         # Check which data has changed to emit customs events.
         diff = self._diff(event)
 
diff --git a/src/relations/backend_database.py b/src/relations/backend_database.py
index 5c33c17fd..1c82fff3c 100644
--- a/src/relations/backend_database.py
+++ b/src/relations/backend_database.py
@@ -181,6 +181,7 @@ def _on_database_created(self, event: DatabaseCreatedEvent) -> None:
             self.charm.render_pgb_config(reload_pgbouncer=True)
             self.charm.render_prometheus_service()
             self.charm.update_status()
+            self.charm.client_relation.set_ready()
             return
 
         logger.info("initialising pgbouncer backend relation")
@@ -215,6 +216,7 @@ def _on_database_created(self, event: DatabaseCreatedEvent) -> None:
 
         self.charm.render_pgb_config(reload_pgbouncer=True)
         self.charm.render_prometheus_service()
+        self.charm.client_relation.set_ready()
 
         self.charm.update_status()
 
diff --git a/src/relations/pgbouncer_provider.py b/src/relations/pgbouncer_provider.py
index d87e7b87b..77c79e9ba 100644
--- a/src/relations/pgbouncer_provider.py
+++ b/src/relations/pgbouncer_provider.py
@@ -49,7 +49,11 @@
     PostgreSQLDeleteUserError,
     PostgreSQLGetPostgreSQLVersionError,
 )
-from ops.charm import CharmBase, RelationBrokenEvent, RelationDepartedEvent
+from ops.charm import (
+    CharmBase,
+    RelationBrokenEvent,
+    RelationDepartedEvent,
+)
 from ops.framework import Object
 from ops.model import (
     Application,
@@ -107,6 +111,9 @@ def _on_database_requested(self, event: DatabaseRequestedEvent) -> None:
         Deferrals:
             - If backend relation is not fully initialised
         """
+        rel_id = event.relation.id
+        self.database_provides.set_subordinated(rel_id)
+
         if not self.charm.backend.check_backend():
             event.defer()
             return
@@ -121,7 +128,6 @@ def _on_database_requested(self, event: DatabaseRequestedEvent) -> None:
         # Retrieve the database name and extra user roles using the charm library.
         database = event.database
         extra_user_roles = event.extra_user_roles or ""
-        rel_id = event.relation.id
 
         dbs = self.charm.generate_relation_databases()
         dbs[str(event.relation.id)] = {"name": database, "legacy": False}
@@ -168,6 +174,7 @@ def _on_database_requested(self, event: DatabaseRequestedEvent) -> None:
             return
 
         self.charm.render_pgb_config(reload_pgbouncer=True)
+        self.set_ready()
 
         # Share the credentials and updated connection info with the client application.
         self.database_provides.set_credentials(rel_id, user, password)
@@ -251,6 +258,11 @@ def update_connection_info(self, relation):
         self.charm.unit.status = initial_status
         self.charm.update_status()
 
+    def set_ready(self) -> None:
+        """Marks the unit as ready for all database relations."""
+        for relation in self.model.relations[self.relation_name]:
+            relation.data[self.charm.unit].update({"state": "ready"})
+
     def update_read_only_endpoints(self, event: DatabaseRequestedEvent = None) -> None:
         """Set the read-only endpoint only if there are replicas."""
         if not self.charm.unit.is_leader():