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

* Sync relations with pg_hba

* Port from K8s

* Increase timeouts

* Switch observer to httpx

* Bump coverage

* Tactical sleep

* Try to clean up triggers

* Use edge for spaces test

* Blocking test app

* Wrong host

* Drop second trigger

Author: dragomirp

lib/charms/postgresql_k8s/v1/postgresql.py

@@ -898,91 +898,6 @@ def set_up_database(self, temp_location: Optional[str] = None) -> None:
                 self.set_up_login_hook_function()
                 self.set_up_predefined_catalog_roles_function()
 
-                # Create database function and event trigger to identify users created by PgBouncer.
-                cursor.execute("""
-CREATE OR REPLACE FUNCTION update_pg_hba()
-    RETURNS event_trigger
-    LANGUAGE plpgsql
-    AS $$
-        DECLARE
-          temp_schema TEXT;
-          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.
-            SELECT nspname INTO temp_schema FROM pg_namespace WHERE  oid = pg_my_temp_schema();
-            IF temp_schema != '' THEN
-                PERFORM TRUE FROM pg_tables WHERE schemaname = temp_schema AND tablename = 'pg_hba';
-                IF FOUND THEN
-                    DROP TABLE pg_hba;
-                END IF;
-                PERFORM TRUE FROM pg_tables WHERE schemaname = temp_schema AND tablename = 'relation_users';
-                IF FOUND THEN
-                    DROP TABLE relation_users;
-                END IF;
-            END IF;
-            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.
-                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;
-    $$ SECURITY DEFINER;
-                """)
-                cursor.execute(
-                    "SELECT TRUE FROM pg_event_trigger WHERE evtname = 'update_pg_hba_on_create_schema';"
-                )
-                if cursor.fetchone() is None:
-                    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();
-                    """)
-
             connection.close()
             connection = None
 
@@ -1403,3 +1318,40 @@ def is_user_in_hba(self, username: str) -> bool:
         finally:
             if connection:
                 connection.close()
+
+    def drop_hba_triggers(self) -> None:
+        """Drop pg_hba triggers on schema change."""
+        try:
+            with self._connect_to_database() as connection, connection.cursor() as cursor:
+                cursor.execute(
+                    SQL(
+                        "SELECT datname FROM pg_database WHERE datname <> 'template0' AND datname <>'postgres';"
+                    )
+                )
+                databases = [row[0] for row in cursor.fetchall()]
+        except psycopg2.Error as e:
+            logger.warning(f"Failed to get databases when removing hba trigger: {e}")
+            return
+        finally:
+            if connection:
+                connection.close()
+
+            # Existing objects need to be reassigned in each database
+            # before the user can be deleted.
+
+        for database in databases:
+            try:
+                with self._connect_to_database(
+                    database
+                ) as connection, connection.cursor() as cursor:
+                    cursor.execute(
+                        SQL("DROP EVENT TRIGGER IF EXISTS update_pg_hba_on_create_schema;")
+                    )
+                    cursor.execute(
+                        SQL("DROP EVENT TRIGGER IF EXISTS update_pg_hba_on_drop_schema;")
+                    )
+            except psycopg2.Error as e:
+                logger.warning(f"Failed to remove hba trigger for {database}: {e}")
+            finally:
+                if connection:
+                    connection.close()

scripts/cluster_topology_observer.py

@@ -3,25 +3,25 @@
 
 """Cluster topology changes observer."""
 
-import json
 import subprocess
 import sys
-from asyncio import as_completed, get_running_loop, run, wait
+from asyncio import as_completed, create_task, run, wait
 from contextlib import suppress
-from os import environ
 from ssl import create_default_context
 from time import sleep
 from urllib.parse import urljoin
-from urllib.request import urlopen
 
+import psycopg2
 import yaml
+from httpx import AsyncClient
 
 API_REQUEST_TIMEOUT = 5
 PATRONI_CLUSTER_STATUS_ENDPOINT = "cluster"
 TLS_CA_BUNDLE_FILE = "peer_ca_bundle.pem"
 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"
@@ -31,40 +31,16 @@ class UnreachableUnitsError(Exception):
     """Cannot reach any known cluster member."""
 
 
-def call_url(url, context):
-    """Task handler for calling an url."""
-    try:
-        # Scheme is generated by the charm
-        resp = urlopen(  # noqa: S310
-            url,
-            timeout=API_REQUEST_TIMEOUT,
-            context=context,
-        )
-        return json.loads(resp.read())
-    except Exception as e:
-        print(f"Failed to contact {url} with {e}")
-
-
-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
+async def _httpx_get_request(url: str):
+    ssl_ctx = create_default_context()
+    with suppress(FileNotFoundError):
+        ssl_ctx.load_verify_locations(cafile=f"{PATRONI_CONF_PATH}/{TLS_CA_BUNDLE_FILE}")
+    async with AsyncClient(timeout=API_REQUEST_TIMEOUT, verify=ssl_ctx) as client:
+        try:
+            return (await client.get(url)).json()
+        except Exception as e:
+            print(f"Failed to contact {url} with {e}")
+            return None
 
 
 def check_for_cluster_topology_changes(
@@ -90,34 +66,29 @@ def check_for_database_changes(run_cmd, unit, charm_dir, previous_databases):
 
     If changes are detected, dispatch an event to handle them.
     """
-    conf_file_path = "/var/snap/charmed-postgresql/current/etc/patroni/patroni.yaml"
-    with open(conf_file_path) as conf_file:
+    with open(PATRONI_CONF_FILE_PATH) as conf_file:
         conf_file_contents = yaml.safe_load(conf_file)
-        username = conf_file_contents["postgresql"]["authentication"]["superuser"]["username"]
         password = conf_file_contents["postgresql"]["authentication"]["superuser"]["password"]
-    env = environ.copy()
-    env["PGPASSWORD"] = password
-    # Fake cert location for patronictl
-    env["PGSSLCERT"] = "/var/snap/charmed-postgresql/current/etc/patroni/nonexistent_cert.pem"
-    command = [
-        "sudo",
-        "-E",
-        "-H",
-        "charmed-postgresql.patronictl",
-        "-c",
-        conf_file_path,
-        "query",
-        "-c",
-        "SELECT datname,datacl FROM pg_database;",
-        "-U",
-        username,
-        "-d",
-        "postgres",
-    ]
+        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
-        current_databases = subprocess.check_output(command, env=env)  # noqa: S603
-    except subprocess.CalledProcessError as e:
+        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
@@ -131,6 +102,9 @@ def check_for_database_changes(run_cmd, unit, charm_dir, 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):
@@ -148,7 +122,6 @@ async 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("/", "-")
@@ -159,8 +132,7 @@ async def main():
             context.load_verify_locations(cafile=f"{PATRONI_CONF_PATH}/{TLS_CA_BUNDLE_FILE}")
 
         cluster_status = None
-        loop = get_running_loop()
-        tasks = [loop.run_in_executor(None, call_url, url, context) for url in urls]
+        tasks = [create_task(_httpx_get_request(url)) for url in urls]
         for task in as_completed(tasks):
             if result := await task:
                 for task in tasks:
@@ -190,9 +162,6 @@ async 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
             )