[DPE-1454] Update database ownership (#227)

* Bump libs

* Update database ownership

* Fix delete auth query check

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 = 34
+LIBPATCH = 35
 
 PYDEPS = ["ops>=2.0.0"]
 
@@ -3016,7 +3016,7 @@ class KafkaRequiresEvents(CharmEvents):
 # Kafka Provides and Requires
 
 
-class KafkaProvidesData(ProviderData):
+class KafkaProviderData(ProviderData):
     """Provider-side of the Kafka relation."""
 
     def __init__(self, model: Model, relation_name: str) -> None:
@@ -3059,12 +3059,12 @@ def set_zookeeper_uris(self, relation_id: int, zookeeper_uris: str) -> None:
         self.update_relation_data(relation_id, {"zookeeper-uris": zookeeper_uris})
 
 
-class KafkaProvidesEventHandlers(EventHandlers):
+class KafkaProviderEventHandlers(EventHandlers):
     """Provider-side of the Kafka relation."""
 
     on = KafkaProvidesEvents()  # pyright: ignore [reportAssignmentType]
 
-    def __init__(self, charm: CharmBase, relation_data: KafkaProvidesData) -> None:
+    def __init__(self, charm: CharmBase, relation_data: KafkaProviderData) -> None:
         super().__init__(charm, relation_data)
         # Just to keep lint quiet, can't resolve inheritance. The same happened in super().__init__() above
         self.relation_data = relation_data
@@ -3086,15 +3086,15 @@ def _on_relation_changed_event(self, event: RelationChangedEvent) -> None:
             )
 
 
-class KafkaProvides(KafkaProvidesData, KafkaProvidesEventHandlers):
+class KafkaProvides(KafkaProviderData, KafkaProviderEventHandlers):
     """Provider-side of the Kafka relation."""
 
     def __init__(self, charm: CharmBase, relation_name: str) -> None:
-        KafkaProvidesData.__init__(self, charm.model, relation_name)
-        KafkaProvidesEventHandlers.__init__(self, charm, self)
+        KafkaProviderData.__init__(self, charm.model, relation_name)
+        KafkaProviderEventHandlers.__init__(self, charm, self)
 
 
-class KafkaRequiresData(RequirerData):
+class KafkaRequirerData(RequirerData):
     """Requirer-side of the Kafka relation."""
 
     def __init__(
@@ -3124,12 +3124,12 @@ def topic(self, value):
         self._topic = value
 
 
-class KafkaRequiresEventHandlers(RequirerEventHandlers):
+class KafkaRequirerEventHandlers(RequirerEventHandlers):
     """Requires-side of the Kafka relation."""
 
     on = KafkaRequiresEvents()  # pyright: ignore [reportAssignmentType]
 
-    def __init__(self, charm: CharmBase, relation_data: KafkaRequiresData) -> None:
+    def __init__(self, charm: CharmBase, relation_data: KafkaRequirerData) -> None:
         super().__init__(charm, relation_data)
         # Just to keep lint quiet, can't resolve inheritance. The same happened in super().__init__() above
         self.relation_data = relation_data
@@ -3142,10 +3142,13 @@ def _on_relation_created_event(self, event: RelationCreatedEvent) -> None:
             return
 
         # Sets topic, extra user roles, and "consumer-group-prefix" in the relation
-        relation_data = {
-            f: getattr(self, f.replace("-", "_"), "")
-            for f in ["consumer-group-prefix", "extra-user-roles", "topic"]
-        }
+        relation_data = {"topic": self.relation_data.topic}
+
+        if self.relation_data.extra_user_roles:
+            relation_data["extra-user-roles"] = self.relation_data.extra_user_roles
+
+        if self.relation_data.consumer_group_prefix:
+            relation_data["consumer-group-prefix"] = self.relation_data.consumer_group_prefix
 
         self.relation_data.update_relation_data(event.relation.id, relation_data)
 
@@ -3188,7 +3191,7 @@ def _on_relation_changed_event(self, event: RelationChangedEvent) -> None:
             return
 
 
-class KafkaRequires(KafkaRequiresData, KafkaRequiresEventHandlers):
+class KafkaRequires(KafkaRequirerData, KafkaRequirerEventHandlers):
     """Provider-side of the Kafka relation."""
 
     def __init__(
@@ -3200,7 +3203,7 @@ def __init__(
         consumer_group_prefix: Optional[str] = None,
         additional_secret_fields: Optional[List[str]] = [],
     ) -> None:
-        KafkaRequiresData.__init__(
+        KafkaRequirerData.__init__(
             self,
             charm.model,
             relation_name,
@@ -3209,7 +3212,7 @@ def __init__(
             consumer_group_prefix,
             additional_secret_fields,
         )
-        KafkaRequiresEventHandlers.__init__(self, charm, self)
+        KafkaRequirerEventHandlers.__init__(self, charm, self)
 
 
 # Opensearch related events

lib/charms/postgresql_k8s/v0/postgresql.py

@@ -36,7 +36,7 @@
 
 # Increment this PATCH version before using `charmcraft publish-lib` or reset
 # to 0 if you are raising the major API version
-LIBPATCH = 25
+LIBPATCH = 26
 
 INVALID_EXTRA_USER_ROLE_BLOCKING_MESSAGE = "invalid role(s) for extra user roles"
 
@@ -476,11 +476,11 @@ def set_up_database(self) -> None:
         """Set up postgres database with the right permissions."""
         connection = None
         try:
-            self.create_user(
-                "admin",
-                extra_user_roles="pg_read_all_data,pg_write_all_data",
-            )
             with self._connect_to_database() as connection, connection.cursor() as cursor:
+                cursor.execute("SELECT TRUE FROM pg_roles WHERE rolname='admin';")
+                if cursor.fetchone() is not None:
+                    return
+
                 # Allow access to the postgres database only to the system users.
                 cursor.execute("REVOKE ALL PRIVILEGES ON DATABASE postgres FROM PUBLIC;")
                 cursor.execute("REVOKE CREATE ON SCHEMA public FROM PUBLIC;")
@@ -490,6 +490,10 @@ def set_up_database(self) -> None:
                             sql.Identifier(user)
                         )
                     )
+                self.create_user(
+                    "admin",
+                    extra_user_roles="pg_read_all_data,pg_write_all_data",
+                )
                 cursor.execute("GRANT CONNECT ON DATABASE postgres TO admin;")
         except psycopg2.Error as e:
             logger.error(f"Failed to set up databases: {e}")

src/charm.py

@@ -29,6 +29,7 @@
     BlockedStatus,
     MaintenanceStatus,
     ModelError,
+    Relation,
     WaitingStatus,
 )
 
@@ -803,6 +804,15 @@ def leader_ip(self) -> str:
         """Gets leader ip."""
         return self.peers.leader_ip
 
+    @property
+    def client_relations(self) -> List[Relation]:
+        """Return the list of established client relations."""
+        relations = []
+        for relation_name in ["database", "db", "db-admin"]:
+            for relation in self.model.relations.get(relation_name, []):
+                relations.append(relation)
+        return relations
+
 
 if __name__ == "__main__":
     main(PgBouncerCharm)

src/relations/db.py

@@ -258,7 +258,9 @@ def _on_relation_joined(self, join_event: RelationJoinedEvent):
             logger.info(init_msg)
 
             self.charm.backend.postgres.create_user(user, password, admin=self.admin)
-            self.charm.backend.postgres.create_database(database, user)
+            self.charm.backend.postgres.create_database(
+                database, user, client_relations=self.charm.client_relations
+            )
 
             created_msg = f"database and user for {self.relation_name} relation created"
             self.charm.unit.status = initial_status
@@ -271,6 +273,7 @@ def _on_relation_joined(self, join_event: RelationJoinedEvent):
             return
 
         # set up auth function
+        self.charm.backend.remove_auth_function(dbs=[database])
         self.charm.backend.initialise_auth_function([database])
 
     def _on_relation_changed(self, change_event: RelationChangedEvent):
@@ -417,24 +420,13 @@ def _on_relation_broken(self, broken_event: RelationBrokenEvent):
             broken_event.defer()
             return
 
-        dbs = self.charm.generate_relation_databases()
-        dbs.pop(str(broken_event.relation.id), None)
+        dbs = self.charm.get_relation_databases()
+        database = dbs.pop(str(broken_event.relation.id), {}).get("name")
         self.charm.set_relation_databases(dbs)
         if self.charm.unit.is_leader():
             # check database can be deleted from pgb config, and if so, delete it. Database is kept on
             # postgres application because we don't want to delete all user data with one command.
-            delete_db = True
-            relations = self.model.relations.get("db", []) + self.model.relations.get(
-                "db-admin", []
-            )
-            for relation in relations:
-                if relation.id == broken_event.relation.id:
-                    continue
-                if relation.data[self.charm.unit].get("database") == database:
-                    # There's multiple applications using this database, so don't remove it until
-                    # we can guarantee this is the last one.
-                    delete_db = False
-                    break
+            delete_db = database not in [db.get("name") for db in dbs.values()]
 
             if delete_db:
                 self.charm.backend.remove_auth_function([database])

src/relations/pgbouncer_provider.py

@@ -136,8 +136,11 @@ def _on_database_requested(self, event: DatabaseRequestedEvent) -> None:
                 user, password, extra_user_roles=extra_user_roles
             )
             logger.debug("creating database")
-            self.charm.backend.postgres.create_database(database, user)
+            self.charm.backend.postgres.create_database(
+                database, user, client_relations=self.charm.client_relations
+            )
             # set up auth function
+            self.charm.backend.remove_auth_function(dbs=[database])
             self.charm.backend.initialise_auth_function(dbs=[database])
         except (
             PostgreSQLCreateDatabaseError,
@@ -186,14 +189,17 @@ def _on_relation_broken(self, event: RelationBrokenEvent) -> None:
             self.charm.peers.unit_databag.pop(self._depart_flag(event.relation), None)
             return
 
-        dbs = self.charm.generate_relation_databases()
-        dbs.pop(str(event.relation.id), None)
+        dbs = self.charm.get_relation_databases()
+        database = dbs.pop(str(event.relation.id), {}).get("name")
         self.charm.set_relation_databases(dbs)
 
         # Delete the user.
         try:
             user = f"relation_id_{event.relation.id}"
             self.charm.backend.postgres.delete_user(user)
+            delete_db = database not in [db.get("name") for db in dbs.values()]
+            if database and delete_db:
+                self.charm.backend.remove_auth_function(dbs=[database])
         except PostgreSQLDeleteUserError as e:
             logger.exception(e)
             self.charm.unit.status = BlockedStatus(

tests/integration/relations/pgbouncer_provider/helpers.py

@@ -177,8 +177,8 @@ def check_exposed_connection(credentials, tls):
     connection = psycopg2.connect(connstr)
     connection.autocommit = True
     smoke_query = (
-        # TODO fix ownership of DB objects on rerelation in PG to be able to drop
-        f"CREATE TABLE IF NOT EXISTS {table_name}(data TEXT);"
+        f"DROP TABLE IF EXISTS {table_name};"
+        f"CREATE TABLE {table_name}(data TEXT);"
         f"INSERT INTO {table_name}(data) VALUES('{smoke_val}');"
         f"SELECT data FROM {table_name} WHERE data = '{smoke_val}';"
     )

tests/unit/relations/test_db.py

@@ -2,7 +2,7 @@
 # See LICENSE file for licensing details.
 
 import unittest
-from unittest.mock import Mock, PropertyMock, patch
+from unittest.mock import Mock, PropertyMock, patch, sentinel
 
 from ops.model import Unit
 from ops.testing import Harness
@@ -56,6 +56,11 @@ def test_correct_admin_perms_set_in_constructor(self):
         assert self.charm.legacy_db_admin_relation.relation_name == "db-admin"
         assert self.charm.legacy_db_admin_relation.admin is True
 
+    @patch(
+        "charm.PgBouncerCharm.client_relations",
+        new_callable=PropertyMock,
+        return_value=sentinel.client_rels,
+    )
     @patch("relations.backend_database.BackendDatabaseRequires.check_backend", return_value=True)
     @patch(
         "relations.backend_database.BackendDatabaseRequires.postgres", new_callable=PropertyMock
@@ -64,6 +69,7 @@ def test_correct_admin_perms_set_in_constructor(self):
     @patch("charms.postgresql_k8s.v0.postgresql.PostgreSQL")
     @patch("charms.postgresql_k8s.v0.postgresql.PostgreSQL.create_user")
     @patch("charms.postgresql_k8s.v0.postgresql.PostgreSQL.create_database")
+    @patch("relations.backend_database.BackendDatabaseRequires.remove_auth_function")
     @patch("relations.backend_database.BackendDatabaseRequires.initialise_auth_function")
     @patch("charm.PgBouncerCharm.set_relation_databases")
     @patch("charm.PgBouncerCharm.generate_relation_databases")
@@ -72,12 +78,14 @@ def test_on_relation_joined(
         _gen_rel_dbs,
         _set_rel_dbs,
         _init_auth,
+        _remove_auth,
         _create_database,
         _create_user,
         _postgres,
         _gen_pw,
         _backend_pg,
         _check_backend,
+        _,
     ):
         self.harness.set_leader(True)
 
@@ -106,7 +114,7 @@ def test_on_relation_joined(
         _set_rel_dbs.assert_called_once_with({"1": {"name": "test_db", "legacy": True}})
 
         _create_user.assert_called_with(user, password, admin=True)
-        _create_database.assert_called_with(database, user)
+        _create_database.assert_called_with(database, user, client_relations=sentinel.client_rels)
         _init_auth.assert_called_with([database])
 
         for dbag in [relation_data[self.charm.unit], relation_data[self.charm.app]]:

tests/unit/relations/test_pgbouncer_provider.py

@@ -2,7 +2,7 @@
 # See LICENSE file for licensing details.
 
 import unittest
-from unittest.mock import MagicMock, PropertyMock, patch
+from unittest.mock import MagicMock, PropertyMock, patch, sentinel
 
 from ops.testing import Harness
 
@@ -37,6 +37,11 @@ def setUp(self):
         self.client_rel_id = self.harness.add_relation(CLIENT_RELATION_NAME, "application")
         self.harness.add_relation_unit(self.client_rel_id, "application/0")
 
+    @patch(
+        "charm.PgBouncerCharm.client_relations",
+        new_callable=PropertyMock,
+        return_value=sentinel.client_rels,
+    )
     @patch("relations.backend_database.BackendDatabaseRequires.check_backend", return_value=True)
     @patch(
         "relations.backend_database.BackendDatabaseRequires.postgres", new_callable=PropertyMock
@@ -73,6 +78,7 @@ def test_on_database_requested(
         _pg_databag,
         _pg,
         _check_backend,
+        _,
     ):
         self.harness.set_leader()
         _gen_rel_dbs.return_value = {}
@@ -95,7 +101,9 @@ def test_on_database_requested(
         _pg().create_user.assert_called_with(
             user, _password(), extra_user_roles=event.extra_user_roles
         )
-        _pg().create_database.assert_called_with(database, user)
+        _pg().create_database.assert_called_with(
+            database, user, client_relations=sentinel.client_rels
+        )
         _dbp_set_credentials.assert_called_with(rel_id, user, _password())
         _dbp_set_version.assert_called_with(rel_id, _pg().get_postgresql_version())
         _dbp_set_endpoints.assert_called_with(