Add support for exec "service-context" (#957)

This is Pebble PR canonical/pebble#246

Also add working_dir field to pebble.Service and flesh out types for
HttpDict, TcpDict, and ExecDict.

Bump up to latest Pyright version while we're at it

It will be supported on Juju 3.1.6+ and 3.2.1+ and 3.3+

ops/jujuversion.py

@@ -128,3 +128,15 @@ def supports_open_port_on_k8s(self) -> bool:
         """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

ops/model.py

@@ -2360,6 +2360,7 @@ def exec(  # noqa
         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 @@ def exec(  # noqa
         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 @@ def exec(
         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 @@ def exec(
         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,

ops/pebble.py

@@ -79,6 +79,7 @@
                                 '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],
@@ -88,9 +89,25 @@
                                 },
                                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,6 +817,7 @@ def __init__(self, name: str, raw: Optional['ServiceDict'] = None):
         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', {}))
@@ -823,6 +841,7 @@ def to_dict(self) -> 'ServiceDict':
             ('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),
@@ -2136,6 +2155,7 @@ def exec(  # noqa
         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 +2177,7 @@ def exec(  # noqa
         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 +2197,7 @@ def exec(
         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 +2282,11 @@ def exec(
         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 +2352,7 @@ def exec(
 
         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,

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/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,

test/test_jujuversion.py

@@ -80,6 +80,17 @@ def test_supports_open_port_on_k8s(self):
         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",

test/test_model.py

@@ -24,6 +24,7 @@
 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 @@ def raise_error():
         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 @@ def test_exec(self):
         )
         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 @@ def test_exec(self):
         ])
         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')

test/test_pebble.py

@@ -681,6 +681,7 @@ def _assert_empty(self, service, name):
         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, {})
@@ -711,6 +712,7 @@ def test_dict(self):
             'user-id': 1000,
             'group': 'staff',
             'group-id': 2000,
+            'working-dir': '/working/dir',
             'on-success': 'restart',
             'on-failure': 'ignore',
             'on-check-failure': {'chk1': 'halt'},
@@ -732,6 +734,7 @@ def test_dict(self):
         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'})
@@ -2639,10 +2642,11 @@ def add_responses(self, change_id, exit_code, change_err=None):
         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,