From b0dd5fb35cd37af38993ffb838ffb6b0ad687e4f Mon Sep 17 00:00:00 2001
From: Yannan <73408381+YannanGao-gs@users.noreply.github.com>
Date: Fri, 16 Aug 2024 10:04:07 -0400
Subject: [PATCH] support search dataspace query using execution context
 (#3021)

---
 .../query/api/QueryStoreManager.java          | 36 +++++++++++++++++--
 .../query/api/TestQueryStoreManager.java      | 15 ++++++++
 2 files changed, 48 insertions(+), 3 deletions(-)

diff --git a/legend-engine-application-query/src/main/java/org/finos/legend/engine/application/query/api/QueryStoreManager.java b/legend-engine-application-query/src/main/java/org/finos/legend/engine/application/query/api/QueryStoreManager.java
index f5aa2396cbf..f59227252b4 100644
--- a/legend-engine-application-query/src/main/java/org/finos/legend/engine/application/query/api/QueryStoreManager.java
+++ b/legend-engine-application-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/legend-engine-application-query/src/test/java/org/finos/legend/engine/application/query/api/TestQueryStoreManager.java b/legend-engine-application-query/src/test/java/org/finos/legend/engine/application/query/api/TestQueryStoreManager.java
index 19cbe7fb63a..c83b5485541 100644
--- a/legend-engine-application-query/src/test/java/org/finos/legend/engine/application/query/api/TestQueryStoreManager.java
+++ b/legend-engine-application-query/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
     {