From 496d21de5468d4bb46af140a9c2b857ec7430943 Mon Sep 17 00:00:00 2001
From: Christian Kellner <christian@kellner.me>
Date: Tue, 1 Jun 2021 14:08:35 +0200
Subject: [PATCH] objectstore: sub-tree support for read_at

Add the ability to only read a sub-tree of a tree via `Object.read_at`.
Expose the functionality via the `Store{Server,Client}.read_tree_at`.
Extend the tests to check this new functionality.
---
 osbuild/objectstore.py       | 19 ++++++++++++++-----
 test/mod/test_objectstore.py | 24 ++++++++++++++++++++++++
 2 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/osbuild/objectstore.py b/osbuild/objectstore.py
index 6deb9d92..f7fcba9c 100644
--- a/osbuild/objectstore.py
+++ b/osbuild/objectstore.py
@@ -126,11 +126,18 @@ class Object:
                 yield path
 
     @contextlib.contextmanager
-    def read_at(self, target: PathLike) -> str:
+    def read_at(self, target: PathLike, path: str = "/") -> str:
+        """Read the object or a part of it at given location
+
+        Map the tree or a part of it specified via `path` at the
+        specified path `target`.
+        """
         self._check_writable()
         self._check_writer()
 
-        mount(self._path, target)
+        path = os.path.join(self._path, path.lstrip("/"))
+
+        mount(path, target)
         try:
             self._readers += 1
             yield target
@@ -410,6 +417,7 @@ class StoreServer(api.BaseAPI):
     def _read_tree_at(self, msg, sock):
         object_id = msg["object-id"]
         target = msg["target"]
+        subtree = msg["subtree"]
 
         obj = self.store.get(object_id)
         if not obj:
@@ -417,7 +425,7 @@ class StoreServer(api.BaseAPI):
             return
 
         try:
-            reader = obj.read_at(target)
+            reader = obj.read_at(target, subtree)
             path = self._stack.enter_context(reader)
         # pylint: disable=broad-except
         except Exception as e:
@@ -486,11 +494,12 @@ class StoreClient:
 
         return msg["path"]
 
-    def read_tree_at(self, object_id: str, target: str):
+    def read_tree_at(self, object_id: str, target: str, path="/"):
         msg = {
             "method": "read-tree-at",
             "object-id": object_id,
-            "target": os.fspath(target)
+            "target": os.fspath(target),
+            "subtree": os.fspath(path)
         }
 
         self.client.send(msg)
diff --git a/test/mod/test_objectstore.py b/test/mod/test_objectstore.py
index 91274d6f..addce9c1 100644
--- a/test/mod/test_objectstore.py
+++ b/test/mod/test_objectstore.py
@@ -337,10 +337,34 @@ class TestObjectStore(unittest.TestCase):
             txt = filepath.read_text()
             assert txt == "osbuild"
 
+            # check we can mount subtrees via `read_tree_at`
+
+            filemount = Path(tmpdir, "file")
+            filemount.touch()
+
+            path = client.read_tree_at("42", filemount, "/file.txt")
+            filepath = Path(path)
+            assert filepath.is_file()
+            txt = filepath.read_text()
+            assert txt == "osbuild"
+
+            dirmount = Path(tmpdir, "dir")
+            dirmount.mkdir()
+
+            path = client.read_tree_at("42", dirmount, "/directory")
+            dirpath = Path(path)
+            assert dirpath.is_dir()
+
+            # check proper exceptions are raised for non existent
+            # mount points and sub-trees
+
             with self.assertRaises(RuntimeError):
                 nonexistent = os.path.join(tmpdir, "nonexistent")
                 _ = client.read_tree_at("42", nonexistent)
 
+            with self.assertRaises(RuntimeError):
+                _ = client.read_tree_at("42", tmpdir, "/nonexistent")
+
             # The tree is being read via the client, should
             # not be able to write to it
             with self.assertRaises(ValueError):