Repository: dask/distributed

Issue:
`Client.get_metadata` using list of keys raises `KeyError` despite `default` if intermediate level does not exist
**Reproducer**

```python3
from distributed import Client
client = Client()
client.get_metadata(["x", "y"], False)
```

**Expected result**
`Client.get_metadata` should return the `default` on a `KeyError` at _any_ level. This is what the docstring suggests:
https://github.com/dask/distributed/blob/61353b70c177e09782a2911d26d3b03680e9184c/distributed/client.py#L4057-L4059

Note that this code returns `False` if we do `client.set_metadata(["x"], {})` before `client.get_metadata(["x", "y"], False)`, i.e., it only catches a `KeyError` on the last level:
https://github.com/dask/distributed/blob/61353b70c177e09782a2911d26d3b03680e9184c/distributed/scheduler.py#L6941-L6950

**Actual behavior**
The reproducer raises `KeyError: 'x'`

<details><summary>Full traceback</summary>
<p>

```
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
Input In [5], in <cell line: 1>()
----> 1 client.get_metadata(["x", "y"], False)

File ~/projects/dask/distributed/distributed/client.py:4067, in Client.get_metadata(self, keys, default)
   4065 if not isinstance(keys, (list, tuple)):
   4066     keys = (keys,)
-> 4067 return self.sync(self.scheduler.get_metadata, keys=keys, default=default)

File ~/projects/dask/distributed/distributed/utils.py:339, in SyncMethodMixin.sync(self, func, asynchronous, callback_timeout, *args, **kwargs)
    337     return future
    338 else:
--> 339     return sync(
    340         self.loop, func, *args, callback_timeout=callback_timeout, **kwargs
    341     )

File ~/projects/dask/distributed/distributed/utils.py:406, in sync(loop, func, callback_timeout, *args, **kwargs)
    404 if error:
    405     typ, exc, tb = error
--> 406     raise exc.with_traceback(tb)
    407 else:
    408     return result

File ~/projects/dask/distributed/distributed/utils.py:379, in sync.<locals>.f()
    377         future = asyncio.wait_for(future, callback_timeout)
    378     future = asyncio.ensure_future(future)
--> 379     result = yield future
    380 except Exception:
    381     error = sys.exc_info()

File /opt/homebrew/Caskroom/mambaforge/base/envs/dask-distributed-py3.9/lib/python3.9/site-packages/tornado/gen.py:762, in Runner.run(self)
    759 exc_info = None
    761 try:
--> 762     value = future.result()
    763 except Exception:
    764     exc_info = sys.exc_info()

File ~/projects/dask/distributed/distributed/core.py:1154, in PooledRPCCall.__getattr__.<locals>.send_recv_from_rpc(**kwargs)
   1152 prev_name, comm.name = comm.name, "ConnectionPool." + key
   1153 try:
-> 1154     return await send_recv(comm=comm, op=key, **kwargs)
   1155 finally:
   1156     self.pool.reuse(self.addr, comm)

File ~/projects/dask/distributed/distributed/core.py:944, in send_recv(comm, reply, serializers, deserializers, **kwargs)
    942     _, exc, tb = clean_exception(**response)
    943     assert exc
--> 944     raise exc.with_traceback(tb)
    945 else:
    946     raise Exception(response["exception_text"])

File ~/projects/dask/distributed/distributed/core.py:768, in _handle_comm()
    766     result = handler(comm, **msg)
    767 else:
--> 768     result = handler(**msg)
    769 if inspect.iscoroutine(result):
    770     result = await result

File ~/projects/dask/distributed/distributed/scheduler.py:6938, in get_metadata()
   6936 metadata = self.task_metadata
   6937 for key in keys[:-1]:
-> 6938     metadata = metadata[key]
   6939 try:
   6940     return metadata[keys[-1]]

KeyError: 'x'
```
</p>
</details>



Discussion:

Resolving patch:
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -6939,10 +6939,10 @@ def set_metadata(self, keys: list[str], value: object = None) -> None:
 
     def get_metadata(self, keys: list[str], default=no_default):
         metadata = self.task_metadata
-        for key in keys[:-1]:
-            metadata = metadata[key]
         try:
-            return metadata[keys[-1]]
+            for key in keys:
+                metadata = metadata[key]
+            return metadata
         except KeyError:
             if default != no_default:
                 return default


Test patch:
diff --git a/distributed/tests/test_client.py b/distributed/tests/test_client.py
--- a/distributed/tests/test_client.py
+++ b/distributed/tests/test_client.py
@@ -5583,10 +5583,29 @@ def fib(x):
 
 @gen_cluster(client=True)
 async def test_task_metadata(c, s, a, b):
+    with pytest.raises(KeyError):
+        await c.get_metadata("x")
+    with pytest.raises(KeyError):
+        await c.get_metadata(["x"])
+    result = await c.get_metadata("x", None)
+    assert result is None
+    result = await c.get_metadata(["x"], None)
+    assert result is None
+
+    with pytest.raises(KeyError):
+        await c.get_metadata(["x", "y"])
+    result = await c.get_metadata(["x", "y"], None)
+    assert result is None
+
     await c.set_metadata("x", 1)
     result = await c.get_metadata("x")
     assert result == 1
 
+    with pytest.raises(TypeError):
+        await c.get_metadata(["x", "y"])
+    with pytest.raises(TypeError):
+        await c.get_metadata(["x", "y"], None)
+
     future = c.submit(inc, 1)
     key = future.key
     await wait(future)