Merge branch 'main' into sync16

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 = 52
+LIBPATCH = 53
 
 # Groups to distinguish HBA access
 ACCESS_GROUP_IDENTITY = "identity_access"
@@ -113,6 +113,10 @@ class PostgreSQLGetPostgreSQLVersionError(Exception):
     """Exception raised when retrieving PostgreSQL version fails."""
 
 
+class PostgreSQLListAccessibleDatabasesForUserError(Exception):
+    """Exception raised when retrieving the accessible databases for a user fails."""
+
+
 class PostgreSQLListGroupsError(Exception):
     """Exception raised when retrieving PostgreSQL groups list fails."""
 
@@ -190,6 +194,11 @@ def create_access_groups(self) -> None:
         try:
             with self._connect_to_database() as connection, connection.cursor() as cursor:
                 for group in ACCESS_GROUPS:
+                    cursor.execute(
+                        SQL("SELECT TRUE FROM pg_roles WHERE rolname={};").format(Literal(group))
+                    )
+                    if cursor.fetchone() is not None:
+                        continue
                     cursor.execute(
                         SQL("CREATE ROLE {} NOLOGIN;").format(
                             Identifier(group),
@@ -622,15 +631,22 @@ def is_tls_enabled(self, check_current_host: bool = False) -> bool:
             # Connection errors happen when PostgreSQL has not started yet.
             return False
 
-    def list_access_groups(self) -> Set[str]:
+    def list_access_groups(self, current_host=False) -> Set[str]:
         """Returns the list of PostgreSQL database access groups.
 
+        Args:
+            current_host: whether to check the current host
+                instead of the primary host.
+
         Returns:
             List of PostgreSQL database access groups.
         """
         connection = None
+        host = self.current_host if current_host else None
         try:
-            with self._connect_to_database() as connection, connection.cursor() as cursor:
+            with self._connect_to_database(
+                database_host=host
+            ) as connection, connection.cursor() as cursor:
                 cursor.execute(
                     "SELECT groname FROM pg_catalog.pg_group WHERE groname LIKE '%_access';"
                 )
@@ -643,16 +659,69 @@ def list_access_groups(self) -> Set[str]:
             if connection is not None:
                 connection.close()
 
-    def list_users(self) -> Set[str]:
+    def list_accessible_databases_for_user(self, user: str, current_host=False) -> Set[str]:
+        """Returns the list of accessible databases for a specific user.
+
+        Args:
+            user: the user to check.
+            current_host: whether to check the current host
+                instead of the primary host.
+
+        Returns:
+            List of accessible database (the ones where
+                the user has the CONNECT privilege).
+        """
+        connection = None
+        host = self.current_host if current_host else None
+        try:
+            with self._connect_to_database(
+                database_host=host
+            ) as connection, connection.cursor() as cursor:
+                cursor.execute(
+                    SQL(
+                        "SELECT TRUE FROM pg_catalog.pg_user WHERE usename = {} AND usesuper;"
+                    ).format(Literal(user))
+                )
+                if cursor.fetchone() is not None:
+                    return {"all"}
+                cursor.execute(
+                    SQL(
+                        "SELECT datname FROM pg_catalog.pg_database WHERE has_database_privilege({}, datname, 'CONNECT') AND NOT datistemplate;"
+                    ).format(Literal(user))
+                )
+                databases = cursor.fetchall()
+                return {database[0] for database in databases}
+        except psycopg2.Error as e:
+            logger.error(f"Failed to list accessible databases for user {user}: {e}")
+            raise PostgreSQLListAccessibleDatabasesForUserError() from e
+        finally:
+            if connection is not None:
+                connection.close()
+
+    def list_users(self, group: Optional[str] = None, current_host=False) -> Set[str]:
         """Returns the list of PostgreSQL database users.
 
+        Args:
+            group: optional group to filter the users.
+            current_host: whether to check the current host
+                instead of the primary host.
+
         Returns:
             List of PostgreSQL database users.
         """
         connection = None
+        host = self.current_host if current_host else None
         try:
-            with self._connect_to_database() as connection, connection.cursor() as cursor:
-                cursor.execute("SELECT usename FROM pg_catalog.pg_user;")
+            with self._connect_to_database(
+                database_host=host
+            ) as connection, connection.cursor() as cursor:
+                if group:
+                    query = SQL(
+                        "SELECT usename FROM (SELECT UNNEST(grolist) AS user_id FROM pg_catalog.pg_group WHERE groname = {}) AS g JOIN pg_catalog.pg_user AS u ON g.user_id = u.usesysid;"
+                    ).format(Literal(group))
+                else:
+                    query = "SELECT usename FROM pg_catalog.pg_user;"
+                cursor.execute(query)
                 usernames = cursor.fetchall()
                 return {username[0] for username in usernames}
         except psycopg2.Error as e:
@@ -662,19 +731,27 @@ def list_users(self) -> Set[str]:
             if connection is not None:
                 connection.close()
 
-    def list_users_from_relation(self) -> Set[str]:
+    def list_users_from_relation(self, current_host=False) -> Set[str]:
         """Returns the list of PostgreSQL database users that were created by a relation.
 
+        Args:
+            current_host: whether to check the current host
+                instead of the primary host.
+
         Returns:
             List of PostgreSQL database users.
         """
         connection = None
+        host = self.current_host if current_host else None
         try:
-            with self._connect_to_database() as connection, connection.cursor() as cursor:
+            with self._connect_to_database(
+                database_host=host
+            ) as connection, connection.cursor() as cursor:
                 cursor.execute(
                     "SELECT usename "
                     "FROM pg_catalog.pg_user "
-                    "WHERE usename LIKE 'relation_id_%' OR usename LIKE 'relation-%';"
+                    "WHERE usename LIKE 'relation_id_%' OR usename LIKE 'relation-%' "
+                    "OR usename LIKE 'pgbouncer_auth_relation_id_%' OR usename LIKE '%_user_%_%';"
                 )
                 usernames = cursor.fetchall()
                 return {username[0] for username in usernames}
@@ -705,31 +782,100 @@ def set_up_database(self, temp_location: Optional[str] = None) -> None:
         connection = None
         cursor = None
         try:
-            connection = self._connect_to_database()
-            cursor = connection.cursor()
-
-            if temp_location is not None:
-                cursor.execute("SELECT TRUE FROM pg_tablespace WHERE spcname='temp';")
+            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(f"CREATE TABLESPACE temp LOCATION '{temp_location}';")
-                    cursor.execute("GRANT CREATE ON TABLESPACE temp TO public;")
-
-            cursor.execute("SELECT TRUE FROM pg_roles WHERE rolname='admin';")
-            if cursor.fetchone() is None:
-                # 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;")
-                for user in self.system_users:
-                    cursor.execute(
-                        SQL("GRANT ALL PRIVILEGES ON DATABASE postgres TO {};").format(
-                            Identifier(user)
+                    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_id_%' 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:
+                    # 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;")
+                    for user in self.system_users:
+                        cursor.execute(
+                            SQL("GRANT ALL PRIVILEGES ON DATABASE postgres TO {};").format(
+                                Identifier(user)
+                            )
                         )
+                    self.create_user(
+                        PERMISSIONS_GROUP_ADMIN,
+                        extra_user_roles=["pg_read_all_data", "pg_write_all_data"],
                     )
-                self.create_user(
-                    PERMISSIONS_GROUP_ADMIN,
-                    extra_user_roles=["pg_read_all_data", "pg_write_all_data"],
-                )
-                cursor.execute("GRANT CONNECT ON DATABASE postgres TO admin;")
+                    cursor.execute("GRANT CONNECT ON DATABASE postgres TO admin;")
         except psycopg2.Error as e:
             logger.error(f"Failed to set up databases: {e}")
             raise PostgreSQLDatabasesSetupError() from e

scripts/authorisation_rules_observer.py

@@ -0,0 +1,92 @@
+# Copyright 2025 Canonical Ltd.
+# See LICENSE file for licensing details.
+
+"""Authorisation rules changes observer."""
+
+import json
+import subprocess
+import sys
+from ssl import CERT_NONE, create_default_context
+from time import sleep
+from urllib.parse import urljoin
+from urllib.request import urlopen
+
+API_REQUEST_TIMEOUT = 5
+PATRONI_CLUSTER_STATUS_ENDPOINT = "cluster"
+PATRONI_CONFIG_STATUS_ENDPOINT = "config"
+
+
+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"
+    # Input is generated by the charm
+    subprocess.run([run_cmd, "-u", unit, dispatch_sub_cmd.format(charm_dir)])  # noqa: S603
+
+
+def main():
+    """Main watch and dispatch loop.
+
+    Watch the Patroni API cluster info. When changes are detected, dispatch the change event.
+    """
+    patroni_urls, run_cmd, unit, charm_dir = sys.argv[1:]
+
+    previous_authorisation_rules = []
+    urls = [urljoin(url, PATRONI_CLUSTER_STATUS_ENDPOINT) for url in patroni_urls.split(",")]
+    member_name = unit.replace("/", "-")
+    while True:
+        # Disable TLS chain verification
+        context = create_default_context()
+        context.check_hostname = False
+        context.verify_mode = CERT_NONE
+
+        cluster_status = None
+        for url in urls:
+            try:
+                # Scheme is generated by the charm
+                resp = urlopen(  # noqa: S310
+                    url,
+                    timeout=API_REQUEST_TIMEOUT,
+                    context=context,
+                )
+                cluster_status = json.loads(resp.read())
+                break
+            except Exception as e:
+                print(f"Failed to contact {url} with {e}")
+                continue
+        if not cluster_status:
+            raise UnreachableUnitsError("Unable to reach cluster members")
+        is_primary = False
+        for member in cluster_status["members"]:
+            # Check if the current member is the primary.
+            if member["name"] == member_name and member["role"] == "leader":
+                is_primary = True
+                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)
+
+        # Wait some time before checking again for a authorisation rules change.
+        sleep(30)
+
+
+if __name__ == "__main__":
+    main()