Add support for totalPagedResultsPolicy in scripts (WrenSecurity#188)

orchitech · Jun 28, 2024 · 78e9d1a · 78e9d1a
1 parent 402d3b2
commit 78e9d1a
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 14 deletions.
diff --git a/openidm-script/src/main/java/org/forgerock/openidm/script/ResourceFunctions.java b/openidm-script/src/main/java/org/forgerock/openidm/script/ResourceFunctions.java
@@ -24,28 +24,19 @@
 
 package org.forgerock.openidm.script;
 
-import java.lang.IllegalArgumentException;
-import java.lang.NoSuchMethodException;
-import java.lang.Object;
-import java.lang.Override;
-import java.lang.String;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.HashMap;
 import java.util.LinkedHashMap;
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
-
-import org.forgerock.script.scope.Function;
-import org.forgerock.script.scope.FunctionFactory;
-import org.forgerock.script.scope.Parameter;
-import org.forgerock.services.context.Context;
 import org.forgerock.json.JsonValue;
 import org.forgerock.json.resource.ActionRequest;
 import org.forgerock.json.resource.ActionResponse;
 import org.forgerock.json.resource.BadRequestException;
 import org.forgerock.json.resource.ConnectionFactory;
+import org.forgerock.json.resource.CountPolicy;
 import org.forgerock.json.resource.CreateRequest;
 import org.forgerock.json.resource.DeleteRequest;
 import org.forgerock.json.resource.NotFoundException;
@@ -58,9 +49,13 @@
 import org.forgerock.json.resource.ReadRequest;
 import org.forgerock.json.resource.Request;
 import org.forgerock.json.resource.Requests;
-import org.forgerock.json.resource.ResourceResponse;
 import org.forgerock.json.resource.ResourceException;
+import org.forgerock.json.resource.ResourceResponse;
 import org.forgerock.json.resource.UpdateRequest;
+import org.forgerock.script.scope.Function;
+import org.forgerock.script.scope.FunctionFactory;
+import org.forgerock.script.scope.Parameter;
+import org.forgerock.services.context.Context;
 import org.forgerock.util.Factory;
 import org.forgerock.util.LazyMap;
 
@@ -619,6 +614,7 @@ public JsonValue call(Parameter scope, final Function<?> callback, Object... arg
             if (null != queryResponse) {
                 result.put("pagedResultsCookie", queryResponse.getPagedResultsCookie());
                 result.put("totalPagedResults", queryResponse.getTotalPagedResults());
+                result.put("totalPagedResultsPolicy", queryResponse.getTotalPagedResultsPolicy().name());
             }
             if (null != results) {
                 result.put("result", results);
@@ -680,6 +676,15 @@ private final QueryResponse query(final Parameter scope, String resourceContaine
                         qr.setPagedResultsOffset(params.get(name).required().asInteger());
                     } else if (name.equalsIgnoreCase("_pageSize")) {
                         qr.setPageSize(params.get(name).required().asInteger());
+                    } else if (name.equalsIgnoreCase("_totalPagedResultsPolicy")) {
+                        final String s = params.get(name).required().asString();
+                        try {
+                            qr.setTotalPagedResultsPolicy(CountPolicy.valueOf(s));
+                        } catch (IllegalArgumentException e) {
+                            // FIXME: i18n.
+                            throw new BadRequestException("The value '" + s + "' for parameter '"
+                                    + name + "' could not be parsed as a valid count policy");
+                        }
                     } else if (name.equalsIgnoreCase("_queryFilter")) {
                         final String s = params.get(name).required().asString();
                         try {
@@ -688,7 +693,6 @@ private final QueryResponse query(final Parameter scope, String resourceContaine
                             // FIXME: i18n.
                             throw new BadRequestException("The value '" + s + "' for parameter '"
                                     + name + "' could not be parsed as a valid query filter");
-
                         }
                     } else {
                         setAdditionalParameter(qr, name, params.get(name));

diff --git a/openidm-script/src/test/resources/container/resource.groovy b/openidm-script/src/test/resources/container/resource.groovy
@@ -90,15 +90,20 @@ def queryParams1 = [
         _queryId: "query-all-ids"
 ]
 def queryParams2 = [
-        _queryFilter: "nickName eq \"Babs\""
+        _queryFilter: "nickName eq \"Babs\"",
+        _pageSize: 42,
+        _totalPagedResultsPolicy: "EXACT"
 ]
 
 try {
     def userQ1 = router.query("Users", queryParams1)
 } catch (NotSupportedException e) {
     //expected
 }
+
 def userQ2 = router.query("Users", queryParams2)
+assert userQ2.totalPagedResultsPolicy == "EXACT"
+
 def userQ3 = router.query("Users", queryParams2, callback)
 
 def printResult = { resource, error ->

diff --git a/openidm-script/src/test/resources/container/resource.js b/openidm-script/src/test/resources/container/resource.js
@@ -121,15 +121,22 @@ var queryParams1 = {
     "_queryId": "query-all-ids"
 }
 var queryParams2 = {
-    "_queryFilter": "nickName eq \"Babs\""
+    "_queryFilter": "nickName eq \"Babs\"",
+    "_pageSize": 42,
+    "_totalPagedResultsPolicy": "EXACT"
 }
 
 try {
     var userQ1 = router.query("Users", queryParams1)
 } catch (e) {
     //expected
 }
+
 var userQ2 = router.query("Users", queryParams2)
+if (userQ2.totalPagedResultsPolicy != "EXACT") {
+    throw { "message": "Expected total paged results policy not set" };
+}
+
 var userQ3 = router.query("Users", queryParams2, callback)
 
 var printResult = function (resource, error) {