support search dataspace query using execution context (#3021)

finos · Aug 16, 2024 · b0dd5fb · b0dd5fb
1 parent fc9e98b
commit b0dd5fb
Show file tree

Hide file tree

Showing 2 changed files with 48 additions and 3 deletions.
diff --git a/...-query/src/main/java/org/finos/legend/engine/application/query/api/QueryStoreManager.java b/...-query/src/main/java/org/finos/legend/engine/application/query/api/QueryStoreManager.java
@@ -43,8 +43,10 @@
 import java.time.Instant;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.List;
 import java.util.regex.Pattern;
+import java.util.stream.Collectors;
 
 public class QueryStoreManager
 {
@@ -242,9 +244,37 @@ public List<Query> searchQueries(QuerySearchSpecification searchSpecification, S
         }
         if (searchSpecification.taggedValues != null && !searchSpecification.taggedValues.isEmpty())
         {
-            List taggedValueFilters = ListIterate.collect(searchSpecification.taggedValues, taggedValue ->
-                    Filters.and(Filters.eq("taggedValues.tag.profile", taggedValue.tag.profile), Filters.eq("taggedValues.tag.value", taggedValue.tag.value), Filters.eq("taggedValues.value", taggedValue.value)));
-            filters.add(searchSpecification.combineTaggedValuesCondition != null && searchSpecification.combineTaggedValuesCondition ? Filters.and(taggedValueFilters) : Filters.or(taggedValueFilters));
+            List<Bson> taggedValueFilters = ListIterate.collect(searchSpecification.taggedValues, taggedValue ->
+                    Filters.and(
+                            Filters.eq("taggedValues.tag.profile", taggedValue.tag.profile),
+                            Filters.eq("taggedValues.tag.value", taggedValue.tag.value),
+                            Filters.eq("taggedValues.value", taggedValue.value)
+                    )
+            );
+            Bson taggedValuesFilter = searchSpecification.combineTaggedValuesCondition != null && searchSpecification.combineTaggedValuesCondition
+                    ? Filters.and(taggedValueFilters)
+                    : Filters.or(taggedValueFilters);
+
+            List<String> dataspaceTaggedValues = searchSpecification.taggedValues != null
+                    ? searchSpecification.taggedValues.stream()
+                    .filter(taggedValue -> QUERY_PROFILE_PATH.equals(taggedValue.tag.profile) &&
+                            QUERY_PROFILE_TAG_DATA_SPACE.equals(taggedValue.tag.value))
+                    .map(taggedValue ->  taggedValue.value)
+                    .collect(Collectors.toList())
+                    : Collections.emptyList();
+
+            if (!dataspaceTaggedValues.isEmpty())
+            {
+                Bson executionContextFilter = Filters.and(
+                        Filters.eq("executionContext._type", "dataSpaceExecutionContext"),
+                        Filters.in("executionContext.dataSpacePath", dataspaceTaggedValues)
+                );
+                filters.add(Filters.or(taggedValuesFilter, executionContextFilter));
+            }
+            else
+            {
+                filters.add(taggedValuesFilter);
+            }
         }
         if (searchSpecification.stereotypes != null && !searchSpecification.stereotypes.isEmpty())
         {

diff --git a/...ry/src/test/java/org/finos/legend/engine/application/query/api/TestQueryStoreManager.java b/...ry/src/test/java/org/finos/legend/engine/application/query/api/TestQueryStoreManager.java
@@ -673,6 +673,21 @@ public void testGetQueriesWithTaggedValues() throws Exception
         Assert.assertEquals(3, queryStoreManager.searchQueries(new TestQuerySearchSpecificationBuilder().withTaggedValues(Lists.fixedSize.of(taggedValue1, taggedValue2, taggedValue3)).build(), currentUser).size());
     }
 
+    @Test
+    public void testGetDataSpaceQueriesWithExecutionContext() throws Exception
+    {
+        String currentUser = "testUser";
+        String dataspacePath = "test::Dataspace";
+        TaggedValue taggedValue = createTestTaggedValue("meta::pure::profiles::query", "dataSpace", dataspacePath);
+        Query testQuery1 = TestQueryBuilder.create("1", "query1", currentUser).withDataSpaceExecution(dataspacePath).build();
+        Query testQuery2 = TestQueryBuilder.create("2", "query2", currentUser).withTaggedValues(Lists.fixedSize.of(taggedValue)).withDataSpaceExecution(dataspacePath).build();
+        queryStoreManager.createQuery(testQuery1, currentUser);
+        queryStoreManager.createQuery(testQuery2, currentUser);
+
+        Assert.assertEquals(2, queryStoreManager.searchQueries(new TestQuerySearchSpecificationBuilder().build(), currentUser).size());
+        Assert.assertEquals(2, queryStoreManager.searchQueries(new TestQuerySearchSpecificationBuilder().withTaggedValues(Lists.fixedSize.of(taggedValue)).build(), currentUser).size());
+    }
+
     @Test
     public void testGetQueriesWithGridConfigs() throws Exception
     {