From 99e69fdc9d350c1247142f6474c6985620334d4b Mon Sep 17 00:00:00 2001
From: Guillaume Boutry <guillaume.boutry@canonical.com>
Date: Mon, 24 Feb 2025 11:06:54 +0100
Subject: [PATCH] [ops-sunbeam] Allow post-init to throw status exceptions

When the setup of relation handlers throws an ops sunbeam status
exception, the charm is put to error while this is a supported patterns
for developping charms. The reason is that the exception is not thrown
from within a guard. But it is reasonable, for example, for
`OSBaseOperatorAPICharm.internal_url` to raise a WaitingExceptionError
instead of returning None.

Change-Id: Ide137421308733784b6aca7e247eb3e13485d2ff
Signed-off-by: Guillaume Boutry <guillaume.boutry@canonical.com>
---
 ops-sunbeam/ops_sunbeam/core.py  | 20 +++++++++++++++++++-
 ops-sunbeam/ops_sunbeam/guard.py | 26 +++++++++++++++++---------
 2 files changed, 36 insertions(+), 10 deletions(-)

diff --git a/ops-sunbeam/ops_sunbeam/core.py b/ops-sunbeam/ops_sunbeam/core.py
index 4017bfd0..e2fb7b40 100644
--- a/ops-sunbeam/ops_sunbeam/core.py
+++ b/ops-sunbeam/ops_sunbeam/core.py
@@ -25,6 +25,9 @@ from typing import (
 )
 
 import ops_sunbeam.tracing as sunbeam_tracing
+from ops_sunbeam.guard import (
+    BaseStatusExceptionError,
+)
 
 if TYPE_CHECKING:
     from ops_sunbeam.charm import (
@@ -101,5 +104,20 @@ class PostInitMeta(type):
     def __call__(cls, *args, **kw):
         """Call __post_init__ after __init__."""
         instance = super().__call__(*args, **kw)
-        instance.__post_init__()
+        try:
+            instance.__post_init__()
+        except BaseStatusExceptionError as e:
+            # Allow post init to raise an ops_sunbeam status
+            # exception without causing the charm to error.
+            # This status will be collected and set on the
+            # unit.
+            # import here to avoid circular import
+            from ops_sunbeam.charm import (
+                OSBaseOperatorCharm,
+            )
+
+            if isinstance(instance, OSBaseOperatorCharm):
+                instance.status.set(e.to_status())
+            else:
+                raise e
         return instance
diff --git a/ops-sunbeam/ops_sunbeam/guard.py b/ops-sunbeam/ops_sunbeam/guard.py
index 331d3057..c32120a0 100644
--- a/ops-sunbeam/ops_sunbeam/guard.py
+++ b/ops-sunbeam/ops_sunbeam/guard.py
@@ -20,9 +20,11 @@ from contextlib import (
     contextmanager,
 )
 
-from ops.model import (
+from ops import (
+    ActiveStatus,
     BlockedStatus,
     MaintenanceStatus,
+    StatusBase,
     WaitingStatus,
 )
 
@@ -43,27 +45,33 @@ class GuardExceptionError(Exception):
 class BaseStatusExceptionError(Exception):
     """Charm is blocked."""
 
-    def __init__(self, msg):
+    status_type: type[StatusBase] = ActiveStatus
+
+    def __init__(self, msg: str):
+        super().__init__(msg)
         self.msg = msg
-        super().__init__(self.msg)
+
+    def to_status(self):
+        """Convert the exception to an ops status."""
+        return self.status_type(self.msg)
 
 
 class BlockedExceptionError(BaseStatusExceptionError):
     """Charm is blocked."""
 
-    pass
+    status_type = BlockedStatus
 
 
 class MaintenanceExceptionError(BaseStatusExceptionError):
     """Charm is performing maintenance."""
 
-    pass
+    status_type = MaintenanceStatus
 
 
 class WaitingExceptionError(BaseStatusExceptionError):
     """Charm is waiting."""
 
-    pass
+    status_type = WaitingStatus
 
 
 @contextmanager
@@ -103,19 +111,19 @@ def guard(
         logger.warning(
             "Charm is blocked in section '%s' due to '%s'", section, str(e)
         )
-        charm.status.set(BlockedStatus(e.msg))
+        charm.status.set(e.to_status())
     except WaitingExceptionError as e:
         logger.warning(
             "Charm is waiting in section '%s' due to '%s'", section, str(e)
         )
-        charm.status.set(WaitingStatus(e.msg))
+        charm.status.set(e.to_status())
     except MaintenanceExceptionError as e:
         logger.warning(
             "Charm performing maintenance in section '%s' due to '%s'",
             section,
             str(e),
         )
-        charm.status.set(MaintenanceStatus(e.msg))
+        charm.status.set(e.to_status())
     except Exception as e:
         # something else went wrong
         if handle_exception: