[DPE-7594] Sync up pg_hba changes and remove trigger (#1070)

* Port user hash

* Blocking test app

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 = 53
+LIBPATCH = 55
 
 # Groups to distinguish HBA access
 ACCESS_GROUP_IDENTITY = "identity_access"
@@ -267,7 +267,8 @@ def create_database(
             raise PostgreSQLCreateDatabaseError() from e
 
         # Enable preset extensions
-        self.enable_disable_extensions(dict.fromkeys(plugins, True), database)
+        if plugins:
+            self.enable_disable_extensions(dict.fromkeys(plugins, True), database)
 
     def create_user(
         self,
@@ -782,83 +783,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/cluster_topology_observer.py

@@ -11,8 +11,15 @@
 from urllib.parse import urljoin
 from urllib.request import urlopen
 
+import psycopg2
+import yaml
+
 API_REQUEST_TIMEOUT = 5
 PATRONI_CLUSTER_STATUS_ENDPOINT = "cluster"
+SNAP_CURRENT_PATH = "/var/snap/charmed-postgresql/current"
+SNAP_CONF_PATH = f"{SNAP_CURRENT_PATH}/etc"
+PATRONI_CONF_PATH = f"{SNAP_CONF_PATH}/patroni"
+PATRONI_CONF_FILE_PATH = f"{PATRONI_CONF_PATH}/patroni.yaml"
 
 # File path for the spawned cluster topology observer process to write logs.
 LOG_FILE_PATH = "/var/log/cluster_topology_observer.log"
@@ -22,28 +29,6 @@ class UnreachableUnitsError(Exception):
     """Cannot reach any known cluster member."""
 
 
-def check_for_authorisation_rules_changes(run_cmd, unit, charm_dir, previous_authorisation_rules):
-    """Check for changes in the authorisation rules.
-
-    If changes are detected, dispatch an event to handle them.
-    """
-    # Read contents from the pg_hba.conf file.
-    with open("/var/snap/charmed-postgresql/common/var/lib/postgresql/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, "authorisation_rules_change")
-    return previous_authorisation_rules
-
-
 def check_for_cluster_topology_changes(
     run_cmd, unit, charm_dir, previous_cluster_topology, current_cluster_topology
 ):
@@ -62,6 +47,52 @@ def check_for_cluster_topology_changes(
     return previous_cluster_topology
 
 
+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"]
+        listen_addr = conf_file_contents["postgresql"]["listen"]
+        if not listen_addr or ":" not in listen_addr:
+            with open(LOG_FILE_PATH, "a") as log_file:
+                log_file.write("Failed to retrieve databases: Host not set yet.\n")
+            return previous_databases
+        host, _ = listen_addr.split(":")
+
+    connection = None
+    try:
+        # Input is generated by the charm
+        with (
+            psycopg2.connect(
+                f"dbname='postgres' user='operator' host='{host}' "
+                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 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"
@@ -77,7 +108,7 @@ def main():
     patroni_urls, run_cmd, unit, charm_dir = sys.argv[1:]
 
     previous_cluster_topology = {}
-    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:
@@ -122,8 +153,8 @@ def main():
         )
 
         if is_primary:
-            previous_authorisation_rules = check_for_authorisation_rules_changes(
-                run_cmd, unit, charm_dir, previous_authorisation_rules
+            previous_databases = check_for_database_changes(
+                run_cmd, unit, charm_dir, previous_databases
             )
 
         # Wait some time before checking again for a cluster topology change.

src/charm.py

@@ -14,6 +14,7 @@
 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
@@ -185,10 +186,8 @@ def __init__(self, *args):
         )
         self._observer = ClusterTopologyObserver(self, run_cmd)
         self._rotate_logs = RotateLogs(self)
-        self.framework.observe(
-            self.on.authorisation_rules_change, self._on_authorisation_rules_change
-        )
         self.framework.observe(self.on.cluster_topology_change, self._on_cluster_topology_change)
+        self.framework.observe(self.on.databases_change, self._on_databases_change)
         self.framework.observe(self.on.install, self._on_install)
         self.framework.observe(self.on.leader_elected, self._on_leader_elected)
         self.framework.observe(self.on.config_changed, self._on_config_changed)
@@ -239,8 +238,10 @@ def __init__(self, *args):
         )
         self._tracing_endpoint_config, _ = charm_tracing_config(self._grafana_agent, None)
 
-    def _on_authorisation_rules_change(self, _):
-        """Handle authorisation rules change event."""
+    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}")
@@ -2072,6 +2073,9 @@ def update_config(self, is_creating_backup: bool = False, no_peers: bool = False
         self._restart_metrics_service(postgres_snap)
         self._restart_ldap_sync_service(postgres_snap)
 
+        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:
@@ -2190,11 +2194,34 @@ def relations_user_databases_map(self) -> dict:
                     REPLICATION_USER: "all",
                     REWIND_USER: "all",
                 })
+
+            # 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-{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
         except PostgreSQLListUsersError:
             logger.debug("relations_user_databases_map: Unable to get users")
             return {USER: "all", REPLICATION_USER: "all", REWIND_USER: "all"}
 
+    @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_restart_condition(
         self, new_condition: str, repeat_cause: str | None
     ) -> bool:

src/cluster_topology_observer.py

@@ -7,6 +7,8 @@
 import os
 import signal
 import subprocess
+from pathlib import Path
+from sys import version_info
 
 from ops.charm import CharmBase, CharmEvents
 from ops.framework import EventBase, EventSource, Object
@@ -22,8 +24,8 @@ class ClusterTopologyChangeEvent(EventBase):
     """A custom event for cluster topology changes."""
 
 
-class AuthorisationRulesChangeEvent(EventBase):
-    """A custom event for authorisation rules changes."""
+class DatabasesChangeEvent(EventBase):
+    """A custom event for databases changes."""
 
 
 class ClusterTopologyChangeCharmEvents(CharmEvents):
@@ -32,8 +34,8 @@ class ClusterTopologyChangeCharmEvents(CharmEvents):
     Includes :class:`ClusterTopologyChangeEvent` in those that can be handled.
     """
 
-    authorisation_rules_change = EventSource(AuthorisationRulesChangeEvent)
     cluster_topology_change = EventSource(ClusterTopologyChangeEvent)
+    databases_change = EventSource(DatabasesChangeEvent)
 
 
 class ClusterTopologyObserver(Object):
@@ -74,6 +76,20 @@ def start_observer(self):
         new_env = os.environ.copy()
         if "JUJU_CONTEXT_ID" in new_env:
             new_env.pop("JUJU_CONTEXT_ID")
+        # 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] + [
             self._charm._patroni._patroni_url.replace(self._charm._patroni.unit_ip, peer)