Define validate shape of descriptor

databroker/mongo_normalized.py

@@ -514,6 +514,7 @@ def __init__(
         event_collection,
         root_map,
         sub_dict,
+        validate_shape,
     ):
         self._run = run
         self._stream_name = stream_name
@@ -522,6 +523,7 @@ def __init__(
         self._event_collection = event_collection
         self._sub_dict = sub_dict
         self.root_map = root_map
+        self.validate_shape = validate_shape
 
         # metadata should look like
         # {
@@ -851,7 +853,7 @@ def populate_columns(keys, min_seq_num, max_seq_num):
                 if expected_shape and (not is_external):
                     validated_column = list(
                         map(
-                            lambda item: _validate_shape(
+                            lambda item: self.validate_shape(
                                 key, numpy.asarray(item), expected_shape
                             ),
                             result[key],
@@ -936,7 +938,7 @@ def populate_columns(keys, min_seq_num, max_seq_num):
                         last_datum_id=None,
                     )
                     filled_data = filled_mock_event["data"][key]
-                    validated_filled_data = _validate_shape(
+                    validated_filled_data = self.validate_shape(
                         key, filled_data, expected_shape
                     )
                     filled_column.append(validated_filled_data)
@@ -1047,6 +1049,7 @@ def from_uri(
         access_policy=None,
         cache_ttl_complete=60,  # seconds
         cache_ttl_partial=2,  # seconds
+        validate_shape=None
     ):
         """
         Create a MongoAdapter from MongoDB with the "normalized" (original) layout.
@@ -1094,6 +1097,9 @@ def from_uri(
         cache_ttl_complete : float
             Time (in seconds) to cache a *complete* BlueskyRun before checking
             the database for updates. Default 60.
+        validate_shape: func
+            function that will be used to validate that the shape of the data matches
+            the shape in the descriptor document
         """
         metadatastore_db = _get_database(uri)
         if asset_registry_uri is None:
@@ -1122,6 +1128,7 @@ def from_uri(
             cache_of_partial_bluesky_runs=cache_of_partial_bluesky_runs,
             metadata=metadata,
             access_policy=access_policy,
+            validate_shape=validate_shape,
         )
 
     @classmethod
@@ -1135,6 +1142,7 @@ def from_mongomock(
         access_policy=None,
         cache_ttl_complete=60,  # seconds
         cache_ttl_partial=2,  # seconds
+        validate_shape=None
     ):
         """
         Create a transient MongoAdapter from backed by "mongomock".
@@ -1178,6 +1186,9 @@ def from_mongomock(
         cache_ttl_complete : float
             Time (in seconds) to cache a *complete* BlueskyRun before checking
             the database for updates. Default 60.
+        validate_shape: func
+            function that will be used to validate that the shape of the data matches
+            the shape in the descriptor document
         """
         import mongomock
 
@@ -1205,6 +1216,7 @@ def from_mongomock(
             cache_of_partial_bluesky_runs=cache_of_partial_bluesky_runs,
             metadata=metadata,
             access_policy=access_policy,
+            validate_shape=validate_shape,
         )
 
     def __init__(
@@ -1220,6 +1232,7 @@ def __init__(
         queries=None,
         sorting=None,
         access_policy=None,
+        validate_shape=None,
     ):
         "This is not user-facing. Use MongoAdapter.from_uri."
         self._run_start_collection = metadatastore_db.get_collection("run_start")
@@ -1249,6 +1262,11 @@ def __init__(
         self._sorting = sorting
         self.access_policy = access_policy
         self._serializer = None
+        if validate_shape is None:
+            validate_shape = default_validate_shape
+        elif isinstance(validate_shape, str):
+            validate_shape = import_object(validate_shape)
+        self.validate_shape = validate_shape
         super().__init__()
 
     @property
@@ -1441,6 +1459,7 @@ def _build_event_stream(self, *, run_start_uid, stream_name, is_complete):
                     event_collection=self._event_collection,
                     root_map=self.root_map,
                     sub_dict="data",
+                    validate_shape=self.validate_shape,
                 ),
                 "timestamps": lambda: DatasetFromDocuments(
                     run=run,
@@ -1450,6 +1469,7 @@ def _build_event_stream(self, *, run_start_uid, stream_name, is_complete):
                     event_collection=self._event_collection,
                     root_map=self.root_map,
                     sub_dict="timestamps",
+                    validate_shape=self.validate_shape,
                 ),
                 "config": lambda: Config(
                     OneShotCachedMap(
@@ -2095,7 +2115,7 @@ class BadShapeMetadata(Exception):
     pass
 
 
-def _validate_shape(key, data, expected_shape):
+def default_validate_shape(key, data, expected_shape):
     """
     Check that data.shape == expected.shape.
 

databroker/tests/test_validate_shape.py

@@ -0,0 +1,31 @@
+from bluesky import RunEngine
+from bluesky.plans import count
+from ophyd.sim import img
+from tiled.client import Context, from_context
+from tiled.server.app import build_app
+
+from ..mongo_normalized import MongoAdapter
+
+
+def test_validate_shape(tmpdir):
+    # custom_validate_shape will mutate this to show it has been called
+    shapes = []
+
+    def custom_validate_shape(key, data, expected_shape):
+        shapes.append(expected_shape)
+        return data
+
+    adapter = MongoAdapter.from_mongomock(validate_shape=custom_validate_shape)
+
+    with Context.from_app(build_app(adapter), token_cache=tmpdir) as context:
+        client = from_context(context)
+
+        def post_document(name, doc):
+            client.post_document(name, doc)
+
+        RE = RunEngine()
+        RE.subscribe(post_document)
+        (uid,) = RE(count([img]))
+        assert not shapes
+        client[uid]["primary"]["data"]["img"][:]
+        assert shapes