authorisation_rules_observer.py

# Copyright 2025 Canonical Ltd.
# See LICENSE file for licensing details.

"""Authorisation rules changes observer."""

import subprocess
import sys
from asyncio import as_completed, create_task, run, wait
from contextlib import suppress
from ssl import create_default_context
from time import sleep
from urllib.parse import urljoin

import psycopg2
import yaml
from httpx import AsyncClient

API_REQUEST_TIMEOUT = 5
PATRONI_CLUSTER_STATUS_ENDPOINT = "cluster"
PATRONI_CONFIG_STATUS_ENDPOINT = "config"
TLS_CA_BUNDLE_FILE = "peer_ca_bundle.pem"
PATRONI_CONF_FILE_PATH = "/var/lib/pg/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."""


async def _httpx_get_request(url: str):
    ssl_ctx = create_default_context()
    with suppress(FileNotFoundError):
        # Certificate bundle is not secret
        ssl_ctx.load_verify_locations(cafile=f"/tmp/{TLS_CA_BUNDLE_FILE}")  # noqa: S108
    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 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(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()


async 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_databases = None
    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()
        with suppress(FileNotFoundError):
            # Certificate bundle is not secret
            context.load_verify_locations(cafile=f"/tmp/{TLS_CA_BUNDLE_FILE}")  # noqa: S108

        cluster_status = None
        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:
                    task.cancel()
                await wait(tasks)
                cluster_status = result
                break
        if not cluster_status:
            raise UnreachableUnitsError("Unable to reach cluster members")
        is_primary = False
        for member in cluster_status["members"]:
            member_url = urljoin(member["api_url"], PATRONI_CLUSTER_STATUS_ENDPOINT)
            # Call the current unit first
            if member["name"] == member_name:
                urls.insert(0, member_url)
                # Check if the current member is the primary.
                if member["role"] == "leader":
                    is_primary = True
            else:
                urls.append(member_url)

        if is_primary:
            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)


if __name__ == "__main__":
    run(main())