From 4f45fdc86d77c44c33834c5080893704f3948b58 Mon Sep 17 00:00:00 2001 From: Raymond Wiker Date: Tue, 26 Nov 2024 11:52:27 +0100 Subject: [PATCH] Implement .select on SearchContext, to specify what is returned as metadata from sumo. --- src/fmu/sumo/explorer/cache.py | 5 ++ .../sumo/explorer/objects/_search_context.py | 82 ++++++++++++------- tests/test_explorer.py | 16 ++++ 3 files changed, 73 insertions(+), 30 deletions(-) diff --git a/src/fmu/sumo/explorer/cache.py b/src/fmu/sumo/explorer/cache.py index 0204b63f..f76eacd8 100644 --- a/src/fmu/sumo/explorer/cache.py +++ b/src/fmu/sumo/explorer/cache.py @@ -33,3 +33,8 @@ def put(self, key, value): def has(self, key): return key in self.cache + + def clear(self): + with self.lock: + self.cache.clear() + self.access.clear() diff --git a/src/fmu/sumo/explorer/objects/_search_context.py b/src/fmu/sumo/explorer/objects/_search_context.py index 5afe0afd..7f96e282 100644 --- a/src/fmu/sumo/explorer/objects/_search_context.py +++ b/src/fmu/sumo/explorer/objects/_search_context.py @@ -10,14 +10,6 @@ from fmu.sumo.explorer.cache import LRUCache -_CASE_FIELDS = {"include": [], "exclude": []} - -_CHILD_FIELDS = { - "include": [], - "exclude": ["data.spec.columns", "fmu.realization.parameters"], -} - - def _gen_filter_none(): def _fn(value): return None, None @@ -289,6 +281,9 @@ def __init__( self._hits = None self._cache = LRUCache(capacity=200) self._length = None + self._select = { + "excludes": ["fmu.realization.parameters"], + } return @property @@ -492,7 +487,44 @@ async def getitem_async(self, index): uuid = self._hits[index] return await self.get_object_async(uuid) - def get_object(self, uuid: str, select: List[str] = None) -> Dict: + def select(self, sel): + """Specify what should be returned from elasticsearch. + Has the side effect of clearing the lru cache. + sel is either a single string value, a list of string value, + or a dictionary with keys "includes" and/or "excludes" and + the values are lists of strings. The string values are nested + property names. + + Args: + sel (str | List(str) | Dict(str, List[str]): select specification + Returns: + None + """ + + required = set(["class"]) + def extreq(lst): + if isinstance(lst, str): + lst = [lst] + return list(set(lst) | required) + if isinstance(sel, str): + self._select = extreq([sel]) + elif isinstance(sel, list): + self._select = extreq(sel) + elif isinstance(sel, dict): + inc = sel.get("includes") + exc = sel.get("excludes") + slct = {} + if inc is not None: + slct["includes"] = extreq(inc) + pass + if exc is not None: + slct["excludes"] = exc + pass + self._select = slct + pass + self._cache.clear() + + def get_object(self, uuid: str) -> Dict: """Get metadata object by uuid Args: @@ -507,11 +539,9 @@ def get_object(self, uuid: str, select: List[str] = None) -> Dict: query = { "query": {"ids": {"values": [uuid]}}, "size": 1, + "_source": self._select } - if select is not None: - query["_source"] = select - res = self._sumo.post("/search", json=query) hits = res.json()["hits"]["hits"] @@ -522,9 +552,7 @@ def get_object(self, uuid: str, select: List[str] = None) -> Dict: return self._to_sumo(obj) - async def get_object_async( - self, uuid: str, select: List[str] = None - ) -> Dict: + async def get_object_async(self, uuid: str) -> Dict: """Get metadata object by uuid Args: @@ -540,11 +568,9 @@ async def get_object_async( query = { "query": {"ids": {"values": [uuid]}}, "size": 1, + "_source": self._select } - if select is not None: - query["_source"] = select - res = await self._sumo.post_async("/search", json=query) hits = res.json()["hits"]["hits"] @@ -563,9 +589,7 @@ def _maybe_prefetch(self, index): uuids = [uuid for uuid in uuids if not self._cache.has(uuid)] hits = self.__search_all( {"ids": {"values": uuids}}, - select={ - "excludes": ["fmu.realization.parameters"], - }, + select=self._select, ) if len(hits) == 0: return @@ -581,9 +605,7 @@ async def _maybe_prefetch_async(self, index): uuids = [uuid for uuid in uuids if not self._cache.has(uuid)] hits = await self.__search_all_async( {"ids": {"values": uuids}}, - select={ - "excludes": ["fmu.realization.parameters"], - }, + select=self._select, ) if len(hits) == 0: return @@ -1131,7 +1153,7 @@ def get_surface_by_uuid(self, uuid: str): Returns: Surface: surface object """ - metadata = self.get_object(uuid, _CHILD_FIELDS) + metadata = self.get_object(uuid) return objects.Surface(self._sumo, metadata) async def get_surface_by_uuid_async(self, uuid: str): @@ -1143,7 +1165,7 @@ async def get_surface_by_uuid_async(self, uuid: str): Returns: Surface: surface object """ - metadata = await self.get_object_async(uuid, _CHILD_FIELDS) + metadata = await self.get_object_async(uuid) return objects.Surface(self._sumo, metadata) def get_polygons_by_uuid(self, uuid: str): @@ -1155,7 +1177,7 @@ def get_polygons_by_uuid(self, uuid: str): Returns: Polygons: polygons object """ - metadata = self.get_object(uuid, _CHILD_FIELDS) + metadata = self.get_object(uuid) return objects.Polygons(self._sumo, metadata) async def get_polygons_by_uuid_async(self, uuid: str): @@ -1167,7 +1189,7 @@ async def get_polygons_by_uuid_async(self, uuid: str): Returns: Polygons: polygons object """ - metadata = await self.get_object_async(uuid, _CHILD_FIELDS) + metadata = await self.get_object_async(uuid) return objects.Polygons(self._sumo, metadata) def get_table_by_uuid(self, uuid: str): @@ -1179,7 +1201,7 @@ def get_table_by_uuid(self, uuid: str): Returns: Table: table object """ - metadata = self.get_object(uuid, _CHILD_FIELDS) + metadata = self.get_object(uuid) return objects.Table(self._sumo, metadata) async def get_table_by_uuid_async(self, uuid: str): @@ -1191,7 +1213,7 @@ async def get_table_by_uuid_async(self, uuid: str): Returns: Table: table object """ - metadata = await self.get_object_async(uuid, _CHILD_FIELDS) + metadata = await self.get_object_async(uuid) return objects.Table(self._sumo, metadata) def _verify_aggregation_operation(self): diff --git a/tests/test_explorer.py b/tests/test_explorer.py index 01053ca5..293d6f3d 100644 --- a/tests/test_explorer.py +++ b/tests/test_explorer.py @@ -347,3 +347,19 @@ def test_grids_and_properties(explorer: Explorer): assert xtgridprop.nrow == gridpropspec["nrow"] assert xtgridprop.ncol == gridpropspec["ncol"] +def test_search_context_select(test_case: Case): + surfs = test_case.surfaces.filter(realization=True) + assert "_sumo" in surfs[0].metadata + surfs.select("fmu") + assert "_sumo" not in surfs[0].metadata + assert "fmu" in surfs[0].metadata + surfs.select(["fmu"]) + assert "_sumo" not in surfs[0].metadata + assert "fmu" in surfs[0].metadata + surfs.select({"excludes": ["fmu"]}) + assert "_sumo" in surfs[0].metadata + assert "fmu" not in surfs[0].metadata + surfs.select({"includes": ["_sumo"], "excludes": ["_sumo.timestamp"]}) + assert "_sumo" in surfs[0].metadata + assert "fmu" not in surfs[0].metadata + assert "timestamp" not in surfs[0].metadata["_sumo"]