Use DataBlockResponse model to validate block db r/w

Author: ml-evs

pydatalab/schemas/cell.json

@@ -214,6 +214,20 @@
             "type": "string"
           }
         },
+        "errors": {
+          "title": "Errors",
+          "type": "array",
+          "items": {
+            "type": "string"
+          }
+        },
+        "warnings": {
+          "title": "Warnings",
+          "type": "array",
+          "items": {
+            "type": "string"
+          }
+        },
         "b64_encoded_image": {
           "title": "B64 Encoded Image",
           "type": "object",
@@ -223,7 +237,11 @@
         },
         "bokeh_plot_data": {
           "title": "Bokeh Plot Data",
-          "type": "string"
+          "type": "object"
+        },
+        "processed_data": {
+          "title": "Processed Data",
+          "type": "object"
         }
       },
       "required": [

pydatalab/schemas/equipment.json

@@ -178,6 +178,20 @@
             "type": "string"
           }
         },
+        "errors": {
+          "title": "Errors",
+          "type": "array",
+          "items": {
+            "type": "string"
+          }
+        },
+        "warnings": {
+          "title": "Warnings",
+          "type": "array",
+          "items": {
+            "type": "string"
+          }
+        },
         "b64_encoded_image": {
           "title": "B64 Encoded Image",
           "type": "object",
@@ -187,7 +201,11 @@
         },
         "bokeh_plot_data": {
           "title": "Bokeh Plot Data",
-          "type": "string"
+          "type": "object"
+        },
+        "processed_data": {
+          "title": "Processed Data",
+          "type": "object"
         }
       },
       "required": [

pydatalab/schemas/sample.json

@@ -267,6 +267,20 @@
             "type": "string"
           }
         },
+        "errors": {
+          "title": "Errors",
+          "type": "array",
+          "items": {
+            "type": "string"
+          }
+        },
+        "warnings": {
+          "title": "Warnings",
+          "type": "array",
+          "items": {
+            "type": "string"
+          }
+        },
         "b64_encoded_image": {
           "title": "B64 Encoded Image",
           "type": "object",
@@ -276,7 +290,11 @@
         },
         "bokeh_plot_data": {
           "title": "Bokeh Plot Data",
-          "type": "string"
+          "type": "object"
+        },
+        "processed_data": {
+          "title": "Processed Data",
+          "type": "object"
         }
       },
       "required": [

pydatalab/schemas/startingmaterial.json

@@ -320,6 +320,20 @@
             "type": "string"
           }
         },
+        "errors": {
+          "title": "Errors",
+          "type": "array",
+          "items": {
+            "type": "string"
+          }
+        },
+        "warnings": {
+          "title": "Warnings",
+          "type": "array",
+          "items": {
+            "type": "string"
+          }
+        },
         "b64_encoded_image": {
           "title": "B64 Encoded Image",
           "type": "object",
@@ -329,7 +343,11 @@
         },
         "bokeh_plot_data": {
           "title": "Bokeh Plot Data",
-          "type": "string"
+          "type": "object"
+        },
+        "processed_data": {
+          "title": "Processed Data",
+          "type": "object"
         }
       },
       "required": [

pydatalab/src/pydatalab/blocks/base.py

@@ -1,11 +1,13 @@
 import functools
+import json
 import random
 import warnings
 from collections.abc import Callable, Sequence
 from typing import Any
 
 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 +102,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.
@@ -191,15 +195,12 @@ def __init__(
             collection_id,
         )
 
-    def to_db(self):
+    def to_db(self) -> dict:
         """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):
@@ -251,7 +252,9 @@ def to_web(self) -> dict[str, Any]:
         else:
             self.data.pop("warnings", None)
 
-        return self.data
+        model = self.block_db_model(**self.data)
+        serialized = json.loads(model.json(exclude_unset=True, exclude_none=True))
+        return serialized
 
     def process_events(self, events: list[dict] | dict):
         """Handle any supported events passed to the block."""
@@ -339,6 +342,10 @@ 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={"bokeh_plot_data", "b64_processed_image"}, exclude_unset=True
+            )
+        )
 
         return self

pydatalab/src/pydatalab/models/blocks.py

@@ -32,13 +32,22 @@ 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
+    errors: list[str] | None = None
+    """Any errors that occurred during block processing."""
+
+    warnings: list[str] | None = None
+    """Any warnings that occurred during block processing."""
+
+    b64_encoded_image: dict[PyObjectId, str] | None
     """Any base64-encoded image data associated with the block, keyed by file_id, if any."""
 
-    bokeh_plot_data: str | None = None
+    bokeh_plot_data: dict | None
     """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"

pydatalab/tests/server/test_blocks.py

@@ -75,7 +75,7 @@ def test_invalid_block_type(admin_client, default_sample_dict):
         },
     )
 
-    assert response.status_code == 400
+    assert response.status_code == 501
     assert "Invalid block type" in response.json["message"]
 
 
@@ -259,7 +259,9 @@ def test_xrd_block_lifecycle(admin_client, default_sample_dict):
     assert "bokeh_plot_data" in web_block
     assert "processed_data" in web_block
     assert "peak_data" in web_block["processed_data"]
-    assert web_block.get("errors", []) == []
+    assert "file_id" in web_block
+    assert web_block["file_id"] == file_id
+    assert web_block.get("errors") is None
 
     block = XRDBlock.from_web(web_block)
     db = block.to_db()
@@ -306,7 +308,7 @@ def test_comment_block_manipulation(admin_client, default_sample_dict, database)
     assert response.json["new_block_data"]["blocktype"] == block_type
     assert response.json["new_block_data"]["freeform_comment"] == "This is a test comment block."
     assert response.json["new_block_data"]["title"] == "Test Comment Block"
-    assert "errors" not in response.json["new_block_data"]
+    assert response.json["new_block_data"].get("errors") is None
 
     # Check that this result was actually stored
     response = admin_client.get(f"/get-item-data/{sample_id}")
@@ -319,7 +321,7 @@ def test_comment_block_manipulation(admin_client, default_sample_dict, database)
     assert "errors" not in response.json["item_data"]["blocks_obj"][block_id]
 
     # Try to add some bad data
-    block_data["bokeh_plot_data"] = '{"bokeh": "json"}'
+    block_data["bokeh_plot_data"] = {"bokeh": "json"}
     block_data["random_new_key"] = "test new key"
     response = admin_client.post("/update-block/", json={"block_data": block_data})
     assert response.status_code == 200
@@ -419,7 +421,7 @@ def test_create_sample_with_example_files(
     if block_type == "xrd":
         assert response.json["new_block_data"]["processed_data"]["peak_data"] is not None
 
-    response = admin_client.get(f"/get-item-data/{sample_id}?load_blocks=1")
+    response = admin_client.get(f"/get-item-data/{sample_id}")
     assert response.status_code == 200
     assert response.json["status"] == "success"