diff --git a/databroker/experimental/server_ext.py b/databroker/experimental/server_ext.py index f085bc5ba..f1e16b76a 100644 --- a/databroker/experimental/server_ext.py +++ b/databroker/experimental/server_ext.py @@ -19,6 +19,9 @@ Comparison, Eq, FullText, + In, + NotEq, + NotIn, Operator, Regex, ) @@ -435,6 +438,10 @@ def eq(query, catalog): return catalog.apply_mongo_query({f"metadata.{query.key}": query.value}) +def noteq(query, catalog): + return catalog.apply_mongo_query({f"metadata.{query.key}": {"$ne": query.value}}) + + def regex(query, catalog): options = "" if query.case_sensitive else "i" return catalog.apply_mongo_query( @@ -442,6 +449,18 @@ def regex(query, catalog): ) +def _in(query, catalog): + if not isinstance(query.value, list): + query.value = [query.value] + return catalog.apply_mongo_query({f"metadata.{query.key}": {"$in": query.value}}) + + +def notin(query, catalog): + if not isinstance(query.value, list): + query.value = [query.value] + return catalog.apply_mongo_query({f"metadata.{query.key}": {"$nin": query.value}}) + + def full_text_search(query, catalog): # First if this catalog is backed by mongomock, which does not support $text queries. # Avoid importing mongomock if it is not already imported. @@ -463,5 +482,8 @@ def full_text_search(query, catalog): MongoAdapter.register_query(Contains, contains) MongoAdapter.register_query(Comparison, comparison) MongoAdapter.register_query(Eq, eq) -MongoAdapter.register_query(FullText, full_text_search) +MongoAdapter.register_query(NotEq, noteq) MongoAdapter.register_query(Regex, regex) +MongoAdapter.register_query(In, _in) +MongoAdapter.register_query(NotIn, notin) +MongoAdapter.register_query(FullText, full_text_search) diff --git a/databroker/tests/test_experimental.py b/databroker/tests/test_experimental.py index e61e2a0f1..1fa21cd96 100644 --- a/databroker/tests/test_experimental.py +++ b/databroker/tests/test_experimental.py @@ -1,5 +1,15 @@ from tiled.client import from_tree -from tiled.queries import Contains, Comparison, Eq, FullText, Key, Regex +from tiled.queries import ( + Contains, + Comparison, + Eq, + FullText, + In, + Key, + NotEq, + NotIn, + Regex, +) from ..experimental.server_ext import MongoAdapter @@ -98,3 +108,18 @@ def test_queries(tmpdir): numpy.testing.assert_equal( test5.values()[0].read(), test5.values()[0].metadata["number"] * numpy.ones(10) ) + test6 = client.search(NotEq("letter", "a")) + # The first result should not be "a" + assert test6.values()[0].metadata["letter"] != "a" + + test7 = client.search(In("letter", ["a", "b"])) + numpy.testing.assert_equal( + test7.values()[0].read(), test7.values()[0].metadata["number"] * numpy.ones(10) + ) + numpy.testing.assert_equal( + test7.values()[1].read(), test7.values()[1].metadata["number"] * numpy.ones(10) + ) + + test8 = client.search(NotIn("letter", ["a"])) + # The first result should not be "a" + assert test8.values()[0].metadata["letter"] != "a"