Repository: canonical/operator
Pull request #483 (2021-02-25 01:33:22)

Dynamic add_layer API and get_plan API
The Python / Operator Framework end of https://github.com/canonical/pebble/pull/10

This adds `add_layer` and `get_plan` to both the `Pebble` instance as well as `model.Container`.

Note that these changes would be a breaking change, but the functions in question currently raise `NotImplementedError`, and nobody except me is using it yet anyway.


Patch:
diff --git a/ops/model.py b/ops/model.py
index d03a387..9c5e85c 100644
--- a/ops/model.py
+++ b/ops/model.py
@@ -1047,21 +1047,26 @@ class Container:
         """Stop given service(s) by name."""
         self._pebble.stop_services(service_names)
 
-    def add_layer(self, layer: typing.Union[str, typing.Dict, 'pebble.Layer']):
-        """Dynamically add a setup layer.
+    def add_layer(self, label: str, layer: typing.Union[str, typing.Dict, 'pebble.Layer'],
+                  combine: bool = False):
+        """Dynamically add a new layer onto the Pebble configuration layers.
 
         Args:
-            layer: A YAML string, setup layer dict, or pebble.Layer object
-                containing the Pebble layer to add.
+            label: Label for new layer (and label of layer to merge with if
+                combining).
+            layer: A YAML string, configuration layer dict, or pebble.Layer
+                object containing the Pebble layer to add.
+            combine: If combine is False (the default), append the new layer
+                as the top layer with the given label (must be unique). If
+                combine is True and the label already exists, the two layers
+                are combined into a single one considering the layer override
+                rules; if the layer doesn't exist, it is added as usual.
         """
-        if isinstance(layer, dict):
-            layer = pebble.Layer(layer)
-        self._pebble.add_layer(layer)
-
-    def get_layer(self) -> 'pebble.Layer':
-        """Fetch the flattened setup layers as a pebble.Layer object."""
-        setup_yaml = self._pebble.get_layer()
-        return pebble.Layer(setup_yaml)
+        self._pebble.add_layer(label, layer, combine=combine)
+
+    def get_plan(self) -> 'pebble.Plan':
+        """Get the current effective pebble configuration."""
+        return self._pebble.get_plan()
 
 
 class ContainerMapping(Mapping):
diff --git a/ops/pebble.py b/ops/pebble.py
index 1fddfe8..b95c969 100644
--- a/ops/pebble.py
+++ b/ops/pebble.py
@@ -386,11 +386,35 @@ class Change:
                 ).format(self=self)
 
 
+class Plan:
+    """Represents the effective Pebble configuration."""
+
+    def __init__(self, raw: str):
+        d = yaml.safe_load(raw) or {}
+        self._raw = raw
+        self._services = {name: Service(name, service)
+                          for name, service in d.get('services', {}).items()}
+
+    @property
+    def services(self):
+        """This plan's services mapping (maps service name to Service).
+
+        This property is currently read-only.
+        """
+        return self._services
+
+    def to_yaml(self) -> str:
+        """Return this plan's YAML representation."""
+        return self._raw
+
+    __str__ = to_yaml
+
+
 class Layer:
-    """Represents a Pebble setup layer (or flattened setup).
+    """Represents a Pebble configuration layer.
 
     The format of this is not documented, but is captured in code here:
-    https://github.com/canonical/pebble/blob/master/internal/setup/setup.go
+    https://github.com/canonical/pebble/blob/master/internal/plan/plan.go
     """
 
     def __init__(self, raw: Union[str, Dict] = None):
@@ -423,7 +447,7 @@ class Layer:
 
 
 class Service:
-    """Represents a service description in a Pebble setup layer."""
+    """Represents a service description in a Pebble configuration layer."""
 
     def __init__(self, name: str, raw: Dict = None):
         self.name = name
@@ -627,19 +651,37 @@ class Client:
         raise TimeoutError(
             'timed out waiting for change {} ({} seconds)'.format(change_id, timeout))
 
-    def add_layer(self, layer: Union[str, dict, Layer]):
-        """Dynamically add a layer to the Pebble setup."""
+    def add_layer(self, label: str, layer: Union[str, dict, Layer], combine: bool = False):
+        """Dynamically add a new layer onto the Pebble configuration layers.
+
+        If combine is False (the default), append the new layer as the top
+        layer with the given label. If combine is True and the label already
+        exists, the two layers are combined into a single one considering the
+        layer override rules; if the layer doesn't exist, it is added as usual.
+        """
+        if not isinstance(label, str):
+            raise TypeError('label must be a str, not {}'.format(type(label).__name__))
+
         if isinstance(layer, str):
             layer_yaml = layer
         elif isinstance(layer, dict):
             layer_yaml = Layer(layer).to_yaml()
-        else:
+        elif isinstance(layer, Layer):
             layer_yaml = layer.to_yaml()
-        _ = layer_yaml
-        # TODO(benhoyt) - send layer_yaml to Pebble when that API is implemented
-        raise NotImplementedError('add_layer not yet implemented in Pebble')
-
-    def get_layer(self) -> str:
-        """Get the flattened setup layers as a YAML string."""
-        # TODO(benhoyt) - fetch setup YAML from Pebble when that API is implemented
-        raise NotImplementedError('get_layer not yet implemented in Pebble')
+        else:
+            raise TypeError('layer must be str, dict, or pebble.Layer, not {}'.format(
+                type(layer).__name__))
+
+        body = {
+            'action': 'add',
+            'combine': combine,
+            'label': label,
+            'format': 'yaml',
+            'layer': layer_yaml,
+        }
+        self._request('POST', '/v1/layers', body=body)
+
+    def get_plan(self) -> Plan:
+        """Get the Pebble plan (currently contains only combined services)."""
+        result = self._request('GET', '/v1/plan', {'format': 'yaml'})
+        return Plan(result['result'])


Test patch:
diff --git a/test/pebble_cli.py b/test/pebble_cli.py
index cbe8721..706ce51 100644
--- a/test/pebble_cli.py
+++ b/test/pebble_cli.py
@@ -42,6 +42,11 @@ def main():
                                        'format), default current time',
                    type=pebble._parse_timestamp)
 
+    p = subparsers.add_parser('add', help='add a configuration layer dynamically')
+    p.add_argument('--combine', action='store_true', help='combine layer instead of appending')
+    p.add_argument('label', help='label for new layer')
+    p.add_argument('layer_path', help='path of layer YAML file')
+
     p = subparsers.add_parser('autostart', help='autostart default service(s)')
 
     p = subparsers.add_parser('change', help='show a single change by ID')
@@ -52,6 +57,8 @@ def main():
                    choices=[s.value for s in pebble.ChangeState], default='all')
     p.add_argument('--service', help='optional service name to filter on')
 
+    p = subparsers.add_parser('plan', help='show configuration plan (combined layers)')
+
     p = subparsers.add_parser('start', help='start service(s)')
     p.add_argument('service', help='name of service to start (can specify multiple)', nargs='+')
 
@@ -85,6 +92,15 @@ def main():
         elif args.command == 'ack':
             timestamp = args.timestamp or datetime.datetime.now(tz=datetime.timezone.utc)
             result = client.ack_warnings(timestamp)
+        elif args.command == 'add':
+            try:
+                with open(args.layer_path, encoding='utf-8') as f:
+                    layer_yaml = f.read()
+            except IOError as e:
+                parser.error('cannot read layer YAML: {}'.format(e))
+            client.add_layer(args.label, layer_yaml, combine=bool(args.combine))
+            result = 'Layer {!r} added successfully from {!r}'.format(
+                args.label, args.layer_path)
         elif args.command == 'autostart':
             result = client.autostart_services()
         elif args.command == 'change':
@@ -92,6 +108,8 @@ def main():
         elif args.command == 'changes':
             result = client.get_changes(select=pebble.ChangeState(args.select),
                                         service=args.service)
+        elif args.command == 'plan':
+            result = client.get_plan().raw_yaml
         elif args.command == 'start':
             result = client.start_services(args.service)
         elif args.command == 'stop':
diff --git a/test/test_model.py b/test/test_model.py
index afe18d6..809d679 100755
--- a/test/test_model.py
+++ b/test/test_model.py
@@ -849,21 +849,24 @@ containers:
             container.stop(['foo'])
 
     def test_add_layer(self):
-        self.container.add_layer('summary: str\n')
-        self.container.add_layer({'summary': 'dict'})
-        self.container.add_layer(ops.pebble.Layer('summary: Layer'))
+        self.container.add_layer('a', 'summary: str\n')
+        self.container.add_layer('b', {'summary': 'dict'})
+        self.container.add_layer('c', ops.pebble.Layer('summary: Layer'))
+        self.container.add_layer('d', 'summary: str\n', combine=True)
         self.assertEqual(self.pebble.requests, [
-            ('add_layer', 'summary: str\n'),
-            ('add_layer', 'summary: dict\n'),
-            ('add_layer', 'summary: Layer\n'),
+            ('add_layer', 'a', 'summary: str\n', False),
+            ('add_layer', 'b', 'summary: dict\n', False),
+            ('add_layer', 'c', 'summary: Layer\n', False),
+            ('add_layer', 'd', 'summary: str\n', True),
         ])
 
-    def test_get_layer(self):
-        self.pebble.responses.append('summary: foo')
-        layer = self.container.get_layer()
-        self.assertEqual(self.pebble.requests, [('get_layer',)])
-        self.assertIsInstance(layer, ops.pebble.Layer)
-        self.assertEqual(layer.summary, 'foo')
+    def test_get_plan(self):
+        plan_yaml = 'services:\n foo:\n  override: replace\n  command: bar'
+        self.pebble.responses.append(ops.pebble.Plan(plan_yaml))
+        plan = self.container.get_plan()
+        self.assertEqual(self.pebble.requests, [('get_plan',)])
+        self.assertIsInstance(plan, ops.pebble.Plan)
+        self.assertEqual(plan.to_yaml(), plan_yaml)
 
 
 class MockPebbleBackend(ops.model._ModelBackend):
@@ -886,13 +889,15 @@ class MockPebbleClient:
     def stop_services(self, service_names):
         self.requests.append(('stop', service_names))
 
-    def add_layer(self, layer):
-        if isinstance(layer, ops.pebble.Layer):
+    def add_layer(self, label, layer, combine=False):
+        if isinstance(layer, dict):
+            layer = ops.pebble.Layer(layer).to_yaml()
+        elif isinstance(layer, ops.pebble.Layer):
             layer = layer.to_yaml()
-        self.requests.append(('add_layer', layer))
+        self.requests.append(('add_layer', label, layer, combine))
 
-    def get_layer(self):
-        self.requests.append(('get_layer',))
+    def get_plan(self):
+        self.requests.append(('get_plan',))
         return self.responses.pop(0)
 
 
diff --git a/test/test_pebble.py b/test/test_pebble.py
index d1255b1..6539381 100644
--- a/test/test_pebble.py
+++ b/test/test_pebble.py
@@ -338,6 +338,37 @@ class TestTypes(unittest.TestCase):
         self.assertEqual(change.spawn_time, datetime_utc(2021, 1, 28, 14, 37, 2, 247202))
 
 
+class TestPlan(unittest.TestCase):
+    def test_no_args(self):
+        with self.assertRaises(TypeError):
+            pebble.Plan()
+
+    def test_services(self):
+        plan = pebble.Plan('')
+        self.assertEqual(plan.services, {})
+
+        plan = pebble.Plan('services:\n foo:\n  override: replace\n  command: echo foo')
+
+        self.assertEqual(len(plan.services), 1)
+        self.assertEqual(plan.services['foo'].name, 'foo')
+        self.assertEqual(plan.services['foo'].override, 'replace')
+        self.assertEqual(plan.services['foo'].command, 'echo foo')
+
+        # Should be read-only ("can't set attribute")
+        with self.assertRaises(AttributeError):
+            plan.services = {}
+
+    def test_yaml(self):
+        plan = pebble.Plan('')
+        self.assertEqual(plan.to_yaml(), '')
+        self.assertEqual(str(plan), '')
+
+        yaml = 'services:\n foo:\n  override: replace\n  command: echo foo'
+        plan = pebble.Plan(yaml)
+        self.assertEqual(plan.to_yaml(), yaml)
+        self.assertEqual(str(plan), yaml)
+
+
 class TestLayer(unittest.TestCase):
     def _assert_empty(self, layer):
         self.assertEqual(layer.summary, '')
@@ -788,6 +819,11 @@ class TestClient(unittest.TestCase):
         self.assertIsInstance(cm.exception.change, pebble.Change)
         self.assertEqual(cm.exception.change.id, '70')
 
+        self.assertEqual(self.client.requests, [
+            ('POST', '/v1/services', None, {'action': 'autostart', 'services': []}),
+            ('GET', '/v1/changes/70', None, None),
+        ])
+
     def test_wait_change_timeout(self):
         with unittest.mock.patch('ops.pebble.time', MockTime()):
             change = self.build_mock_change_dict()
@@ -825,6 +861,80 @@ class TestClient(unittest.TestCase):
         self.assertEqual(response.id, '70')
         self.assertEqual(response.err, 'Some kind of service error')
 
+        self.assertEqual(self.client.requests, [
+            ('GET', '/v1/changes/70', None, None),
+        ])
+
+    def test_add_layer(self):
+        okay_response = {
+            "result": True,
+            "status": "OK",
+            "status-code": 200,
+            "type": "sync"
+        }
+        self.client.responses.append(okay_response)
+        self.client.responses.append(okay_response)
+        self.client.responses.append(okay_response)
+        self.client.responses.append(okay_response)
+
+        layer_yaml = """
+services:
+  foo:
+    command: echo bar
+    override: replace
+"""[1:]
+        layer = pebble.Layer(layer_yaml)
+
+        self.client.add_layer('a', layer)
+        self.client.add_layer('b', layer.to_yaml())
+        self.client.add_layer('c', layer.to_dict())
+        self.client.add_layer('d', layer, combine=True)
+
+        def build_expected(label, combine):
+            return {
+                'action': 'add',
+                'combine': combine,
+                'label': label,
+                'format': 'yaml',
+                'layer': layer_yaml,
+            }
+
+        self.assertEqual(self.client.requests, [
+            ('POST', '/v1/layers', None, build_expected('a', False)),
+            ('POST', '/v1/layers', None, build_expected('b', False)),
+            ('POST', '/v1/layers', None, build_expected('c', False)),
+            ('POST', '/v1/layers', None, build_expected('d', True)),
+        ])
+
+    def test_add_layer_invalid_type(self):
+        with self.assertRaises(TypeError):
+            self.client.add_layer('foo', 42)
+        with self.assertRaises(TypeError):
+            self.client.add_layer(42, 'foo')
+
+    def test_get_plan(self):
+        plan_yaml = """
+services:
+  foo:
+    command: echo bar
+    override: replace
+"""[1:]
+        self.client.responses.append({
+            "result": plan_yaml,
+            "status": "OK",
+            "status-code": 200,
+            "type": "sync"
+        })
+        plan = self.client.get_plan()
+        self.assertEqual(plan.to_yaml(), plan_yaml)
+        self.assertEqual(len(plan.services), 1)
+        self.assertEqual(plan.services['foo'].command, 'echo bar')
+        self.assertEqual(plan.services['foo'].override, 'replace')
+
+        self.assertEqual(self.client.requests, [
+            ('GET', '/v1/plan', {'format': 'yaml'}, None),
+        ])
+
 
 class TestSocketClient(unittest.TestCase):
     @unittest.skipIf(sys.platform == 'win32', "Unix sockets don't work on Windows")