[DPE-2290] Upgrade integration tests

.github/workflows/ci.yaml

@@ -83,6 +83,7 @@ jobs:
           - password-rotation-integration
           - plugins-integration
           - tls-integration
+          - upgrade-integration
         agent-versions:
           - "2.9.44"  # renovate: latest juju 2
           - "3.1.5"  # renovate: latest juju 3

src/charm.py

@@ -141,6 +141,14 @@ def __init__(self, *args):
             log_slots=[f"{POSTGRESQL_SNAP_NAME}:logs"],
         )
 
+    @property
+    def app_units(self) -> set[Unit]:
+        """The peer-related units in the application."""
+        if not self._peers:
+            return set()
+
+        return {self.unit, *self._peers.units}
+
     @property
     def app_peer_data(self) -> Dict:
         """Application peer relation data object."""
@@ -935,6 +943,12 @@ def _can_start(self, event: StartEvent) -> bool:
             self._reboot_on_detached_storage(event)
             return False
 
+        # Safeguard against starting while upgrading.
+        if not self.upgrade.idle:
+            logger.debug("Defer on_start: Cluster is upgrading")
+            event.defer()
+            return False
+
         # Doesn't try to bootstrap the cluster if it's in a blocked state
         # caused, for example, because a failed installation of packages.
         if self.is_blocked:
@@ -981,6 +995,10 @@ def _setup_exporter(self) -> None:
         cache = snap.SnapCache()
         postgres_snap = cache[POSTGRESQL_SNAP_NAME]
 
+        if postgres_snap.revision != list(filter(lambda snap_package: snap_package[0] == POSTGRESQL_SNAP_NAME, SNAP_PACKAGES))[0][1]["revision"]:
+            logger.debug("Early exit _setup_exporter: snap was not refreshed to the right version yet")
+            return
+
         postgres_snap.set(
             {
                 "exporter.user": MONITORING_USER,

src/upgrade.py

@@ -12,12 +12,13 @@
     DependencyModel,
     UpgradeGrantedEvent,
 )
-from ops.model import ActiveStatus, MaintenanceStatus, WaitingStatus
+from ops.model import ActiveStatus, MaintenanceStatus, RelationDataContent, WaitingStatus
 from pydantic import BaseModel
 from tenacity import RetryError, Retrying, stop_after_attempt, wait_fixed
 from typing_extensions import override
 
-from constants import SNAP_PACKAGES
+from constants import APP_SCOPE, MONITORING_PASSWORD_KEY, MONITORING_USER, SNAP_PACKAGES
+from utils import new_password
 
 logger = logging.getLogger(__name__)
 
@@ -26,6 +27,7 @@ class PostgreSQLDependencyModel(BaseModel):
     """PostgreSQL dependencies model."""
 
     charm: DependencyModel
+    snap: DependencyModel
 
 
 def get_postgresql_dependencies_model() -> PostgreSQLDependencyModel:
@@ -42,6 +44,9 @@ def __init__(self, charm, model: BaseModel, **kwargs) -> None:
         """Initialize the class."""
         super().__init__(charm, model, **kwargs)
         self.charm = charm
+        # self.framework.observe(self.charm.on.upgrade_charm, self._on_upgrade_charm_check_legacy)
+        logger.error("running")
+        self._on_upgrade_charm_check_legacy(None)
 
     @override
     def build_upgrade_stack(self) -> List[int]:
@@ -77,6 +82,28 @@ def log_rollback_instructions(self) -> None:
             "Run `juju refresh --revision <previous-revision> postgresql` to initiate the rollback"
         )
 
+    def _on_upgrade_charm_check_legacy(self, event) -> None:
+        if not self.peer_relation or len(self.app_units) < len(self.charm.app_units):
+            logger.debug("Wait all units join the upgrade relation")
+            return
+
+        if self.state:
+            # Do nothing - if state set, upgrade is supported
+            return
+
+        if not self.charm.unit.is_leader():
+            # set ready state on non-leader units
+            self.unit_upgrade_data.update({"state": "ready"})
+            return
+
+        peers_state = list(filter(lambda state: state != "", self.unit_states))
+
+        if len(peers_state) == len(self.peer_relation.units) and (set(peers_state) == {"ready"} or len(peers_state) == 0):
+            # All peers have set the state to ready
+            self.unit_upgrade_data.update({"state": "ready"})
+            self._prepare_upgrade_from_legacy()
+            getattr(self.on, "upgrade_charm").emit()
+
     @override
     def _on_upgrade_granted(self, event: UpgradeGrantedEvent) -> None:
         # Refresh the charmed PostgreSQL snap and restart the database.
@@ -91,6 +118,14 @@ def _on_upgrade_granted(self, event: UpgradeGrantedEvent) -> None:
         self.charm._setup_exporter()
         self.charm.backup.start_stop_pgbackrest_service()
 
+        try:
+            self.charm.unit.set_workload_version(
+                self.charm._patroni.get_postgresql_version() or "unset"
+            )
+        except TypeError:
+            # Don't fail on this, just log it.
+            logger.warning("Failed to get PostgreSQL version")
+
         # Wait until the database initialise.
         self.charm.unit.status = WaitingStatus("waiting for database initialisation")
         try:
@@ -144,3 +179,35 @@ def pre_upgrade_check(self) -> None:
                 "a backup is being created",
                 "wait for the backup creation to finish before starting the upgrade",
             )
+
+    def _prepare_upgrade_from_legacy(self) -> None:
+        """Prepare upgrade from legacy charm without upgrade support.
+
+        Assumes run on leader unit only.
+        """
+        logger.warning("Upgrading from unsupported version")
+
+        # Populate app upgrade databag to allow upgrade procedure
+        logger.debug("Building upgrade stack")
+        upgrade_stack = self.build_upgrade_stack()
+        logger.debug(f"Upgrade stack: {upgrade_stack}")
+        self.upgrade_stack = upgrade_stack
+        logger.debug("Persisting dependencies to upgrade relation data...")
+        self.peer_relation.data[self.charm.app].update(
+            {"dependencies": json.dumps(self.dependency_model.dict())}
+        )
+        if self.charm.get_secret(APP_SCOPE, MONITORING_PASSWORD_KEY) is None:
+            self.charm.set_secret(APP_SCOPE, MONITORING_PASSWORD_KEY, new_password())
+        users = self.charm.postgresql.list_users()
+        if MONITORING_USER not in users:
+            # Create the monitoring user.
+            self.charm.postgresql.create_user(
+                MONITORING_USER,
+                self.charm.get_secret(APP_SCOPE, MONITORING_PASSWORD_KEY),
+                extra_user_roles="pg_monitor",
+            )
+
+    @property
+    def unit_upgrade_data(self) -> RelationDataContent:
+        """Return the application upgrade data."""
+        return self.peer_relation.data[self.charm.unit]

tests/integration/ha_tests/test_upgrade.py

@@ -0,0 +1,168 @@
+#!/usr/bin/env python3
+# Copyright 2023 Canonical Ltd.
+# See LICENSE file for licensing details.
+import json
+import logging
+import subprocess
+import zipfile
+
+import pytest as pytest
+from pytest_operator.plugin import OpsTest
+
+from tests.integration.ha_tests.conftest import APPLICATION_NAME
+from tests.integration.ha_tests.helpers import (
+    app_name,
+    are_writes_increasing,
+    check_writes,
+    start_continuous_writes,
+)
+from tests.integration.helpers import get_primary
+
+logger = logging.getLogger(__name__)
+
+
+@pytest.mark.abort_on_fail
+async def test_build_and_deploy(ops_test: OpsTest) -> None:
+    """Build and deploy three unit of PostgreSQL."""
+    wait_for_apps = False
+    # Check if there is a pre-existing cluster.
+    if not await app_name(ops_test):
+        wait_for_apps = True
+        charm = await ops_test.build_charm(".")
+        async with ops_test.fast_forward():
+            await ops_test.model.deploy(charm, num_units=3)
+    # Deploy the continuous writes application charm if it wasn't already deployed.
+    if not await app_name(ops_test, APPLICATION_NAME):
+        wait_for_apps = True
+        async with ops_test.fast_forward():
+            charm = await ops_test.build_charm("tests/integration/ha_tests/application-charm")
+            await ops_test.model.deploy(charm, application_name=APPLICATION_NAME)
+
+    if wait_for_apps:
+        async with ops_test.fast_forward():
+            await ops_test.model.wait_for_idle(status="active", timeout=1000)
+
+
+async def test_successful_upgrade(ops_test: OpsTest, continuous_writes) -> None:
+    # Start an application that continuously writes data to the database.
+    logger.info("starting continuous writes to the database")
+    app = await app_name(ops_test)
+    await start_continuous_writes(ops_test, app)
+
+    # Check whether writes are increasing.
+    logger.info("checking whether writes are increasing")
+    any_unit_name = next(iter(ops_test.model.applications[app].units)).name
+    primary_name = await get_primary(ops_test, any_unit_name)
+    await are_writes_increasing(ops_test, primary_name)
+
+    # Run the pre-upgrade-check action.
+    logger.info("running pre-upgrade check")
+    leader_unit_name = None
+    for unit in ops_test.model.applications[app].units:
+        if await unit.is_leader_from_status():
+            leader_unit_name = unit.name
+            break
+    action = await ops_test.model.units.get(leader_unit_name).run_action("pre-upgrade-check")
+    await action.wait()
+    assert action.results["Code"] == "0"
+
+    # Run juju refresh.
+    logger.info("refreshing the charm")
+    application = ops_test.model.applications[app]
+    charm = await ops_test.build_charm(".")
+    await application.refresh(path=charm)
+    async with ops_test.fast_forward(fast_interval="30s"):
+        await ops_test.model.wait_for_idle(
+            apps=[app], status="active", idle_period=15, raise_on_blocked=True
+        )
+
+    # Check whether writes are increasing.
+    logger.info("checking whether writes are increasing")
+    primary_name = await get_primary(ops_test, any_unit_name)
+    await are_writes_increasing(ops_test, primary_name)
+
+    # Verify that no writes to the database were missed after stopping the writes
+    # (check that all the units have all the writes).
+    logger.info("checking whether no writes were lost")
+    await check_writes(ops_test)
+
+
+async def test_failed_upgrade(ops_test: OpsTest) -> None:
+    # Run the pre-upgrade-check action.
+    logger.info("running pre-upgrade check")
+    app = await app_name(ops_test)
+    leader_unit_name = None
+    for unit in ops_test.model.applications[app].units:
+        if await unit.is_leader_from_status():
+            leader_unit_name = unit.name
+            break
+    action = await ops_test.model.units.get(leader_unit_name).run_action("pre-upgrade-check")
+    await action.wait()
+    assert action.results["Code"] == "0"
+
+    # Run juju refresh.
+    logger.info("refreshing the charm")
+    charm = await ops_test.build_charm(".")
+    print(f"charm: {charm}")
+    modified_charm = f"{charm}.modified"
+    print(f"modified_charm: {modified_charm}")
+    with zipfile.ZipFile(charm, "r") as charm_file, zipfile.ZipFile(
+        modified_charm, "w"
+    ) as modified_charm_file:
+        # Iterate the input files
+        unix_attributes = {}
+        for charm_info in charm_file.infolist():
+            # Read input file
+            with charm_file.open(charm_info) as file:
+                print(f"charm_info.filename: {charm_info.filename}")
+                if charm_info.filename == "src/dependency.json":
+                    content = json.loads(file.read())
+                    # Modify the content of the file by replacing a string
+                    content["snap"]["upgrade_supported"] = "^15"
+                    content["snap"]["version"] = "15.1"
+                    # Write content.
+                    modified_charm_file.writestr(charm_info.filename, json.dumps(content))
+                else:  # Other file, don't want to modify => just copy it.
+                    content = file.read()
+                    modified_charm_file.writestr(charm_info.filename, content)
+                unix_attributes[charm_info.filename] = charm_info.external_attr >> 16
+
+        for modified_charm_info in modified_charm_file.infolist():
+            modified_charm_info.external_attr = unix_attributes[modified_charm_info.filename] << 16
+    process = subprocess.run(
+        f"juju refresh --model {ops_test.model.info.name} {app} --path {modified_charm}".split(),
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+    )
+    if process.returncode != 0:
+        raise Exception(
+            f"Expected juju refresh command to succeed instead it failed: {process.returncode} - {process.stderr.decode()}"
+        )
+    async with ops_test.fast_forward(fast_interval="30s"):
+        await ops_test.model.wait_for_idle(
+            apps=[app], status="blocked", idle_period=15, raise_on_blocked=False
+        )
+
+
+async def test_rollback(ops_test: OpsTest) -> None:
+    # Run the pre-upgrade-check action.
+    logger.info("running pre-upgrade check")
+    app = await app_name(ops_test)
+    leader_unit_name = None
+    for unit in ops_test.model.applications[app].units:
+        if await unit.is_leader_from_status():
+            leader_unit_name = unit.name
+            break
+    action = await ops_test.model.units.get(leader_unit_name).run_action("pre-upgrade-check")
+    await action.wait()
+    assert action.results["Code"] == "0"
+
+    # Run juju refresh.
+    logger.info("refreshing the charm")
+    application = ops_test.model.applications[app]
+    charm = await ops_test.build_charm(".")
+    await application.refresh(path=charm)
+    async with ops_test.fast_forward(fast_interval="30s"):
+        await ops_test.model.wait_for_idle(
+            apps=[app], status="active", idle_period=15, raise_on_blocked=True
+        )

tox.ini

@@ -182,6 +182,17 @@ commands =
     pip install juju=={env:LIBJUJU}
     poetry run pytest -v --tb native --log-cli-level=INFO -s {posargs} {[vars]tests_path}/integration/test_tls.py
 
+[testenv:upgrade-integration-{juju2, juju3}]
+description = Run upgrade integration tests
+pass_env =
+    {[testenv]pass_env}
+    CI
+    CI_PACKED_CHARMS
+commands =
+    poetry install --with integration
+    pip install juju=={env:LIBJUJU}
+    poetry run pytest -v --tb native --log-cli-level=INFO -s {posargs} {[vars]tests_path}/integration/ha_tests/test_upgrade.py
+
 [testenv:integration-{juju2, juju3}]
 description = Run all integration tests
 pass_env =