Use DataBlockResponse model to validate block db r/w

Author: ml-evs

pydatalab/schemas/cell.json

@@ -223,7 +223,18 @@
         },
         "bokeh_plot_data": {
           "title": "Bokeh Plot Data",
-          "type": "string"
+          "anyOf": [
+            {
+              "type": "string"
+            },
+            {
+              "type": "object"
+            }
+          ]
+        },
+        "processed_data": {
+          "title": "Processed Data",
+          "type": "object"
         }
       },
       "required": [

pydatalab/schemas/equipment.json

@@ -187,7 +187,18 @@
         },
         "bokeh_plot_data": {
           "title": "Bokeh Plot Data",
-          "type": "string"
+          "anyOf": [
+            {
+              "type": "string"
+            },
+            {
+              "type": "object"
+            }
+          ]
+        },
+        "processed_data": {
+          "title": "Processed Data",
+          "type": "object"
         }
       },
       "required": [

pydatalab/schemas/sample.json

@@ -276,7 +276,18 @@
         },
         "bokeh_plot_data": {
           "title": "Bokeh Plot Data",
-          "type": "string"
+          "anyOf": [
+            {
+              "type": "string"
+            },
+            {
+              "type": "object"
+            }
+          ]
+        },
+        "processed_data": {
+          "title": "Processed Data",
+          "type": "object"
         }
       },
       "required": [

pydatalab/schemas/startingmaterial.json

@@ -329,7 +329,18 @@
         },
         "bokeh_plot_data": {
           "title": "Bokeh Plot Data",
-          "type": "string"
+          "anyOf": [
+            {
+              "type": "string"
+            },
+            {
+              "type": "object"
+            }
+          ]
+        },
+        "processed_data": {
+          "title": "Processed Data",
+          "type": "object"
         }
       },
       "required": [

pydatalab/src/pydatalab/blocks/base.py

@@ -6,6 +6,7 @@
 
 from pydatalab import __version__
 from pydatalab.logger import LOGGER
+from pydatalab.models.blocks import DataBlockResponse
 
 __all__ = ("generate_random_id", "DataBlock", "generate_js_callback_single_float_parameter")
 
@@ -100,6 +101,8 @@ def generate_random_id():
 class DataBlock:
     """Base class for a data block."""
 
+    block_db_model = DataBlockResponse
+
     name: str = "base"
     """The human-readable block name specifying which technique
     or file format it pertains to.
@@ -194,12 +197,9 @@ def __init__(
     def to_db(self):
         """returns a dictionary with the data for this
         block, ready to be input into mongodb"""
-        from bson import ObjectId
 
         LOGGER.debug("Casting block %s to database object.", self.__class__.__name__)
-        dct_for_db = {k: v for k, v in self.data.items() if k != "bokeh_plot_data"}
-        dct_for_db["file_id"] = ObjectId(dct_for_db["file_id"]) if "file_id" in dct_for_db else None
-        return dct_for_db
+        return self.block_db_model(**self.data).dict(exclude_unset=True)
 
     @classmethod
     def from_db(cls, block: dict):
@@ -339,6 +339,6 @@ def update_from_web(self, data: dict):
             "Updating block %s from web request",
             self.__class__.__name__,
         )
-        self.data.update(data)
+        self.data.update(self.block_db_model(**data).dict(exclude={}, exclude_unset=True))
 
         return self

pydatalab/src/pydatalab/models/blocks.py

@@ -1,4 +1,4 @@
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 from pydatalab.models.utils import JSON_ENCODERS, PyObjectId
 
@@ -32,13 +32,16 @@ class DataBlockResponse(BaseModel):
     file_ids: list[PyObjectId] | None = None
     """A list of file IDs associated with the block, if any."""
 
-    b64_encoded_image: dict[PyObjectId, str] | None = None
+    b64_encoded_image: dict[PyObjectId, str] | None = Field(exclude=True)
     """Any base64-encoded image data associated with the block, keyed by file_id, if any."""
 
-    bokeh_plot_data: str | None = None
+    bokeh_plot_data: str | dict | None = Field(exclude=True)
     """A JSON-encoded string containing the Bokeh plot data, if any."""
 
+    processed_data: dict | None = None
+    """Any processed data associated with the block, small enough to store."""
+
     class Config:
         allow_population_by_field_name = True
         json_encoders = JSON_ENCODERS
-        extra = "allow"
+        extra = "ignore"