draft

Author: sinclert-canonical

poetry.lock

@@ -0,0 +1,133 @@
+#!/usr/bin/env python3
+# Copyright 2025 Canonical Ltd.
+# See LICENSE file for licensing details.
+
+from collections.abc import Callable
+
+import jubilant
+from jubilant import Juju
+from jubilant.statustypes import Status, UnitStatus
+
+from constants import SERVER_CONFIG_USERNAME
+
+from ..helpers import execute_queries_on_unit
+
+MINUTE_SECS = 60
+
+JujuModelStatusFunc = Callable[[Status], bool]
+JujuAppsStatusFunc = Callable[[Status, str], bool]
+
+
+def get_app_leader(juju: Juju, app_name: str) -> str:
+    """Get the leader unit for the given application."""
+    model_status = juju.status()
+    app_status = model_status.apps[app_name]
+    for name, status in app_status.units.items():
+        if status.leader:
+            return name
+
+    raise Exception("No leader unit found")
+
+
+def get_app_name(juju: Juju, charm_name: str) -> str | None:
+    """Get the application name for the given charm."""
+    model_status = juju.status()
+    app_statuses = model_status.apps
+    for name, status in app_statuses.items():
+        if status.charm_name == charm_name:
+            return name
+
+    raise Exception("No application name found")
+
+
+def get_app_units(juju: Juju, app_name: str) -> dict[str, UnitStatus]:
+    """Get the units for the given application."""
+    model_status = juju.status()
+    app_status = model_status.apps[app_name]
+    return app_status.units
+
+
+def get_unit_ip(juju: Juju, app_name: str, unit_name: str) -> str:
+    """Get the application unit IP."""
+    model_status = juju.status()
+    app_status = model_status.apps[app_name]
+    for name, status in app_status.units.items():
+        if name == unit_name:
+            return status.public_address
+
+    raise Exception("No application unit found")
+
+
+def get_mysql_cluster_status(juju: Juju, unit: str, cluster_set: bool | None = False) -> dict:
+    """Get the cluster status by running the get-cluster-status action.
+
+    Args:
+        juju: The juju instance to use.
+        unit: The unit on which to execute the action on
+        cluster_set: Whether to get the cluster-set instead
+
+    Returns:
+        A dictionary representing the cluster status
+    """
+    task = juju.run(
+        unit=unit,
+        action="get-cluster-status",
+        params={"cluster-set": bool(cluster_set)},
+        wait=5 * MINUTE_SECS,
+    )
+    task.raise_on_failure()
+
+    return task.results.get("status", {})
+
+
+def get_mysql_primary_unit(juju: Juju, app_name: str) -> str:
+    """Get the current primary node of the cluster."""
+    mysql_primary = get_app_leader(juju, app_name)
+    mysql_cluster_status = get_mysql_cluster_status(juju, mysql_primary)
+    mysql_cluster_topology = mysql_cluster_status["defaultreplicaset"]["topology"]
+
+    for label, value in mysql_cluster_topology.items():
+        if value["memberrole"] == "primary":
+            return label.replace("-", "/")
+
+    raise Exception("No MySQL primary node found")
+
+
+async def get_mysql_max_written_value(juju: Juju, app_name: str, unit_name: str) -> int:
+    """Retrieve the max written value in the MySQL database.
+
+    Args:
+        juju: The Juju model.
+        app_name: The application name.
+        unit_name: The unit name.
+    """
+    credentials_task = juju.run(
+        unit=unit_name,
+        action="get-password",
+        params={"username": SERVER_CONFIG_USERNAME},
+    )
+    credentials_task.raise_on_failure()
+
+    output = await execute_queries_on_unit(
+        get_unit_ip(juju, app_name, unit_name),
+        credentials_task.results["username"],
+        credentials_task.results["password"],
+        ["SELECT MAX(number) FROM `continuous_writes`.`data`;"],
+    )
+    return output[0]
+
+
+def wait_for_status(jubilant_status_func: JujuAppsStatusFunc, *apps: str) -> JujuModelStatusFunc:
+    """Waits for Juju agents to be idle, and for applications to reach a certain status.
+
+    Args:
+        jubilant_status_func: The Juju apps status function to wait for.
+        apps: The applications to wait for.
+
+    Returns:
+        Juju model status function.
+    """
+    return lambda status: all((
+        jubilant.all_agents_idle(status, *apps),
+        jubilant_status_func(status, *apps),
+    ))