From fec9dcea970858fd4523ff47ccea93a22c8f038f Mon Sep 17 00:00:00 2001
From: Christian Kellner <christian@kellner.me>
Date: Fri, 25 Nov 2022 13:28:23 +0100
Subject: [PATCH] objectstore: implement a new metadata class

Implement a new class, nested inside `Object`, to read and write
metadata. It is indexed by a key and individual pieces of meta-
data are stored in separate files. Empty files are not created.
---
 osbuild/objectstore.py       | 72 ++++++++++++++++++++++++++++++++++++
 test/mod/test_objectstore.py | 55 +++++++++++++++++++++++++++
 2 files changed, 127 insertions(+)

diff --git a/osbuild/objectstore.py b/osbuild/objectstore.py
index 1fc82f06..3cdc0992 100644
--- a/osbuild/objectstore.py
+++ b/osbuild/objectstore.py
@@ -1,5 +1,6 @@
 import contextlib
 import enum
+import json
 import os
 import subprocess
 import tempfile
@@ -22,6 +23,77 @@ class Object:
         READ = 0
         WRITE = 1
 
+    class Metadata:
+        """store and retrieve metadata for an object"""
+
+        def __init__(self, base, folder: Optional[str] = None) -> None:
+            self.base = base
+            self.folder = folder
+            os.makedirs(self.path, exist_ok=True)
+
+        def _path_for_key(self, key) -> str:
+            assert key
+            name = f"{key}.json"
+            return os.path.join(self.path, name)
+
+        @property
+        def path(self):
+            if not self.folder:
+                return self.base
+            return os.path.join(self.base, self.folder)
+
+        @contextlib.contextmanager
+        def write(self, key):
+
+            tmp = tempfile.NamedTemporaryFile(
+                mode="w",
+                encoding="utf8",
+                dir=self.path,
+                prefix=".",
+                suffix=".tmp.json",
+                delete=True,
+            )
+
+            with tmp as f:
+                yield f
+
+                f.flush()
+
+                # if nothing was written to the file
+                si = os.stat(tmp.name)
+                if si.st_size == 0:
+                    return
+
+                dest = self._path_for_key(key)
+                # ensure it is proper json?
+                os.link(tmp.name, dest)
+
+        @contextlib.contextmanager
+        def read(self, key):
+            dest = self._path_for_key(key)
+            try:
+                with open(dest, "r", encoding="utf8") as f:
+                    yield f
+            except FileNotFoundError:
+                raise KeyError(f"No metadata for '{key}'") from None
+
+        def set(self, key: str, data):
+
+            if data is None:
+                return
+
+            with self.write(key) as f:
+                json.dump(data, f, indent=2)
+
+        def get(self, key: str):
+            with contextlib.suppress(KeyError):
+                with self.read(key) as f:
+                    return json.load(f)
+            return None
+
+        def __fspath__(self):
+            return self.path
+
     def __init__(self, store: "ObjectStore", uid: str, mode: Mode):
         self._mode = mode
         self._workdir = None
diff --git a/test/mod/test_objectstore.py b/test/mod/test_objectstore.py
index 604e387a..4ea3f4aa 100644
--- a/test/mod/test_objectstore.py
+++ b/test/mod/test_objectstore.py
@@ -199,6 +199,61 @@ class TestObjectStore(unittest.TestCase):
             assert store_path(store, "b", "A")
             assert store_path(store, "b", "B")
 
+    def test_metadata(self):
+
+        # test metadata object directly first
+        with tempfile.TemporaryDirectory() as tmp:
+            md = objectstore.Object.Metadata(tmp)
+
+            assert os.fspath(md) == tmp
+
+            with self.assertRaises(KeyError):
+                with md.read("a"):
+                    pass
+
+            # do not write anything to the file, it should not get stored
+            with md.write("a"):
+                pass
+
+            assert len(list(os.scandir(tmp))) == 0
+
+            # also we should not write anything if an exception was raised
+            with self.assertRaises(AssertionError):
+                with md.write("a") as f:
+                    f.write("{}")
+                    raise AssertionError
+
+            with md.write("a") as f:
+                f.write("{}")
+
+            assert len(list(os.scandir(tmp))) == 1
+
+            with md.read("a") as f:
+                assert f.read() == "{}"
+
+        data = {
+            "boolean": True,
+            "number": 42,
+            "string": "yes, please"
+        }
+
+        extra = {
+            "extra": "data"
+        }
+
+        with tempfile.TemporaryDirectory() as tmp:
+            md = objectstore.Object.Metadata(tmp)
+
+            d = md.get("a")
+            assert d is None
+
+            md.set("a", None)
+            with self.assertRaises(KeyError):
+                with md.read("a"):
+                    pass
+
+            md.set("a", data)
+            assert md.get("a") == data
     def test_host_tree(self):
         with objectstore.ObjectStore(self.store) as store:
             host = store.host_tree