Add support for writing to user-editable variables through the API

damnit/api.py

@@ -10,7 +10,8 @@
 import plotly.io as pio
 import xarray as xr
 
-from .backend.db import BlobTypes, DamnitDB
+from .kafka import UpdateProducer
+from .backend.db import BlobTypes, ReducedData, DamnitDB
 
 
 # This is a copy of damnit.ctxsupport.ctxrunner.DataType, purely so that we can
@@ -38,6 +39,13 @@ def find_proposal(propno):
     raise FileNotFoundError("Couldn't find proposal dir for {!r}".format(propno))
 
 
+# This variable is meant to be an instance of UpdateProducer, but we lazily
+# initialize it because creating the producer because takes ~100ms. Which isn't
+# very much, but it may otherwise be created hundreds of times if used in a
+# context file so it's better to avoid it where possible.
+UPDATE_PRODUCER = None
+
+
 class VariableData:
     """Represents a variable for a single run.
 
@@ -48,7 +56,7 @@ class VariableData:
     def __init__(self, name: str, title: str,
                  proposal: int, run: int,
                  h5_path: Path, data_format_version: int,
-                 db: DamnitDB, db_only: bool):
+                 db: DamnitDB, db_only: bool, missing: bool):
         self._name = name
         self._title = title
         self._proposal = proposal
@@ -57,6 +65,7 @@ def __init__(self, name: str, title: str,
         self._data_format_version = data_format_version
         self._db = db
         self._db_only = db_only
+        self._missing = missing
 
     @property
     def name(self) -> str:
@@ -145,6 +154,39 @@ def read(self):
                 # Otherwise, return a Numpy array
                 return group["data"][()]
 
+    def write(self, value, send_update=True):
+        """Write a value to a user-editable variable.
+
+        This may throw an exception if converting `value` to the type of the
+        editable variable fails, e.g. `db[100]["number"] = "foo"` will fail if
+        the `number` variable has a numeric type.
+
+        Args:
+            send_update (bool): Whether or not to send an update after
+                writing to the database. Don't use this unless you know what
+                you're doing, it may disappear in the future.
+        """
+        if not self._db_only:
+            raise RuntimeError(f"Cannot write to variable '{self.name}', it's not a user-editable variable.")
+
+        # Convert the input
+        user_variable = self._db.get_user_variables()[self.name]
+        variable_type = user_variable.get_type_class()
+        value = variable_type.to_db_value(value)
+        if value is None:
+            raise ValueError(f"Forbidden conversion of value '{value!r}' to type '{variable_type.py_type}'")
+
+        # Write to the database
+        self._db.set_variable(self.proposal, self.run, self.name, ReducedData(value))
+
+        if send_update:
+            global UPDATE_PRODUCER
+            if UPDATE_PRODUCER is None:
+                UPDATE_PRODUCER = UpdateProducer(None)
+
+            UPDATE_PRODUCER.variable_set(self.name, self.title, variable_type.type_name,
+                                      flush=True, topic=self._db.kafka_topic)
+
     def summary(self):
         """Read the summary data for a variable.
 
@@ -204,21 +246,40 @@ def file(self) -> Path:
         """The path to the HDF5 file for the run."""
         return self._h5_path
 
-    def __getitem__(self, name):
+    def _get_variable(self, name):
         key_locs = self._key_locations()
         names_to_titles = self._var_titles()
         titles_to_names = { title: name for name, title in names_to_titles.items() }
 
-        if name not in key_locs and name not in titles_to_names:
-            raise KeyError(f"Variable data for '{name!r}' not found for p{self.proposal}, r{self.run}")
-
         if name in titles_to_names:
             name = titles_to_names[name]
 
-        return VariableData(name, names_to_titles[name],
+        missing = name not in key_locs
+        user_variables = self._db.get_user_variables()
+        if missing and name in user_variables:
+            key_locs[name] = True
+        elif missing and name not in user_variables:
+            raise KeyError(f"Variable data for '{name!r}' not found for p{self.proposal}, r{self.run}")
+
+        return VariableData(name, names_to_titles.get(name),
                             self.proposal, self.run,
                             self._h5_path, self._data_format_version,
-                            self._db, key_locs[name])
+                            self._db, key_locs[name],
+                            missing)
+
+    def __getitem__(self, name):
+        variable = self._get_variable(name)
+        if variable._missing:
+            raise KeyError(f"Variable data for '{name!r}' not found for p{self.proposal}, r{self.run}")
+
+        return variable
+
+    def __setitem__(self, name, value):
+        variable = self._get_variable(name)
+
+        # The environment variable is basically only useful for tests
+        send_update = bool(int(os.environ.get("DAMNIT_SEND_UPDATE", 1)))
+        variable.write(value, send_update)
 
     def _key_locations(self):
         # Read keys from the HDF5 file

damnit/backend/db.py

@@ -209,7 +209,7 @@ def get_user_variables(self):
                 attributes=rr["attributes"],
             )
             user_variables[var_name] = new_var
-        log.debug("Loaded %d user variables", len(user_variables))
+
         return user_variables
 
     def update_computed_variables(self, vars: dict):

damnit/backend/user_variables.py

@@ -10,17 +10,23 @@ class ValueType:
 
     examples = None
 
+    py_type = None
+
     def __str__(self):
         return self.type_name
 
     @classmethod
     def parse(cls, input: str):
-        return input
+        return cls.py_type(input)
 
     @classmethod
     def from_db_value(cls, value):
         return value
 
+    @classmethod
+    def to_db_value(cls, value):
+        return cls.py_type(value)
+
 
 class BooleanValueType(ValueType):
     type_name = "boolean"
@@ -29,6 +35,8 @@ class BooleanValueType(ValueType):
 
     examples = ["True", "T", "true", "1", "False", "F", "f", "0"]
 
+    py_type = bool
+
     _valid_values = {
         "true": True,
         "yes": True,
@@ -61,6 +69,10 @@ def from_db_value(cls, value):
             return None
         return bool(value)
 
+    @classmethod
+    def to_db_value(cls, value):
+        return value if isinstance(value, bool) else None
+
 
 class IntegerValueType(ValueType):
     type_name = "integer"
@@ -69,9 +81,7 @@ class IntegerValueType(ValueType):
 
     examples = ["-7", "-2", "0", "10", "34"]
 
-    @classmethod
-    def parse(cls, input: str):
-        return int(input)
+    py_type = int
 
 class NumberValueType(ValueType):
     type_name = "number"
@@ -80,10 +90,7 @@ class NumberValueType(ValueType):
 
     examples = ["-34.1e10", "-7.1", "-4", "0.0", "3.141592653589793", "85.4E7"]
 
-    @classmethod
-    def parse(cls, input: str):
-        return float(input)
-
+    py_type = float
 
 class StringValueType(ValueType):
     type_name = "string"
@@ -92,6 +99,12 @@ class StringValueType(ValueType):
 
     examples = ["Broken", "Dark frame", "test_frame"]
 
+    py_type = str
+
+    @classmethod
+    def to_db_value(cls, value):
+        return value if isinstance(value, str) else None
+
 
 value_types_by_name = {tt.type_name: tt for tt in [
     BooleanValueType(), IntegerValueType(), NumberValueType(), StringValueType()