[DPE-4239] Warn user when deploying charm with wrong architecture (#613)

* check architecture

* update error message

* try checking inside init

* try another approach

* relocate imports

* try fake charm approach

* refactor change

* add unit tests + more refactor

* try adding integration tests

* fix markers

* try fix integration test

* fix typo

* remove self

* use wait for idle instead of block until

* nits

src/arch_utils.py

@@ -0,0 +1,38 @@
+# Copyright 2024 Canonical Ltd.
+# See LICENSE file for licensing details.
+
+"""Utilities for catching and raising architecture errors."""
+
+import os
+import sys
+
+from ops.charm import CharmBase
+from ops.model import BlockedStatus
+
+
+class WrongArchitectureWarningCharm(CharmBase):
+    """A fake charm class that only signals a wrong architecture deploy."""
+
+    def __init__(self, *args):
+        super().__init__(*args)
+        self.unit.status = BlockedStatus(
+            f"Error: Charm version incompatible with {os.uname().machine} architecture"
+        )
+        sys.exit(0)
+
+
+def is_wrong_architecture() -> bool:
+    """Checks if charm was deployed on wrong architecture."""
+    juju_charm_file = f"{os.environ.get('CHARM_DIR')}/.juju-charm"
+    if not os.path.exists(juju_charm_file):
+        return False
+
+    with open(juju_charm_file, "r") as file:
+        ch_platform = file.read()
+    hw_arch = os.uname().machine
+    if ("amd64" in ch_platform and hw_arch == "x86_64") or (
+        "arm64" in ch_platform and hw_arch == "aarch64"
+    ):
+        return False
+
+    return True

src/charm.py

@@ -13,7 +13,21 @@
 from pathlib import Path
 from typing import Dict, List, Literal, Optional, Tuple, get_args
 
-import psycopg2
+# First platform-specific import, will fail on wrong architecture
+try:
+    import psycopg2
+except ModuleNotFoundError:
+    from ops.main import main
+
+    from arch_utils import WrongArchitectureWarningCharm, is_wrong_architecture
+
+    # If the charm was deployed inside a host with different architecture
+    # (possibly due to user specifying an incompatible revision)
+    # then deploy an empty blocked charm with a warning.
+    if is_wrong_architecture() and __name__ == "__main__":
+        main(WrongArchitectureWarningCharm, use_juju_for_storage=True)
+    raise
+
 from charms.data_platform_libs.v0.data_interfaces import DataPeerData, DataPeerUnitData
 from charms.data_platform_libs.v0.data_models import TypedCharmBase
 from charms.grafana_k8s.v0.grafana_dashboard import GrafanaDashboardProvider

tests/integration/test_wrong_arch.py

@@ -0,0 +1,79 @@
+#!/usr/bin/env python3
+# Copyright 2024 Canonical Ltd.
+# See LICENSE file for licensing details.
+
+import logging
+import os
+import pathlib
+import typing
+
+import pytest
+import yaml
+from pytest_operator.plugin import OpsTest
+
+from . import markers
+from .helpers import CHARM_SERIES, DATABASE_APP_NAME, METADATA
+
+logger = logging.getLogger(__name__)
+
+
+async def fetch_charm(
+    charm_path: typing.Union[str, os.PathLike],
+    architecture: str,
+    bases_index: int,
+) -> pathlib.Path:
+    """Fetches packed charm from CI runner without checking for architecture."""
+    charm_path = pathlib.Path(charm_path)
+    charmcraft_yaml = yaml.safe_load((charm_path / "charmcraft.yaml").read_text())
+    assert charmcraft_yaml["type"] == "charm"
+    base = charmcraft_yaml["bases"][bases_index]
+    build_on = base.get("build-on", [base])[0]
+    version = build_on["channel"]
+    packed_charms = list(charm_path.glob(f"*{version}-{architecture}.charm"))
+    return packed_charms[0].resolve(strict=True)
+
+
+@pytest.mark.group(1)
+@markers.amd64_only
+async def test_arm_charm_on_amd_host(ops_test: OpsTest) -> None:
+    """Tries deploying an arm64 charm on amd64 host."""
+    charm = await fetch_charm(".", "arm64", 1)
+    resources = {
+        "postgresql-image": METADATA["resources"]["postgresql-image"]["upstream-source"],
+    }
+    await ops_test.model.deploy(
+        charm,
+        resources=resources,
+        application_name=DATABASE_APP_NAME,
+        trust=True,
+        num_units=1,
+        series=CHARM_SERIES,
+        config={"profile": "testing"},
+    )
+
+    await ops_test.model.wait_for_idle(
+        apps=[DATABASE_APP_NAME], raise_on_error=False, status="blocked"
+    )
+
+
+@pytest.mark.group(1)
+@markers.arm64_only
+async def test_amd_charm_on_arm_host(ops_test: OpsTest) -> None:
+    """Tries deploying an amd64 charm on arm64 host."""
+    charm = await fetch_charm(".", "amd64", 0)
+    resources = {
+        "postgresql-image": METADATA["resources"]["postgresql-image"]["upstream-source"],
+    }
+    await ops_test.model.deploy(
+        charm,
+        resources=resources,
+        application_name=DATABASE_APP_NAME,
+        trust=True,
+        num_units=1,
+        series=CHARM_SERIES,
+        config={"profile": "testing"},
+    )
+
+    await ops_test.model.wait_for_idle(
+        apps=[DATABASE_APP_NAME], raise_on_error=False, status="blocked"
+    )

tests/unit/test_arch_utils.py

@@ -0,0 +1,76 @@
+# Copyright 2024 Canonical Ltd.
+# See LICENSE file for licensing details.
+import builtins
+import sys
+import unittest.mock as mock
+from unittest.mock import patch
+
+import pytest
+
+from arch_utils import is_wrong_architecture
+
+real_import = builtins.__import__
+
+
+def psycopg2_not_found(name, globals=None, locals=None, fromlist=(), level=0):  # noqa: A002
+    """Fake import function to simulate psycopg2 import error."""
+    if name == "psycopg2":
+        raise ModuleNotFoundError(f"Mocked module not found {name}")
+    return real_import(name, globals=globals, locals=locals, fromlist=fromlist, level=level)
+
+
+def test_on_module_not_found_error(monkeypatch):
+    """Checks if is_wrong_architecture is called on ModuleNotFoundError."""
+    with patch("arch_utils.is_wrong_architecture") as _is_wrong_arch:
+        # If psycopg2 not there, charm should check architecture
+        monkeypatch.delitem(sys.modules, "psycopg2", raising=False)
+        monkeypatch.delitem(sys.modules, "charm", raising=False)
+        monkeypatch.setattr(builtins, "__import__", psycopg2_not_found)
+        with pytest.raises(ModuleNotFoundError):
+            import charm  # noqa: F401
+
+        _is_wrong_arch.assert_called_once()
+
+        # If no import errors, charm continues as normal
+        _is_wrong_arch.reset_mock()
+        monkeypatch.setattr(builtins, "__import__", real_import)
+        import charm  # noqa: F401
+
+        _is_wrong_arch.assert_not_called()
+
+
+def test_wrong_architecture_file_not_found():
+    """Tests if the function returns False when the charm file doesn't exist."""
+    with (
+        patch("os.environ.get", return_value="/tmp"),
+        patch("os.path.exists", return_value=False),
+    ):
+        assert not is_wrong_architecture()
+
+
+def test_wrong_architecture_amd64():
+    """Tests if the function correctly identifies arch when charm is AMD."""
+    with (
+        patch("os.environ.get", return_value="/tmp"),
+        patch("os.path.exists", return_value=True),
+        patch("builtins.open", mock.mock_open(read_data="amd64\n")),
+        patch("os.uname") as _uname,
+    ):
+        _uname.return_value = mock.Mock(machine="x86_64")
+        assert not is_wrong_architecture()
+        _uname.return_value = mock.Mock(machine="aarch64")
+        assert is_wrong_architecture()
+
+
+def test_wrong_architecture_arm64():
+    """Tests if the function correctly identifies arch when charm is ARM."""
+    with (
+        patch("os.environ.get", return_value="/tmp"),
+        patch("os.path.exists", return_value=True),
+        patch("builtins.open", mock.mock_open(read_data="arm64\n")),
+        patch("os.uname") as _uname,
+    ):
+        _uname.return_value = mock.Mock(machine="x86_64")
+        assert is_wrong_architecture()
+        _uname.return_value = mock.Mock(machine="aarch64")
+        assert not is_wrong_architecture()