Updated the composite query logic to handle the validateQuery mechani…

…sm (#2744) * Updated the composite query logic to handle the validateQuery mechanism * Updated the filtered query logic and delegating query logic also
NationalSecurityAgency · Feb 10, 2025 · 3c93f74 · 3c93f74
1 parent 17a7625
commit 3c93f74
Show file tree

Hide file tree

Showing 4 changed files with 133 additions and 0 deletions.
diff --git a/core/query/src/main/java/datawave/core/query/logic/DelegatingQueryLogic.java b/core/query/src/main/java/datawave/core/query/logic/DelegatingQueryLogic.java
@@ -8,6 +8,7 @@
 
 import org.apache.accumulo.core.client.AccumuloClient;
 import org.apache.accumulo.core.security.Authorizations;
+import org.apache.commons.collections4.Transformer;
 import org.apache.commons.collections4.iterators.TransformIterator;
 
 import datawave.audit.SelectorExtractor;
@@ -21,6 +22,7 @@
 import datawave.webservice.common.connection.AccumuloClientConfiguration;
 import datawave.webservice.query.exception.QueryException;
 import datawave.webservice.query.result.event.ResponseObjectFactory;
+import datawave.webservice.result.QueryValidationResponse;
 
 /**
  * A delegating query logic that simply passes through to a delegate query logic. Intended to simplify extending classes.
@@ -91,6 +93,11 @@ public final QueryLogicTransformer getEnrichedTransformer(Query settings) {
         return delegate.getEnrichedTransformer(settings);
     }
 
+    @Override
+    public ResultPostprocessor getResultPostprocessor(GenericQueryConfiguration config) {
+        return delegate.getResultPostprocessor(config);
+    }
+
     @Override
     public String getResponseClass(Query query) throws QueryException {
         return delegate.getResponseClass(query);
@@ -368,6 +375,16 @@ public void setServerUser(ProxiedUserDetails serverUser) {
         delegate.setServerUser(serverUser);
     }
 
+    @Override
+    public Object validateQuery(AccumuloClient client, Query query, Set<Authorizations> auths) throws Exception {
+        return delegate.validateQuery(client, query, auths);
+    }
+
+    @Override
+    public Transformer<Object,QueryValidationResponse> getQueryValidationResponseTransformer() {
+        return delegate.getQueryValidationResponseTransformer();
+    }
+
     @Override
     public UserOperations getUserOperations() {
         return delegate.getUserOperations();

diff --git a/core/query/src/main/java/datawave/core/query/logic/composite/CompositeQueryLogic.java b/core/query/src/main/java/datawave/core/query/logic/composite/CompositeQueryLogic.java
@@ -16,6 +16,7 @@
 
 import org.apache.accumulo.core.client.AccumuloClient;
 import org.apache.accumulo.core.security.Authorizations;
+import org.apache.commons.collections4.Transformer;
 import org.apache.commons.collections4.iterators.TransformIterator;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -42,6 +43,7 @@
 import datawave.security.authorization.UserOperations;
 import datawave.webservice.query.result.event.EventBase;
 import datawave.webservice.result.BaseResponse;
+import datawave.webservice.result.QueryValidationResponse;
 
 /**
  * Query Logic implementation that is configured with more than one query logic delegate. The queries are run in parallel unless configured to be sequential.
@@ -726,6 +728,34 @@ public void setServerUser(ProxiedUserDetails user) {
         }
     }
 
+    @Override
+    public Object validateQuery(AccumuloClient client, Query query, Set<Authorizations> auths) throws Exception {
+        // see if we can find a query logic that supports this method
+        for (QueryLogic<?> logic : getQueryLogics().values()) {
+            try {
+                return logic.validateQuery(client, query, auths);
+            } catch (UnsupportedOperationException uoe) {
+                // try the next one
+            }
+        }
+        // ok, call the super method to throw the exception
+        return super.validateQuery(client, query, auths);
+    }
+
+    @Override
+    public Transformer<Object,QueryValidationResponse> getQueryValidationResponseTransformer() {
+        // see if we can find a query logic that supports this method
+        for (QueryLogic<?> logic : getQueryLogics().values()) {
+            try {
+                return logic.getQueryValidationResponseTransformer();
+            } catch (UnsupportedOperationException uoe) {
+                // try the next one
+            }
+        }
+        // ok, call the super method to throw the exception
+        return super.getQueryValidationResponseTransformer();
+    }
+
     /**
      * Setting the page processing start time is called after the logic is created. Pass this on to the children.
      *

diff --git a/core/query/src/main/java/datawave/core/query/logic/filtered/FilteredQueryLogic.java b/core/query/src/main/java/datawave/core/query/logic/filtered/FilteredQueryLogic.java
@@ -131,6 +131,15 @@ public TransformIterator getTransformIterator(Query settings) {
         }
     }
 
+    @Override
+    public Object validateQuery(AccumuloClient client, Query query, Set<Authorizations> auths) throws Exception {
+        if (!isFiltered()) {
+            return super.validateQuery(client, query, auths);
+        } else {
+            throw new UnsupportedOperationException("Query validation not implemented");
+        }
+    }
+
     @Override
     public UserOperations getUserOperations() {
         if (!isFiltered()) {

diff --git a/...uery/src/test/java/datawave/webservice/query/logic/composite/CompositeQueryLogicTest.java b/...uery/src/test/java/datawave/webservice/query/logic/composite/CompositeQueryLogicTest.java
@@ -20,6 +20,7 @@
 import org.apache.accumulo.core.security.ColumnVisibility;
 import org.apache.accumulo.core.security.VisibilityEvaluator;
 import org.apache.accumulo.core.security.VisibilityParseException;
+import org.apache.commons.collections4.Transformer;
 import org.apache.commons.collections4.iterators.TransformIterator;
 import org.junit.Assert;
 import org.junit.Before;
@@ -57,6 +58,7 @@
 import datawave.webservice.query.result.edge.EdgeBase;
 import datawave.webservice.result.BaseQueryResponse;
 import datawave.webservice.result.GenericResponse;
+import datawave.webservice.result.QueryValidationResponse;
 
 public class CompositeQueryLogicTest {
 
@@ -88,6 +90,8 @@ public class CompositeQueryLogicTest {
     private Value valueFailure = new Value(keyFailure.getRowData().getBackingArray());
     private Value valueSpecial = new Value(keySpecial.getRowData().getBackingArray());
 
+    private static final String VALIDATION_MESSAGE = "Light is green, the trap is clean";
+
     public static class TestQueryConfiguration extends GenericQueryConfiguration {
 
     }
@@ -454,6 +458,33 @@ public Set<String> getExampleQueries() {
             return Collections.emptySet();
         }
 
+        @Override
+        public Object validateQuery(AccumuloClient client, Query query, Set<Authorizations> auths) throws Exception {
+            // return something valid
+            return VALIDATION_MESSAGE;
+        }
+
+        @Override
+        public Transformer<Object,QueryValidationResponse> getQueryValidationResponseTransformer() {
+            return new TestQueryValidationResultTransformer();
+        }
+
+    }
+
+    public static class TestQueryValidationResultTransformer implements Transformer<Object,QueryValidationResponse> {
+
+        @Override
+        public QueryValidationResponse transform(Object input) {
+            String validation = String.valueOf(input);
+            QueryValidationResponse.Result result = new QueryValidationResponse.Result();
+            result.setMessages(Collections.singletonList(validation));
+
+            QueryValidationResponse response = new QueryValidationResponse();
+            response.setHasResults(true);
+            response.setResults(Collections.singletonList(result));
+
+            return response;
+        }
     }
 
     public static class TestFilteredQueryLogic extends FilteredQueryLogic {
@@ -1767,4 +1798,50 @@ public void testAuthorizationsUpdate() throws Exception {
         c.close();
     }
 
+    @Test(expected = UnsupportedOperationException.class)
+    public void testValidationFails() throws Exception {
+        Map<String,QueryLogic<?>> logics = new HashMap<>();
+        TestQueryLogic logic1 = new TestQueryLogic();
+        TestQueryLogic2 logic2 = new TestQueryLogic2();
+        logics.put("TestQueryLogic", logic1);
+        logics.put("TestQueryLogic2", logic2);
+
+        QueryImpl settings = new QueryImpl();
+        settings.setPagesize(100);
+        settings.setQueryAuthorizations(auths.toString());
+        settings.setQuery("FOO == 'BAR'");
+        settings.setParameters(new HashSet<>());
+        settings.setId(UUID.randomUUID());
+
+        CompositeQueryLogic c = new CompositeQueryLogic();
+        c.setQueryLogics(logics);
+        c.setCurrentUser(principal);
+
+        c.validateQuery(null, settings, Collections.singleton(auths));
+    }
+
+    @Test
+    public void testValidationHappyPath() throws Exception {
+        Map<String,QueryLogic<?>> logics = new HashMap<>();
+        TestQueryLogic logic1 = new TestQueryLogic();
+        DifferentTestQueryLogic logic2 = new DifferentTestQueryLogic();
+        logics.put("TestQueryLogic", logic1);
+        logics.put("TestQueryLogic2", logic2);
+
+        QueryImpl settings = new QueryImpl();
+        settings.setPagesize(100);
+        settings.setQueryAuthorizations(auths.toString());
+        settings.setQuery("FOO == 'BAR'");
+        settings.setParameters(new HashSet<>());
+        settings.setId(UUID.randomUUID());
+
+        CompositeQueryLogic c = new CompositeQueryLogic();
+        c.setQueryLogics(logics);
+        c.setCurrentUser(principal);
+
+        Object validation = c.validateQuery(null, settings, Collections.singleton(auths));
+        Transformer<Object,QueryValidationResponse> transformer = c.getQueryValidationResponseTransformer();
+        QueryValidationResponse response = transformer.transform(validation);
+        Assert.assertEquals(VALIDATION_MESSAGE, response.getResults().get(0).getMessages().get(0));
+    }
 }