Repository: canonical/operator

Issue:
missing `kill-delay` pebble service support
Pebble merged `kill-delay` support but ops does not yet support it. (See pebble issue [190](https://github.com/canonical/pebble/pull/190/files))



Discussion:

Patch:
diff --git a/ops/jujuversion.py b/ops/jujuversion.py
index bb8151b..57153e3 100755
--- a/ops/jujuversion.py
+++ b/ops/jujuversion.py
@@ -128,3 +128,15 @@ class JujuVersion:
         """Report whether this Juju version supports open-port on Kubernetes."""
         # Support added: https://bugs.launchpad.net/juju/+bug/1920960
         return (self.major, self.minor, self.patch) >= (3, 0, 3)
+
+    @property
+    def supports_exec_service_context(self) -> bool:
+        """Report whether this Juju version supports exec's service_context option."""
+        if (self.major, self.minor, self.patch) < (3, 1, 6):
+            # First released in 3.1.6
+            return False
+        if (self.major, self.minor, self.patch) == (3, 2, 0):
+            # 3.2.0 was released before Pebble was updated, but all other 3.2
+            # releases have the change (3.2.1 tag was never released).
+            return False
+        return True
diff --git a/ops/model.py b/ops/model.py
index 4756b2c..24f8568 100644
--- a/ops/model.py
+++ b/ops/model.py
@@ -2360,6 +2360,7 @@ class Container:
         self,
         command: List[str],
         *,
+        service_context: Optional[str] = None,
         environment: Optional[Dict[str, str]] = None,
         working_dir: Optional[str] = None,
         timeout: Optional[float] = None,
@@ -2381,6 +2382,7 @@ class Container:
         self,
         command: List[str],
         *,
+        service_context: Optional[str] = None,
         environment: Optional[Dict[str, str]] = None,
         working_dir: Optional[str] = None,
         timeout: Optional[float] = None,
@@ -2400,6 +2402,7 @@ class Container:
         self,
         command: List[str],
         *,
+        service_context: Optional[str] = None,
         environment: Optional[Dict[str, str]] = None,
         working_dir: Optional[str] = None,
         timeout: Optional[float] = None,
@@ -2418,8 +2421,14 @@ class Container:
         See :meth:`ops.pebble.Client.exec` for documentation of the parameters
         and return value, as well as examples.
         """
+        if service_context is not None:
+            version = JujuVersion.from_environ()
+            if not version.supports_exec_service_context:
+                raise RuntimeError(
+                    f'exec with service_context not supported on Juju version {version}')
         return self._pebble.exec(
             command,
+            service_context=service_context,
             environment=environment,
             working_dir=working_dir,
             timeout=timeout,
diff --git a/ops/pebble.py b/ops/pebble.py
index 13e3a9f..35ad53c 100644
--- a/ops/pebble.py
+++ b/ops/pebble.py
@@ -79,18 +79,36 @@ ServiceDict = typing.TypedDict('ServiceDict',
                                 'user-id': Optional[int],
                                 'group': str,
                                 'group-id': Optional[int],
+                                'working-dir': str,
                                 'on-success': str,
                                 'on-failure': str,
                                 'on-check-failure': Dict[str, Any],
                                 'backoff-delay': str,
                                 'backoff-factor': Optional[int],
                                 'backoff-limit': str,
+                                'kill-delay': Optional[str],
                                 },
                                total=False)
 
-HttpDict = typing.TypedDict('HttpDict', {'url': str})
-TcpDict = typing.TypedDict('TcpDict', {'port': int})
-ExecDict = typing.TypedDict('ExecDict', {'command': str})
+HttpDict = typing.TypedDict('HttpDict',
+                            {'url': str,
+                             'headers': Dict[str, str]},
+                            total=False)
+TcpDict = typing.TypedDict('TcpDict',
+                           {'port': int,
+                            'host': str},
+                           total=False)
+ExecDict = typing.TypedDict('ExecDict',
+                            {'command': str,
+                             # see JujuVersion.supports_exec_service_context
+                             'service-context': str,
+                             'environment': Dict[str, str],
+                             'user-id': Optional[int],
+                             'user': str,
+                             'group-id': Optional[int],
+                             'group': str,
+                             'working-dir': str},
+                            total=False)
 
 CheckDict = typing.TypedDict('CheckDict',
                              {'override': str,
@@ -800,12 +818,14 @@ class Service:
         self.user_id = dct.get('user-id')
         self.group = dct.get('group', '')
         self.group_id = dct.get('group-id')
+        self.working_dir = dct.get('working-dir', '')
         self.on_success = dct.get('on-success', '')
         self.on_failure = dct.get('on-failure', '')
         self.on_check_failure = dict(dct.get('on-check-failure', {}))
         self.backoff_delay = dct.get('backoff-delay', '')
         self.backoff_factor = dct.get('backoff-factor')
         self.backoff_limit = dct.get('backoff-limit', '')
+        self.kill_delay = dct.get('kill-delay', '')
 
     def to_dict(self) -> 'ServiceDict':
         """Convert this service object to its dict representation."""
@@ -823,12 +843,14 @@ class Service:
             ('user-id', self.user_id),
             ('group', self.group),
             ('group-id', self.group_id),
+            ('working-dir', self.working_dir),
             ('on-success', self.on_success),
             ('on-failure', self.on_failure),
             ('on-check-failure', self.on_check_failure),
             ('backoff-delay', self.backoff_delay),
             ('backoff-factor', self.backoff_factor),
             ('backoff-limit', self.backoff_limit),
+            ('kill-delay', self.kill_delay),
         ]
         dct = {name: value for name, value in fields if value}
         return typing.cast('ServiceDict', dct)
@@ -2136,6 +2158,7 @@ class Client:
         self,
         command: List[str],
         *,
+        service_context: Optional[str] = None,
         environment: Optional[Dict[str, str]] = None,
         working_dir: Optional[str] = None,
         timeout: Optional[float] = None,
@@ -2157,6 +2180,7 @@ class Client:
         self,
         command: List[str],
         *,
+        service_context: Optional[str] = None,
         environment: Optional[Dict[str, str]] = None,
         working_dir: Optional[str] = None,
         timeout: Optional[float] = None,
@@ -2176,6 +2200,7 @@ class Client:
         self,
         command: List[str],
         *,
+        service_context: Optional[str] = None,
         environment: Optional[Dict[str, str]] = None,
         working_dir: Optional[str] = None,
         timeout: Optional[float] = None,
@@ -2260,6 +2285,11 @@ class Client:
         Args:
             command: Command to execute: the first item is the name (or path)
                 of the executable, the rest of the items are the arguments.
+            service_context: If specified, run the command in the context of
+                this service. Specifically, inherit its environment variables,
+                user/group settings, and working directory. The other exec
+                options will override the service context; ``environment``
+                will be merged on top of the service's.
             environment: Environment variables to pass to the process.
             working_dir: Working directory to run the command in. If not set,
                 Pebble uses the target user's $HOME directory (and if the user
@@ -2325,6 +2355,7 @@ class Client:
 
         body = {
             'command': command,
+            'service-context': service_context,
             'environment': environment or {},
             'working-dir': working_dir,
             'timeout': _format_timeout(timeout) if timeout is not None else None,
diff --git a/requirements-dev.txt b/requirements-dev.txt
index 320fe03..0b2ac7f 100644
--- a/requirements-dev.txt
+++ b/requirements-dev.txt
@@ -7,7 +7,7 @@ flake8-builtins==2.1.0
 pyproject-flake8==4.0.1
 pep8-naming==0.13.2
 pytest==7.2.1
-pyright==1.1.313
+pyright==1.1.316
 pytest-operator==0.23.0
 coverage[toml]==7.0.5
 typing_extensions==4.2.0


Test patch:
diff --git a/test/pebble_cli.py b/test/pebble_cli.py
index 05fc422..adeebde 100644
--- a/test/pebble_cli.py
+++ b/test/pebble_cli.py
@@ -64,6 +64,7 @@ def main():
     p.add_argument('name', help='check name(s) to filter on', nargs='*')
 
     p = subparsers.add_parser('exec', help='execute a command')
+    p.add_argument('--context', help='service context')
     p.add_argument('--env', help='environment variables to set', action='append',
                    metavar='KEY=VALUE')
     p.add_argument('--working-dir', help='working directory to run command in')
@@ -196,6 +197,7 @@ def main():
 
             process = client.exec(
                 args.exec_command,
+                service_context=args.context,
                 environment=environment,
                 working_dir=args.working_dir,
                 timeout=args.timeout,
diff --git a/test/test_jujuversion.py b/test/test_jujuversion.py
index 7e04230..a22cb05 100755
--- a/test/test_jujuversion.py
+++ b/test/test_jujuversion.py
@@ -80,6 +80,17 @@ class TestJujuVersion(unittest.TestCase):
         self.assertFalse(ops.JujuVersion('3.0.2').supports_open_port_on_k8s)
         self.assertFalse(ops.JujuVersion('2.9.30').supports_open_port_on_k8s)
 
+    def test_supports_exec_service_context(self):
+        self.assertFalse(ops.JujuVersion('2.9.30').supports_exec_service_context)
+        self.assertTrue(ops.JujuVersion('4.0.0').supports_exec_service_context)
+        self.assertFalse(ops.JujuVersion('3.0.0').supports_exec_service_context)
+        self.assertFalse(ops.JujuVersion('3.1.5').supports_exec_service_context)
+        self.assertTrue(ops.JujuVersion('3.1.6').supports_exec_service_context)
+        self.assertFalse(ops.JujuVersion('3.2.0').supports_exec_service_context)
+        self.assertTrue(ops.JujuVersion('3.2.2').supports_exec_service_context)
+        self.assertTrue(ops.JujuVersion('3.3.0').supports_exec_service_context)
+        self.assertTrue(ops.JujuVersion('3.4.0').supports_exec_service_context)
+
     def test_parsing_errors(self):
         invalid_versions = [
             "xyz",
diff --git a/test/test_model.py b/test/test_model.py
index 2ff3205..d9e86c9 100755
--- a/test/test_model.py
+++ b/test/test_model.py
@@ -24,6 +24,7 @@ import unittest
 from collections import OrderedDict
 from test.test_helpers import fake_script, fake_script_calls
 from textwrap import dedent
+from unittest.mock import patch
 
 import pytest
 
@@ -1777,10 +1778,12 @@ containers:
         self.assertEqual(len(cm.output), 1)
         self.assertRegex(cm.output[0], r'WARNING:ops.model:.*: api error!')
 
+    @patch('model.JujuVersion.from_environ', new=lambda: ops.model.JujuVersion('3.1.6'))
     def test_exec(self):
         self.pebble.responses.append('fake_exec_process')
         p = self.container.exec(
             ['echo', 'foo'],
+            service_context='srv1',
             environment={'K1': 'V1', 'K2': 'V2'},
             working_dir='WD',
             timeout=10.5,
@@ -1796,6 +1799,7 @@ containers:
         )
         self.assertEqual(self.pebble.requests, [
             ('exec', ['echo', 'foo'], dict(
+                service_context='srv1',
                 environment={'K1': 'V1', 'K2': 'V2'},
                 working_dir='WD',
                 timeout=10.5,
@@ -1812,6 +1816,11 @@ containers:
         ])
         self.assertEqual(p, 'fake_exec_process')
 
+    @patch('model.JujuVersion.from_environ', new=lambda: ops.model.JujuVersion('3.1.5'))
+    def test_exec_service_context_not_supported(self):
+        with self.assertRaises(RuntimeError):
+            self.container.exec(['foo'], service_context='srv1')
+
     def test_send_signal(self):
         with self.assertRaises(TypeError):
             self.container.send_signal('SIGHUP')
diff --git a/test/test_pebble.py b/test/test_pebble.py
index 586b5f0..6d172a2 100644
--- a/test/test_pebble.py
+++ b/test/test_pebble.py
@@ -681,12 +681,14 @@ class TestService(unittest.TestCase):
         self.assertIs(service.user_id, None)
         self.assertEqual(service.group, '')
         self.assertIs(service.group_id, None)
+        self.assertEqual(service.working_dir, '')
         self.assertEqual(service.on_success, '')
         self.assertEqual(service.on_failure, '')
         self.assertEqual(service.on_check_failure, {})
         self.assertEqual(service.backoff_delay, '')
         self.assertIs(service.backoff_factor, None)
         self.assertEqual(service.backoff_limit, '')
+        self.assertIs(service.kill_delay, '')
         self.assertEqual(service.to_dict(), {})
 
     def test_name_only(self):
@@ -711,12 +713,14 @@ class TestService(unittest.TestCase):
             'user-id': 1000,
             'group': 'staff',
             'group-id': 2000,
+            'working-dir': '/working/dir',
             'on-success': 'restart',
             'on-failure': 'ignore',
             'on-check-failure': {'chk1': 'halt'},
             'backoff-delay': '1s',
             'backoff-factor': 4,
             'backoff-limit': '10s',
+            'kill-delay': '420s',
         }
         s = pebble.Service('Name 2', d)
         self.assertEqual(s.name, 'Name 2')
@@ -732,12 +736,14 @@ class TestService(unittest.TestCase):
         self.assertEqual(s.user_id, 1000)
         self.assertEqual(s.group, 'staff')
         self.assertEqual(s.group_id, 2000)
+        self.assertEqual(s.working_dir, '/working/dir')
         self.assertEqual(s.on_success, 'restart')
         self.assertEqual(s.on_failure, 'ignore')
         self.assertEqual(s.on_check_failure, {'chk1': 'halt'})
         self.assertEqual(s.backoff_delay, '1s')
         self.assertEqual(s.backoff_factor, 4)
         self.assertEqual(s.backoff_limit, '10s')
+        self.assertEqual(s.kill_delay, '420s')
 
         self.assertEqual(s.to_dict(), d)
 
@@ -2639,10 +2645,11 @@ class TestExec(unittest.TestCase):
         return (stdio, stderr, control)
 
     def build_exec_data(
-            self, command, environment=None, working_dir=None, timeout=None,
+            self, command, service_context=None, environment=None, working_dir=None, timeout=None,
             user_id=None, user=None, group_id=None, group=None, combine_stderr=False):
         return {
             'command': command,
+            'service-context': service_context,
             'environment': environment or {},
             'working-dir': working_dir,
             'timeout': f'{timeout:.3f}s' if timeout is not None else None,