Repository: canonical/operator
Pull request #201 (2020-04-01 06:47:40)

Add an event for the "remove" hook
"remove" is fired after stop when a unit is removed.

https://github.com/juju/juju/pull/11237


Patch:
diff --git a/ops/charm.py b/ops/charm.py
index a85b37c..e4db6aa 100755
--- a/ops/charm.py
+++ b/ops/charm.py
@@ -111,6 +111,15 @@ class StopEvent(HookEvent):
     pass
 
 
+class RemoveEvent(HookEvent):
+    """Represents the `remove` hook from Juju.
+
+    It's triggered by `CharmBase.on.remove`.
+    """
+
+    pass
+
+
 class ConfigChangedEvent(HookEvent):
     """Represents the `config-changed` hook from Juju.
 
@@ -375,6 +384,7 @@ class CharmEvents(EventSetBase):
     install = EventSource(InstallEvent)
     start = EventSource(StartEvent)
     stop = EventSource(StopEvent)
+    remove = EventSource(RemoveEvent)
     update_status = EventSource(UpdateStatusEvent)
     config_changed = EventSource(ConfigChangedEvent)
     upgrade_charm = EventSource(UpgradeCharmEvent)
diff --git a/ops/framework.py b/ops/framework.py
index abdd1fe..03e1c88 100755
--- a/ops/framework.py
+++ b/ops/framework.py
@@ -558,8 +558,18 @@ class Framework(Object):
             raise RuntimeError(
                 'cannot save {} values before registering that type'.format(type(value).__name__))
         data = value.snapshot()
-        # Use marshal as a validator, enforcing the use of simple types.
-        marshal.dumps(data)
+
+        # Use marshal as a validator, enforcing the use of simple types, as we later the
+        # information is really pickled, which is too error prone for future evolution of the
+        # stored data (e.g. if the developer stores a custom object and later changes its
+        # class name; when unpickling the original class will not be there and event
+        # data loading will fail).
+        try:
+            marshal.dumps(data)
+        except ValueError:
+            msg = "unable to save the data for {}, it must contain only simple types: {!r}"
+            raise ValueError(msg.format(value.__class__.__name__, data))
+
         # Use pickle for serialization, so the value remains portable.
         raw_data = pickle.dumps(data)
         self._storage.save_snapshot(value.handle.path, raw_data)
@@ -812,39 +822,78 @@ class BoundStoredState:
 
 
 class StoredState:
+    """A class used to store data the charm needs persisted across invocations.
+
+    Example::
+
+        class MyClass(Object):
+            _stored = StoredState()
+
+    Instances of `MyClass` can transparently save state between invocations by
+    setting attributes on `_stored`. Initial state should be set with
+    `set_default` on the bound object, that is::
+
+        class MyClass(Object):
+            _stored = StoredState()
+
+        def __init__(self, parent, key):
+            super().__init__(parent, key)
+            self._stored.set_default(seen=set())
+            self.framework.observe(self.on.seen, self._on_seen)
+
+        def _on_seen(self, event):
+            self._stored.seen.add(event.uuid)
+
+    """
 
     def __init__(self):
         self.parent_type = None
         self.attr_name = None
 
     def __get__(self, parent, parent_type=None):
-        if self.parent_type is None:
-            self.parent_type = parent_type
-        elif self.parent_type is not parent_type:
+        if self.parent_type is not None and self.parent_type not in parent_type.mro():
+            # the StoredState instance is being shared between two unrelated classes
+            # -> unclear what is exepcted of us -> bail out
             raise RuntimeError(
                 'StoredState shared by {} and {}'.format(
                     self.parent_type.__name__, parent_type.__name__))
 
         if parent is None:
+            # accessing via the class directly (e.g. MyClass.stored)
             return self
 
-        bound = parent.__dict__.get(self.attr_name)
-        if bound is None:
-            for attr_name, attr_value in parent_type.__dict__.items():
-                if attr_value is self:
-                    if self.attr_name and attr_name != self.attr_name:
-                        parent_tname = parent_type.__name__
-                        raise RuntimeError("StoredState shared by {}.{} and {}.{}".format(
-                            parent_tname, self.attr_name, parent_tname, attr_name))
-                    self.attr_name = attr_name
-                    bound = BoundStoredState(parent, attr_name)
-                    parent.__dict__[attr_name] = bound
-                    break
-            else:
-                raise RuntimeError(
-                    'cannot find StoredVariable attribute in type {}'.format(parent_type.__name__))
-
-        return bound
+        bound = None
+        if self.attr_name is not None:
+            bound = parent.__dict__.get(self.attr_name)
+            if bound is not None:
+                # we already have the thing from a previous pass, huzzah
+                return bound
+
+        # need to find ourselves amongst the parent's bases
+        for cls in parent_type.mro():
+            for attr_name, attr_value in cls.__dict__.items():
+                if attr_value is not self:
+                    continue
+                # we've found ourselves! is it the first time?
+                if bound is not None:
+                    # the StoredState instance is being stored in two different
+                    # attributes -> unclear what is expected of us -> bail out
+                    raise RuntimeError("StoredState shared by {0}.{1} and {0}.{2}".format(
+                        cls.__name__, self.attr_name, attr_name))
+                # we've found ourselves for the first time; save where, and bind the object
+                self.attr_name = attr_name
+                self.parent_type = cls
+                bound = BoundStoredState(parent, attr_name)
+
+        if bound is not None:
+            # cache the bound object to avoid the expensive lookup the next time
+            # (don't use setattr, to keep things symmetric with the fast-path lookup above)
+            parent.__dict__[self.attr_name] = bound
+            return bound
+
+        raise AttributeError(
+            'cannot find {} attribute in type {}'.format(
+                self.__class__.__name__, parent_type.__name__))
 
 
 def _wrap_stored(parent_data, value):


Test patch:
diff --git a/test/test_framework.py b/test/test_framework.py
index 7671075..08369ee 100644
--- a/test/test_framework.py
+++ b/test/test_framework.py
@@ -221,23 +221,23 @@ class TestFramework(unittest.TestCase):
 
         class PreCommitObserver(Object):
 
-            state = StoredState()
+            _stored = StoredState()
 
             def __init__(self, parent, key):
                 super().__init__(parent, key)
                 self.seen = []
-                self.state.myinitdata = 40
+                self._stored.myinitdata = 40
 
             def on_pre_commit(self, event):
-                self.state.myinitdata = 41
-                self.state.mydata = 42
+                self._stored.myinitdata = 41
+                self._stored.mydata = 42
                 self.seen.append(type(event))
 
             def on_commit(self, event):
                 # Modifications made here will not be persisted.
-                self.state.myinitdata = 42
-                self.state.mydata = 43
-                self.state.myotherdata = 43
+                self._stored.myinitdata = 42
+                self._stored.mydata = 43
+                self._stored.myotherdata = 43
                 self.seen.append(type(event))
 
         obs = PreCommitObserver(framework, None)
@@ -246,8 +246,8 @@ class TestFramework(unittest.TestCase):
 
         framework.commit()
 
-        self.assertEqual(obs.state.myinitdata, 41)
-        self.assertEqual(obs.state.mydata, 42)
+        self.assertEqual(obs._stored.myinitdata, 41)
+        self.assertEqual(obs._stored.mydata, 42)
         self.assertTrue(obs.seen, [PreCommitEvent, CommitEvent])
         framework.close()
 
@@ -255,11 +255,11 @@ class TestFramework(unittest.TestCase):
 
         new_obs = PreCommitObserver(other_framework, None)
 
-        self.assertEqual(obs.state.myinitdata, 41)
-        self.assertEqual(new_obs.state.mydata, 42)
+        self.assertEqual(obs._stored.myinitdata, 41)
+        self.assertEqual(new_obs._stored.mydata, 42)
 
         with self.assertRaises(AttributeError):
-            new_obs.state.myotherdata
+            new_obs._stored.myotherdata
 
     def test_defer_and_reemit(self):
         framework = self.create_framework()
@@ -761,6 +761,26 @@ class TestFramework(unittest.TestCase):
         self.assertIn('database is locked', str(cm.exception))
         f.close()
 
+    def test_snapshot_saving_restricted_to_simple_types(self):
+        # this can not be saved, as it has not simple types!
+        to_be_saved = {"bar": TestFramework}
+
+        class FooEvent(EventSetBase):
+            def snapshot(self):
+                return to_be_saved
+
+        handle = Handle(None, "a_foo", "some_key")
+        event = FooEvent()
+
+        framework = self.create_framework()
+        framework.register_type(FooEvent, None, handle.kind)
+        with self.assertRaises(ValueError) as cm:
+            framework.save_snapshot(event)
+        expected = (
+            "unable to save the data for FooEvent, it must contain only simple types: "
+            "{'bar': <class 'test.test_framework.TestFramework'>}")
+        self.assertEqual(str(cm.exception), expected)
+
 
 class TestStoredState(unittest.TestCase):
 
@@ -774,61 +794,191 @@ class TestStoredState(unittest.TestCase):
         return framework
 
     def test_basic_state_storage(self):
-        framework = self.create_framework()
+        class SomeObject(Object):
+            _stored = StoredState()
+
+        self._stored_state_tests(SomeObject)
 
+    def test_straight_subclass(self):
         class SomeObject(Object):
-            state = StoredState()
+            _stored = StoredState()
+
+        class Sub(SomeObject):
+            pass
+
+        self._stored_state_tests(Sub)
+
+    def test_straight_sub_subclass(self):
+        class SomeObject(Object):
+            _stored = StoredState()
+
+        class Sub(SomeObject):
+            pass
+
+        class SubSub(SomeObject):
+            pass
 
-        obj = SomeObject(framework, "1")
+        self._stored_state_tests(SubSub)
+
+    def test_two_subclasses(self):
+        class SomeObject(Object):
+            _stored = StoredState()
+
+        class SubA(SomeObject):
+            pass
+
+        class SubB(SomeObject):
+            pass
+
+        self._stored_state_tests(SubA)
+        self._stored_state_tests(SubB)
+
+    def test_the_crazy_thing(self):
+        class NoState(Object):
+            pass
+
+        class StatedObject(NoState):
+            _stored = StoredState()
+
+        class Sibling(NoState):
+            pass
+
+        class FinalChild(StatedObject, Sibling):
+            pass
+
+        self._stored_state_tests(FinalChild)
+
+    def _stored_state_tests(self, cls):
+        framework = self.create_framework()
+        obj = cls(framework, "1")
 
         try:
-            obj.state.foo
+            obj._stored.foo
         except AttributeError as e:
             self.assertEqual(str(e), "attribute 'foo' is not stored")
         else:
             self.fail("AttributeError not raised")
 
         try:
-            obj.state.on = "nonono"
+            obj._stored.on = "nonono"
         except AttributeError as e:
             self.assertEqual(str(e), "attribute 'on' is reserved and cannot be set")
         else:
             self.fail("AttributeError not raised")
 
-        obj.state.foo = 41
-        obj.state.foo = 42
-        obj.state.bar = "s"
-        obj.state.baz = 4.2
-        obj.state.bing = True
+        obj._stored.foo = 41
+        obj._stored.foo = 42
+        obj._stored.bar = "s"
+        obj._stored.baz = 4.2
+        obj._stored.bing = True
 
-        self.assertEqual(obj.state.foo, 42)
+        self.assertEqual(obj._stored.foo, 42)
 
         framework.commit()
 
         # This won't be committed, and should not be seen.
-        obj.state.foo = 43
+        obj._stored.foo = 43
 
         framework.close()
 
         # Since this has the same absolute object handle, it will get its state back.
         framework_copy = self.create_framework()
-        obj_copy = SomeObject(framework_copy, "1")
-        self.assertEqual(obj_copy.state.foo, 42)
-        self.assertEqual(obj_copy.state.bar, "s")
-        self.assertEqual(obj_copy.state.baz, 4.2)
-        self.assertEqual(obj_copy.state.bing, True)
+        obj_copy = cls(framework_copy, "1")
+        self.assertEqual(obj_copy._stored.foo, 42)
+        self.assertEqual(obj_copy._stored.bar, "s")
+        self.assertEqual(obj_copy._stored.baz, 4.2)
+        self.assertEqual(obj_copy._stored.bing, True)
+
+        framework_copy.close()
+
+    def test_two_subclasses_no_conflicts(self):
+        class Base(Object):
+            _stored = StoredState()
+
+        class SubA(Base):
+            pass
+
+        class SubB(Base):
+            pass
+
+        framework = self.create_framework()
+        a = SubA(framework, None)
+        b = SubB(framework, None)
+        z = Base(framework, None)
+
+        a._stored.foo = 42
+        b._stored.foo = "hello"
+        z._stored.foo = {1}
+
+        framework.commit()
+        framework.close()
+
+        framework2 = self.create_framework()
+        a2 = SubA(framework2, None)
+        b2 = SubB(framework2, None)
+        z2 = Base(framework2, None)
+
+        self.assertEqual(a2._stored.foo, 42)
+        self.assertEqual(b2._stored.foo, "hello")
+        self.assertEqual(z2._stored.foo, {1})
+
+    def test_two_names_one_state(self):
+        class Mine(Object):
+            _stored = StoredState()
+            _stored2 = _stored
+
+        framework = self.create_framework()
+        obj = Mine(framework, None)
+
+        with self.assertRaises(RuntimeError):
+            obj._stored.foo = 42
+
+        with self.assertRaises(RuntimeError):
+            obj._stored2.foo = 42
+
+        framework.close()
+
+        # make sure we're not changing the object on failure
+        self.assertNotIn("_stored", obj.__dict__)
+        self.assertNotIn("_stored2", obj.__dict__)
+
+    def test_same_name_two_classes(self):
+        class Base(Object):
+            pass
+
+        class A(Base):
+            _stored = StoredState()
+
+        class B(Base):
+            _stored = A._stored
+
+        framework = self.create_framework()
+        a = A(framework, None)
+        b = B(framework, None)
+
+        # NOTE it's the second one that actually triggers the
+        # exception, but that's an implementation detail
+        a._stored.foo = 42
+
+        with self.assertRaises(RuntimeError):
+            b._stored.foo = "xyzzy"
+
+        framework.close()
+
+        # make sure we're not changing the object on failure
+        self.assertNotIn("_stored", b.__dict__)
 
     def test_mutable_types_invalid(self):
         framework = self.create_framework()
 
         class SomeObject(Object):
-            state = StoredState()
+            _stored = StoredState()
 
         obj = SomeObject(framework, '1')
         try:
             class CustomObject:
                 pass
-            obj.state.foo = CustomObject()
+            obj._stored.foo = CustomObject()
         except AttributeError as e:
             self.assertEqual(
                 str(e),
@@ -994,7 +1144,7 @@ class TestStoredState(unittest.TestCase):
         )]
 
         class SomeObject(Object):
-            state = StoredState()
+            _stored = StoredState()
 
         class WrappedFramework(Framework):
             def __init__(self, data_path, charm_dir, meta, model):
@@ -1002,7 +1152,7 @@ class TestStoredState(unittest.TestCase):
                 self.snapshots = []
 
             def save_snapshot(self, value):
-                if value.handle.path == 'SomeObject[1]/StoredStateData[state]':
+                if value.handle.path == 'SomeObject[1]/StoredStateData[_stored]':
                     self.snapshots.append((type(value), value.snapshot()))
                 return super().save_snapshot(value)
 
@@ -1011,13 +1161,13 @@ class TestStoredState(unittest.TestCase):
             framework = self.create_framework(cls=WrappedFramework)
             obj = SomeObject(framework, '1')
 
-            obj.state.a = get_a()
-            self.assertTrue(isinstance(obj.state, BoundStoredState))
+            obj._stored.a = get_a()
+            self.assertTrue(isinstance(obj._stored, BoundStoredState))
 
-            op(obj.state.a, b)
-            validate_op(obj.state.a, expected_res)
+            op(obj._stored.a, b)
+            validate_op(obj._stored.a, expected_res)
 
-            obj.state.a = get_a()
+            obj._stored.a = get_a()
             framework.commit()
             # We should see an update for initializing a
             self.assertEqual(framework.snapshots, [
@@ -1026,10 +1176,10 @@ class TestStoredState(unittest.TestCase):
             del obj
             gc.collect()
             obj_copy1 = SomeObject(framework, '1')
-            self.assertEqual(obj_copy1.state.a, get_a())
+            self.assertEqual(obj_copy1._stored.a, get_a())
 
-            op(obj_copy1.state.a, b)
-            validate_op(obj_copy1.state.a, expected_res)
+            op(obj_copy1._stored.a, b)
+            validate_op(obj_copy1._stored.a, expected_res)
             framework.commit()
             framework.close()
 
@@ -1037,7 +1187,7 @@ class TestStoredState(unittest.TestCase):
 
             obj_copy2 = SomeObject(framework_copy, '1')
 
-            validate_op(obj_copy2.state.a, expected_res)
+            validate_op(obj_copy2._stored.a, expected_res)
 
             # Commit saves the pre-commit and commit events, and the framework
             # event counter, but shouldn't update the stored state of my object
@@ -1129,22 +1279,22 @@ class TestStoredState(unittest.TestCase):
         )]
 
         class SomeObject(Object):
-            state = StoredState()
+            _stored = StoredState()
 
         framework = self.create_framework()
 
         for i, (a, b, op, op_ab, op_ba) in enumerate(test_operations):
             obj = SomeObject(framework, str(i))
-            obj.state.a = a
-            self.assertEqual(op(obj.state.a, b), op_ab)
-            self.assertEqual(op(b, obj.state.a), op_ba)
+            obj._stored.a = a
+            self.assertEqual(op(obj._stored.a, b), op_ab)
+            self.assertEqual(op(b, obj._stored.a), op_ba)
 
     def test_set_operations(self):
         test_operations = [(
             {"1"},  # A set to test an operation against (other_set).
             lambda a, b: a | b,  # An operation to test.
-            {"1", "a", "b"},  # The expected result of operation(obj.state.set, other_set).
-            {"1", "a", "b"}  # The expected result of operation(other_set, obj.state.set).
+            {"1", "a", "b"},  # The expected result of operation(obj._stored.set, other_set).
+            {"1", "a", "b"}  # The expected result of operation(other_set, obj._stored.set).
         ), (
             {"a", "c"},
             lambda a, b: a - b,
@@ -1168,7 +1318,7 @@ class TestStoredState(unittest.TestCase):
         )]
 
         class SomeObject(Object):
-            state = StoredState()
+            _stored = StoredState()
 
         framework = self.create_framework()
 
@@ -1178,11 +1328,11 @@ class TestStoredState(unittest.TestCase):
         # original sets are not changed or used as a result.
         for i, (variable_operand, operation, ab_res, ba_res) in enumerate(test_operations):
             obj = SomeObject(framework, str(i))
-            obj.state.set = {"a", "b"}
+            obj._stored.set = {"a", "b"}
 
             for a, b, expected in [
-                    (obj.state.set, variable_operand, ab_res),
-                    (variable_operand, obj.state.set, ba_res)]:
+                    (obj._stored.set, variable_operand, ab_res),
+                    (variable_operand, obj._stored.set, ba_res)]:
                 old_a = set(a)
                 old_b = set(b)
 
@@ -1190,7 +1340,7 @@ class TestStoredState(unittest.TestCase):
                 self.assertEqual(result, expected)
 
                 # Common sanity checks
-                self.assertIsNot(obj.state.set._under, result)
+                self.assertIsNot(obj._stored.set._under, result)
                 self.assertIsNot(result, a)
                 self.assertIsNot(result, b)
                 self.assertEqual(a, old_a)
@@ -1200,22 +1350,22 @@ class TestStoredState(unittest.TestCase):
         framework = self.create_framework()
 
         class StatefulObject(Object):
-            state = StoredState()
+            _stored = StoredState()
         parent = StatefulObject(framework, 'key')
-        parent.state.set_default(foo=1)
-        self.assertEqual(parent.state.foo, 1)
-        parent.state.set_default(foo=2)
+        parent._stored.set_default(foo=1)
+        self.assertEqual(parent._stored.foo, 1)
+        parent._stored.set_default(foo=2)
         # foo was already set, so it doesn't get replaced
-        self.assertEqual(parent.state.foo, 1)
-        parent.state.set_default(foo=3, bar=4)
-        self.assertEqual(parent.state.foo, 1)
-        self.assertEqual(parent.state.bar, 4)
+        self.assertEqual(parent._stored.foo, 1)
+        parent._stored.set_default(foo=3, bar=4)
+        self.assertEqual(parent._stored.foo, 1)
+        self.assertEqual(parent._stored.bar, 4)
         # reloading the state still leaves things at the default values
         framework.commit()
         del parent
         parent = StatefulObject(framework, 'key')
-        parent.state.set_default(foo=5, bar=6)
-        self.assertEqual(parent.state.foo, 1)
-        self.assertEqual(parent.state.bar, 4)
+        parent._stored.set_default(foo=5, bar=6)
+        self.assertEqual(parent._stored.foo, 1)
+        self.assertEqual(parent._stored.bar, 4)
         # TODO: jam 2020-01-30 is there a clean way to tell that
-        #       parent.state._data.dirty is False?
+        #       parent._stored._data.dirty is False?