[DPE-7594] Sync up pg_hba changes and remove trigger

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 = 53
+LIBPATCH = 54
 
 # Groups to distinguish HBA access
 ACCESS_GROUP_IDENTITY = "identity_access"
@@ -782,83 +782,6 @@ def set_up_database(self) -> None:
         connection = None
         cursor = None
         try:
-            with self._connect_to_database(
-                database="template1"
-            ) as connection, connection.cursor() as cursor:
-                # Create database function and event trigger to identify users created by PgBouncer.
-                cursor.execute(
-                    "SELECT TRUE FROM pg_event_trigger WHERE evtname = 'update_pg_hba_on_create_schema';"
-                )
-                if cursor.fetchone() is None:
-                    cursor.execute("""
-CREATE OR REPLACE FUNCTION update_pg_hba()
-    RETURNS event_trigger
-    LANGUAGE plpgsql
-    AS $$
-        DECLARE
-          hba_file TEXT;
-          copy_command TEXT;
-          connection_type TEXT;
-          rec record;
-          insert_value TEXT;
-          changes INTEGER = 0;
-        BEGIN
-          -- Don't execute on replicas.
-          IF NOT pg_is_in_recovery() THEN
-            -- Load the current authorisation rules.
-            DROP TABLE IF EXISTS pg_hba;
-            CREATE TEMPORARY TABLE pg_hba (lines TEXT);
-            SELECT setting INTO hba_file FROM pg_settings WHERE name = 'hba_file';
-            IF hba_file IS NOT NULL THEN
-                copy_command='COPY pg_hba FROM ''' || hba_file || '''' ;
-                EXECUTE copy_command;
-                -- Build a list of the relation users and the databases they can access.
-                DROP TABLE IF EXISTS relation_users;
-                CREATE TEMPORARY TABLE relation_users AS
-                  SELECT t.user, STRING_AGG(DISTINCT t.database, ',') AS databases FROM( SELECT u.usename AS user, CASE WHEN u.usesuper THEN 'all' ELSE d.datname END AS database FROM ( SELECT usename, usesuper FROM pg_catalog.pg_user WHERE usename NOT IN ('backup', 'monitoring', 'operator', 'postgres', 'replication', 'rewind')) AS u JOIN ( SELECT datname FROM pg_catalog.pg_database WHERE NOT datistemplate ) AS d ON has_database_privilege(u.usename, d.datname, 'CONNECT') ) AS t GROUP BY 1;
-                IF (SELECT COUNT(lines) FROM pg_hba WHERE lines LIKE 'hostssl %') > 0 THEN
-                  connection_type := 'hostssl';
-                ELSE
-                  connection_type := 'host';
-                END IF;
-                -- Add the new users to the pg_hba file.
-                FOR rec IN SELECT * FROM relation_users
-                LOOP
-                  insert_value := connection_type || ' ' || rec.databases || ' ' || rec.user || ' 0.0.0.0/0 md5';
-                  IF (SELECT COUNT(lines) FROM pg_hba WHERE lines = insert_value) = 0 THEN
-                    INSERT INTO pg_hba (lines) VALUES (insert_value);
-                    changes := changes + 1;
-                  END IF;
-                END LOOP;
-                -- Remove users that don't exist anymore from the pg_hba file.
-                FOR rec IN SELECT h.lines FROM pg_hba AS h LEFT JOIN relation_users AS r ON SPLIT_PART(h.lines, ' ', 3) = r.user WHERE r.user IS NULL AND (SPLIT_PART(h.lines, ' ', 3) LIKE 'relation_id_%' OR SPLIT_PART(h.lines, ' ', 3) LIKE 'pgbouncer_auth_relation_%' OR SPLIT_PART(h.lines, ' ', 3) LIKE '%_user_%_%')
-                LOOP
-                  DELETE FROM pg_hba WHERE lines = rec.lines;
-                  changes := changes + 1;
-                END LOOP;
-                -- Apply the changes to the pg_hba file.
-                IF changes > 0 THEN
-                  copy_command='COPY pg_hba TO ''' || hba_file || '''' ;
-                  EXECUTE copy_command;
-                  PERFORM pg_reload_conf();
-                END IF;
-            END IF;
-          END IF;
-        END;
-    $$;
-                    """)
-                    cursor.execute("""
-CREATE EVENT TRIGGER update_pg_hba_on_create_schema
-    ON ddl_command_end
-    WHEN TAG IN ('CREATE SCHEMA')
-    EXECUTE FUNCTION update_pg_hba();
-                    """)
-                    cursor.execute("""
-CREATE EVENT TRIGGER update_pg_hba_on_drop_schema
-    ON ddl_command_end
-    WHEN TAG IN ('DROP SCHEMA')
-    EXECUTE FUNCTION update_pg_hba();
-                    """)
             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 None:

scripts/authorisation_rules_observer.py

@@ -11,20 +11,66 @@
 from urllib.parse import urljoin
 from urllib.request import urlopen
 
+import psycopg2
+import yaml
+
 API_REQUEST_TIMEOUT = 5
 PATRONI_CLUSTER_STATUS_ENDPOINT = "cluster"
 PATRONI_CONFIG_STATUS_ENDPOINT = "config"
+PATRONI_CONF_FILE_PATH = "/var/lib/postgresql/data/patroni.yml"
+
+# File path for the spawned cluster topology observer process to write logs.
+LOG_FILE_PATH = "/var/log/authorisation_rules_observer.log"
 
 
 class UnreachableUnitsError(Exception):
     """Cannot reach any known cluster member."""
 
 
-def dispatch(run_cmd, unit, charm_dir):
-    """Use the input juju-run command to dispatch a :class:`AuthorisationRulesChangeEvent`."""
-    dispatch_sub_cmd = "JUJU_DISPATCH_PATH=hooks/authorisation_rules_change {}/dispatch"
+def dispatch(run_cmd, unit, charm_dir, custom_event):
+    """Use the input juju-run command to dispatch a custom event."""
+    dispatch_sub_cmd = "JUJU_DISPATCH_PATH=hooks/{} {}/dispatch"
     # Input is generated by the charm
-    subprocess.run([run_cmd, "-u", unit, dispatch_sub_cmd.format(charm_dir)])  # noqa: S603
+    subprocess.run([run_cmd, "-u", unit, dispatch_sub_cmd.format(custom_event, charm_dir)])  # noqa: S603
+
+
+def check_for_database_changes(run_cmd, unit, charm_dir, previous_databases):
+    """Check for changes in the databases.
+
+    If changes are detected, dispatch an event to handle them.
+    """
+    with open(PATRONI_CONF_FILE_PATH) as conf_file:
+        conf_file_contents = yaml.safe_load(conf_file)
+        password = conf_file_contents["postgresql"]["authentication"]["superuser"]["password"]
+    connection = None
+    try:
+        # Input is generated by the charm
+        with (
+            psycopg2.connect(
+                "dbname='postgres' user='operator' host='localhost' "
+                f"password='{password}' connect_timeout=1"
+            ) as connection,
+            connection.cursor() as cursor,
+        ):
+            cursor.execute("SELECT datname, datacl FROM pg_database;")
+            current_databases = cursor.fetchall()
+    except psycopg2.Error as e:
+        with open(LOG_FILE_PATH, "a") as log_file:
+            log_file.write(f"Failed to retrieve databases: {e}\n")
+        return previous_databases
+    else:
+        # If it's the first time the databases were retrieved, then store it and use
+        # it for subsequent checks.
+        if not previous_databases:
+            previous_databases = current_databases
+        # If the databases changed, dispatch a charm event to handle this change.
+        elif current_databases != previous_databases:
+            previous_databases = current_databases
+            dispatch(run_cmd, unit, charm_dir, "databases_change")
+        return previous_databases
+    finally:
+        if connection:
+            connection.close()
 
 
 def main():
@@ -34,7 +80,7 @@ def main():
     """
     patroni_urls, run_cmd, unit, charm_dir = sys.argv[1:]
 
-    previous_authorisation_rules = []
+    previous_databases = None
     urls = [urljoin(url, PATRONI_CLUSTER_STATUS_ENDPOINT) for url in patroni_urls.split(",")]
     member_name = unit.replace("/", "-")
     while True:
@@ -67,22 +113,9 @@ def main():
                 break
 
         if is_primary:
-            # Read contents from the pg_hba.conf file.
-            with open("/var/lib/postgresql/data/pgdata/pg_hba.conf") as file:
-                current_authorisation_rules = file.read()
-                current_authorisation_rules = [
-                    line
-                    for line in current_authorisation_rules.splitlines()
-                    if not line.startswith("#")
-                ]
-            # If it's the first time the authorisation rules were retrieved, then store it and use
-            # it for subsequent checks.
-            if not previous_authorisation_rules:
-                previous_authorisation_rules = current_authorisation_rules
-            # If the authorisation rules changed, dispatch a charm event to handle this change.
-            elif current_authorisation_rules != previous_authorisation_rules:
-                previous_authorisation_rules = current_authorisation_rules
-                dispatch(run_cmd, unit, charm_dir)
+            previous_databases = check_for_database_changes(
+                run_cmd, unit, charm_dir, previous_databases
+            )
 
         # Wait some time before checking again for a authorisation rules change.
         sleep(30)

src/authorisation_rules_observer.py

@@ -8,6 +8,8 @@
 import signal
 import subprocess
 import typing
+from pathlib import Path
+from sys import version_info
 
 from ops.charm import CharmEvents
 from ops.framework import EventBase, EventSource, Object
@@ -22,17 +24,17 @@
 LOG_FILE_PATH = "/var/log/authorisation_rules_observer.log"
 
 
-class AuthorisationRulesChangeEvent(EventBase):
-    """A custom event for authorisation rules changes."""
+class DatabasesChangeEvent(EventBase):
+    """A custom event for databases changes."""
 
 
 class AuthorisationRulesChangeCharmEvents(CharmEvents):
     """A CharmEvents extension for authorisation rules changes.
 
-    Includes :class:`AuthorisationRulesChangeEvent` in those that can be handled.
+    Includes :class:`DatabasesChangeEventt` in those that can be handled.
     """
 
-    authorisation_rules_change = EventSource(AuthorisationRulesChangeEvent)
+    databases_change = EventSource(DatabasesChangeEvent)
 
 
 class AuthorisationRulesObserver(Object):
@@ -74,6 +76,20 @@ def start_authorisation_rules_observer(self):
         # in a hook context, as Juju will disallow use of juju-run.
         new_env = os.environ.copy()
         new_env.pop("JUJU_CONTEXT_ID", None)
+        # Generate the venv path based on the existing lib path
+        for loc in new_env["PYTHONPATH"].split(":"):
+            path = Path(loc)
+            venv_path = (
+                path
+                / ".."
+                / "venv"
+                / "lib"
+                / f"python{version_info.major}.{version_info.minor}"
+                / "site-packages"
+            )
+            if path.stem == "lib":
+                new_env["PYTHONPATH"] = f"{venv_path.resolve()}:{new_env['PYTHONPATH']}"
+                break
 
         urls = [
             self._charm._patroni._patroni_url.replace(self._charm.endpoint, endpoint)

src/charm.py

@@ -4,7 +4,6 @@
 
 """Charmed Kubernetes Operator for the PostgreSQL database."""
 
-import datetime
 import itertools
 import json
 import logging
@@ -13,6 +12,8 @@
 import shutil
 import sys
 import time
+from datetime import datetime
+from hashlib import shake_128
 from pathlib import Path
 from typing import Literal, get_args
 from urllib.parse import urlparse
@@ -221,9 +222,7 @@ def __init__(self, *args):
             "/usr/bin/juju-exec" if self.model.juju_version.major > 2 else "/usr/bin/juju-run"
         )
         self._observer = AuthorisationRulesObserver(self, run_cmd)
-        self.framework.observe(
-            self.on.authorisation_rules_change, self._on_authorisation_rules_change
-        )
+        self.framework.observe(self.on.databases_change, self._on_databases_change)
         self.framework.observe(self.on.config_changed, self._on_config_changed)
         self.framework.observe(self.on.leader_elected, self._on_leader_elected)
         self.framework.observe(self.on[PEER].relation_changed, self._on_peer_relation_changed)
@@ -281,9 +280,11 @@ def __init__(self, *args):
             self, relation_name=TRACING_RELATION_NAME, protocols=[TRACING_PROTOCOL]
         )
 
-    def _on_authorisation_rules_change(self, _):
-        """Handle authorisation rules change event."""
-        timestamp = datetime.datetime.now()
+    def _on_databases_change(self, _):
+        """Handle databases change event."""
+        self.update_config()
+        logger.debug("databases changed")
+        timestamp = datetime.now()
         self._peers.data[self.unit].update({"pg_hba_needs_update_timestamp": str(timestamp)})
         logger.debug(f"authorisation rules changed at {timestamp}")
 
@@ -580,14 +581,14 @@ def _on_peer_relation_changed(self, event: HookEvent) -> None:  # noqa: C901
             if self.unit.is_leader():
                 if self._initialize_cluster(event):
                     logger.debug("Deferring on_peer_relation_changed: Leader initialized cluster")
+                    event.defer()
                 else:
                     logger.debug("_initialized_cluster failed on _peer_relation_changed")
                     return
             else:
                 logger.debug(
-                    "Deferring on_peer_relation_changed: Cluster must be initialized before members can join"
+                    "Early exit on_peer_relation_changed: Cluster must be initialized before members can join"
                 )
-            event.defer()
             return
 
         # If the leader is the one receiving the event, it adds the new members,
@@ -2119,6 +2120,9 @@ def update_config(self, is_creating_backup: bool = False) -> bool:
         self._restart_metrics_service()
         self._restart_ldap_sync_service()
 
+        self.unit_peer_data.update({"user_hash": self.generate_user_hash})
+        if self.unit.is_leader():
+            self.app_peer_data.update({"user_hash": self.generate_user_hash})
         return True
 
     def _validate_config_options(self) -> None:
@@ -2316,8 +2320,30 @@ def relations_user_databases_map(self) -> dict:
                     user, current_host=self.is_connectivity_enabled
                 )
             )
+
+        # Copy relations users directly instead of waiting for them to be created
+        for relation in self.model.relations[self.postgresql_client_relation.relation_name]:
+            user = f"relation_id_{relation.id}"
+            if user not in user_database_map and (
+                database := self.postgresql_client_relation.database_provides.fetch_relation_field(
+                    relation.id, "database"
+                )
+            ):
+                user_database_map[user] = database
         return user_database_map
 
+    @property
+    def generate_user_hash(self) -> str:
+        """Generate expected user and database hash."""
+        user_db_pairs = {}
+        for relation in self.model.relations[self.postgresql_client_relation.relation_name]:
+            if database := self.postgresql_client_relation.database_provides.fetch_relation_field(
+                relation.id, "database"
+            ):
+                user = f"relation_id_{relation.id}"
+                user_db_pairs[user] = database
+        return shake_128(str(user_db_pairs).encode()).hexdigest(16)
+
     def override_patroni_on_failure_condition(
         self, new_condition: str, repeat_cause: str | None
     ) -> bool: