[DPE-3927] Multiple databases (#210)

* Wildcard * Use only unit databags for legacy relation * Remove the reboot flag * Landscape test * Wrong import * Peer data magic * Update test * Use transaction pooling * Reload the service * Revert reloading * Try to keep backend connections below 100 * There's no magic :( * Limit backend conns * Limit the amount of PGB instances up to 4 * Don't use postges as authdb * Inject wildcard db together with admin db * Check for createdb for enabling autodb
canonical · Jun 3, 2024 · 8998a09 · 8998a09
1 parent e293c9e
commit 8998a09
Show file tree

Hide file tree

Showing 13 changed files with 349 additions and 131 deletions.
diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -75,6 +75,7 @@ juju = "<=3.5.0.0"
 tenacity = "*"
 mailmanclient = "^3.3.5"
 psycopg2-binary = "^2.9.9"
+landscape-api-py3 = "^0.9.0"
 allure-pytest = "^2.13.5"
 
 [build-system]

diff --git a/src/charm.py b/src/charm.py
@@ -104,7 +104,7 @@ def __init__(self, *args):
         self.legacy_db_admin_relation = DbProvides(self, admin=True)
         self.tls = PostgreSQLTLS(self, PEER_RELATION_NAME)
 
-        self._cores = os.cpu_count()
+        self._cores = max(min(os.cpu_count(), 4), 2)
         self.service_ids = list(range(self._cores))
         self.pgb_services = [
             f"{PGB}-{self.app.name}@{service_id}" for service_id in self.service_ids
@@ -536,6 +536,7 @@ def generate_relation_databases(self) -> Dict[str, Dict[str, Union[str, bool]]]:
             return {}
 
         databases = {}
+        add_wildcard = False
         for relation in self.model.relations.get("db", []):
             database = self.legacy_db_relation.get_databags(relation)[0].get("database")
             if database:
@@ -551,16 +552,22 @@ def generate_relation_databases(self) -> Dict[str, Dict[str, Union[str, bool]]]:
                     "name": database,
                     "legacy": True,
                 }
+                add_wildcard = True
 
         for rel_id, data in self.client_relation.database_provides.fetch_relation_data(
-            fields=["database"]
+            fields=["database", "extra-user-roles"]
         ).items():
             database = data.get("database")
+            roles = data.get("extra-user-roles", "").lower().split(",")
             if database:
                 databases[str(rel_id)] = {
                     "name": database,
                     "legacy": False,
                 }
+            if "admin" in roles or "superuser" in roles or "createdb" in roles:
+                add_wildcard = True
+        if add_wildcard:
+            databases["*"] = {"name": "*", "auth_dbname": database}
         self.set_relation_databases(databases)
         return databases
 
@@ -586,6 +593,8 @@ def _get_relation_config(self) -> [Dict[str, Dict[str, Union[str, bool]]]]:
 
         for database in databases.values():
             name = database["name"]
+            if name == "*":
+                continue
             pgb_dbs[name] = {
                 "host": host,
                 "dbname": name,
@@ -599,6 +608,13 @@ def _get_relation_config(self) -> [Dict[str, Dict[str, Union[str, bool]]]]:
                     "port": r_port,
                     "auth_user": self.backend.auth_user,
                 }
+        if "*" in databases:
+            pgb_dbs["*"] = {
+                "host": host,
+                "port": port,
+                "auth_user": self.backend.auth_user,
+                "auth_dbname": databases["*"]["auth_dbname"],
+            }
         return pgb_dbs
 
     def render_pgb_config(self, reload_pgbouncer=False):
@@ -641,7 +657,9 @@ def render_pgb_config(self, reload_pgbouncer=False):
                     f"{app_conf_dir}/{INSTANCE_DIR}{service_id}/pgbouncer.ini",
                     template.render(
                         databases=databases,
-                        socket_dir=f"{app_temp_dir}/{INSTANCE_DIR}{service_id}",
+                        peer_id=service_id,
+                        base_socket_dir=f"{app_temp_dir}/{INSTANCE_DIR}",
+                        peers=self.service_ids,
                         log_file=f"{app_log_dir}/{INSTANCE_DIR}{service_id}/pgbouncer.log",
                         pid_file=f"{app_temp_dir}/{INSTANCE_DIR}{service_id}/pgbouncer.pid",
                         listen_addr=addr,

diff --git a/src/relations/db.py b/src/relations/db.py
@@ -84,6 +84,7 @@
 ------------------------------------------------------------------------------------------------------------
 """  # noqa: W505
 
+import json
 import logging
 from typing import Dict, Iterable, List
 
@@ -107,7 +108,6 @@
     MaintenanceStatus,
     Relation,
     Unit,
-    WaitingStatus,
 )
 
 from constants import EXTENSIONS_BLOCKING_MESSAGE
@@ -219,12 +219,10 @@ def _on_relation_joined(self, join_event: RelationJoinedEvent):
         if not (database := remote_app_databag.get("database")) and not (
             database := remote_unit_databag.get("database")
         ):
-            # If there's nothing in either databag, return early.
-            no_db = "No database name provided in app or unit databag"
-            logger.warning(no_db)
-            self.charm.unit.status = WaitingStatus(no_db)
-            join_event.defer()
-            return
+            # Sometimes a relation changed event is triggered, and it doesn't have
+            # a database name in it (like the relation with Landscape server charm),
+            # so create a database with the other application name.
+            database = join_event.relation.app.name
 
         if self._block_on_extensions(join_event.relation, remote_app_databag):
             return
@@ -239,16 +237,17 @@ def _on_relation_joined(self, join_event: RelationJoinedEvent):
 
         dbs = self.charm.generate_relation_databases()
         dbs[str(join_event.relation.id)] = {"name": database, "legacy": True}
+        if self.admin:
+            dbs["*"] = {"name": "*", "auth_dbname": database}
         self.charm.set_relation_databases(dbs)
 
-        self.update_databags(
-            join_event.relation,
-            {
-                "user": user,
-                "password": password,
-                "database": database,
-            },
-        )
+        creds = {
+            "user": user,
+            "password": password,
+            "database": database,
+        }
+        self.charm.peers.app_databag[user] = json.dumps(creds)
+        self.update_databags(join_event.relation, creds)
 
         # Create user and database in backend postgresql database
         try:
@@ -298,7 +297,8 @@ def _on_relation_changed(self, change_event: RelationChangedEvent):
         # No backup values because if databag isn't populated, this relation isn't initialised.
         # This means that the database and user requested in this relation haven't been created,
         # so we defer this event until the databag is populated.
-        databag = self.get_databags(change_event.relation)[0]
+        user = self._generate_username(change_event.relation)
+        databag = json.loads(self.charm.peers.app_databag.get(user, "{}"))
         database = databag.get("database")
         user = databag.get("user")
         password = databag.get("password")
@@ -311,21 +311,20 @@ def _on_relation_changed(self, change_event: RelationChangedEvent):
             return
 
         self.charm.render_pgb_config(reload_pgbouncer=True)
-        if self.charm.unit.is_leader():
-            self.update_connection_info(change_event.relation, self.charm.config["listen_port"])
-            self.update_databags(
-                change_event.relation,
-                {
-                    "allowed-subnets": self.get_allowed_subnets(change_event.relation),
-                    "allowed-units": self.get_allowed_units(change_event.relation),
-                    "version": self.charm.backend.postgres.get_postgresql_version(),
-                    "host": "localhost",
-                    "user": user,
-                    "password": password,
-                    "database": database,
-                    "state": self._get_state(),
-                },
-            )
+        self.update_databags(
+            change_event.relation,
+            {
+                "allowed-subnets": self.get_allowed_subnets(change_event.relation),
+                "allowed-units": self.get_allowed_units(change_event.relation),
+                "version": self.charm.backend.postgres.get_postgresql_version(),
+                "host": "localhost",
+                "user": user,
+                "password": password,
+                "database": database,
+                "state": "master",
+            },
+        )
+        self.update_connection_info(change_event.relation, self.charm.config["listen_port"])
 
     def update_connection_info(self, relation: Relation, port: str = None):
         """Updates databag connection information."""
@@ -451,8 +450,7 @@ def update_databags(self, relation, updates: Dict[str, str]):
         for key, item in updates.items():
             updates[key] = str(item)
 
-        for databag in self.get_databags(relation):
-            databag.update(updates)
+        relation.data[self.charm.unit].update(updates)
 
     def _generate_username(self, relation):
         """Generates a unique username for this relation."""
@@ -461,20 +459,6 @@ def _generate_username(self, relation):
         model_name = self.model.name
         return f"{app_name}_user_{relation_id}_{model_name}".replace("-", "_")
 
-    def _get_state(self) -> str:
-        """Gets the given state for this unit.
-
-        Args:
-            standbys: the comma-separated list of postgres standbys
-
-        Returns:
-            The described state of this unit. Can be 'master' or 'standby'.
-        """
-        if self.charm.unit.is_leader():
-            return "master"
-        else:
-            return "standby"
-
     def get_allowed_subnets(self, relation: Relation) -> str:
         """Gets the allowed subnets from this relation."""
 

diff --git a/src/relations/pgbouncer_provider.py b/src/relations/pgbouncer_provider.py
@@ -155,6 +155,9 @@ def _on_database_requested(self, event: DatabaseRequestedEvent) -> None:
 
         dbs = self.charm.generate_relation_databases()
         dbs[str(event.relation.id)] = {"name": database, "legacy": False}
+        roles = extra_user_roles.lower().split(",")
+        if "admin" in roles or "superuser" in roles:
+            dbs["*"] = {"name": "*", "auth_dbname": database}
         self.charm.set_relation_databases(dbs)
 
         # Share the credentials and updated connection info with the client application.

diff --git a/templates/pgb_config.j2 b/templates/pgb_config.j2
@@ -1,9 +1,15 @@
 [databases]
-{% for name, database in databases.items() %}
-{{ name }} = host={{ database.host }} dbname={{ database.dbname }} port={{ database.port }} auth_user={{ database.auth_user }}
+{% for name, database in databases.items() -%}
+{{ name }} = host={{ database.host }} {% if database.dbname %}dbname={{ database.dbname }}{% else %}auth_dbname={{ database.auth_dbname }}{% endif %} port={{ database.port }} auth_user={{ database.auth_user }}
+{% endfor %}
+
+[peers]
+{% for peer in peers -%}
+{{ peer + 1 }} = host={{ base_socket_dir }}{{ peer }} port={{ listen_port }}
 {% endfor %}
 
 [pgbouncer]
+peer_id = {{ peer_id + 1 }}
 listen_addr = {{ listen_addr }}
 listen_port = {{ listen_port }}
 logfile = {{ log_file }}
@@ -15,7 +21,7 @@ max_client_conn = 10000
 ignore_startup_parameters = extra_float_digits,options
 server_tls_sslmode = prefer
 so_reuseport = 1
-unix_socket_dir = {{ socket_dir }}
+unix_socket_dir = {{ base_socket_dir }}{{ peer_id }}
 pool_mode = {{ pool_mode }}
 max_db_connections = {{ max_db_connections }}
 default_pool_size = {{ default_pool_size }}

diff --git a/templates/pgbouncer.service.j2 b/templates/pgbouncer.service.j2
@@ -4,12 +4,10 @@ After=network.target
 
 [Service]
 Type=simple
-# -R flag lets separate pgbouncer instances reuse sockets from previous instances, so on restart
-# they'll reuse the same sockets, preserving connections.
 ExecStartPre=-/usr/bin/install -o snap_daemon -g snap_daemon -m 700 -d \
         /var/snap/charmed-postgresql/common/var/log/pgbouncer/{{ app_name }}/instance_%i/ \
         {{ snap_tmp_dir }}/{{ app_name }}/instance_%i/
-ExecStart=/snap/bin/charmed-postgresql.pgbouncer-server -R {{ conf_dir }}/{{ app_name }}/instance_%i/pgbouncer.ini
+ExecStart=/snap/bin/charmed-postgresql.pgbouncer-server {{ conf_dir }}/{{ app_name }}/instance_%i/pgbouncer.ini
 KillSignal=SIGINT
 ExecReload=kill -HUP $MAINPID
 Restart=always