countDocuments, add new op and split from legacy

nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/DataApiOpMapper.java

@@ -108,7 +108,6 @@ public class DataApiOpMapper implements OpMapper<DataApiBaseOp,DataApiSpace> {
             case find_one_and_delete -> new DataApiCollectionFindOneAndDeleteOpDispenser(adapter, op, typeAndTarget.targetFunction);
             case replace_one -> new DataApiCollectionReplaceOneOpDispenser(adapter, op, typeAndTarget.targetFunction);
             case find_one_and_replace -> new DataApiCollectionFindOneAndReplaceOpDispenser(adapter, op, typeAndTarget.targetFunction);
-            case count_documents -> new DataApiCollectionCountDocumentsOpDispenser(adapter, op, typeAndTarget.targetFunction);
 
             // LEGACY OPS
             // admin ops:
@@ -116,8 +115,6 @@ public class DataApiOpMapper implements OpMapper<DataApiBaseOp,DataApiSpace> {
             // in-database ops:
             case create_collection_with_class -> new DataApiLegacyCreateCollectionWithClassOpDispenser(adapter, op, typeAndTarget.targetFunction);
             // in-collection ops:
-            case collection_delete_many -> new DataApiCollectionDeleteManyOpDispenser(adapter, op, typeAndTarget.targetFunction);
-            case collection_delete_one -> new DataApiCollectionDeleteOneOpDispenser(adapter, op, typeAndTarget.targetFunction);
             case delete_one -> new DataApiCollectionLegacyDeleteOneOpDispenser(adapter, op, typeAndTarget.targetFunction);
             case delete_many -> new DataApiCollectionLegacyDeleteManyOpDispenser(adapter, op, typeAndTarget.targetFunction);
             case delete_all -> new DataApiCollectionLegacyDeleteAllOpDispenser(adapter, op, typeAndTarget.targetFunction);
@@ -128,15 +125,20 @@ public class DataApiOpMapper implements OpMapper<DataApiBaseOp,DataApiSpace> {
             case find_distinct -> new DataApiCollectionLegacyFindDistinctOpDispenser(adapter, op, typeAndTarget.targetFunction);
             case find_one -> new DataApiCollectionLegacyFindOneOpDispenser(adapter, op, typeAndTarget.targetFunction);
             case find_one_and_update -> new DataApiCollectionLegacyFindOneAndUpdateOpDispenser(adapter, op, typeAndTarget.targetFunction);
+            case count_documents -> new DataApiCollectionLegacyCountDocumentsOpDispenser(adapter, op, typeAndTarget.targetFunction);
             case estimated_document_count ->
                 new DataApiCollectionEstimatedDocumentCountOpDispenser(adapter, op, typeAndTarget.targetFunction);
 
             // NEW-STYLE OPS
+            // in-collection ops:
             case collection_find -> new DataApiCollectionFindOpDispenser(adapter, op, typeAndTarget.targetFunction);
             case collection_find_one -> new DataApiCollectionFindOneOpDispenser(adapter, op, typeAndTarget.targetFunction);
             case collection_find_one_and_update -> new DataApiCollectionFindOneAndUpdateOpDispenser(adapter, op, typeAndTarget.targetFunction);
+            case collection_delete_many -> new DataApiCollectionDeleteManyOpDispenser(adapter, op, typeAndTarget.targetFunction);
+            case collection_delete_one -> new DataApiCollectionDeleteOneOpDispenser(adapter, op, typeAndTarget.targetFunction);
             case collection_estimated_document_count ->
                 new DataApiCollectionEstimatedDocumentCountOpDispenser(adapter, op, typeAndTarget.targetFunction);
+            case collection_count_documents -> new DataApiCollectionCountDocumentsOpDispenser(adapter, op, typeAndTarget.targetFunction);
         };
     }
 

nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiCollectionCountDocumentsOpDispenser.java

@@ -26,6 +26,7 @@ import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 
 import java.util.function.LongFunction;
+import java.util.Optional;
 
 public class DataApiCollectionCountDocumentsOpDispenser extends DataApiOpDispenser {
     private static final Logger logger = LogManager.getLogger(DataApiCollectionCountDocumentsOpDispenser.class);
@@ -40,7 +41,13 @@ public class DataApiCollectionCountDocumentsOpDispenser extends DataApiOpDispens
         return (l) -> {
             Database db = spaceFunction.apply(l).getDatabase();
             Filter filter = getFilterFromOp(op, l);
-            int upperBound = op.getAsRequiredFunction("upperbound", Integer.class).apply(l);
+            Optional<LongFunction<Integer>> ubf = op.getAsOptionalFunction("upper_bound", Integer.class);
+            Optional<Integer> upperBound;
+            if (ubf.isPresent()) {
+                upperBound = Optional.of(ubf.get().apply(l));
+            } else {
+                upperBound = Optional.empty();
+            }
 
             return new DataApiCollectionCountDocumentsOp(
                 db,

nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiCollectionLegacyCountDocumentsOpDispenser.java

@@ -0,0 +1,58 @@
+/*
+ * Copyright (c) nosqlbench
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.nosqlbench.adapter.dataapi.opdispensers;
+
+import com.datastax.astra.client.databases.Database;
+import com.datastax.astra.client.core.query.Filter;
+import io.nosqlbench.adapter.dataapi.DataApiDriverAdapter;
+import io.nosqlbench.adapter.dataapi.ops.DataApiBaseOp;
+import io.nosqlbench.adapter.dataapi.ops.DataApiCollectionLegacyCountDocumentsOp;
+import io.nosqlbench.adapters.api.templating.ParsedOp;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+
+import java.util.function.LongFunction;
+
+public class DataApiCollectionLegacyCountDocumentsOpDispenser extends DataApiOpDispenser {
+    private static final Logger logger = LogManager.getLogger(DataApiCollectionLegacyCountDocumentsOpDispenser.class);
+    private final LongFunction<DataApiCollectionLegacyCountDocumentsOp> opFunction;
+
+    public DataApiCollectionLegacyCountDocumentsOpDispenser(DataApiDriverAdapter adapter, ParsedOp op, LongFunction<String> targetFunction) {
+        super(adapter, op, targetFunction);
+        this.opFunction = createOpFunction(op);
+    }
+
+    private LongFunction<DataApiCollectionLegacyCountDocumentsOp> createOpFunction(ParsedOp op) {
+        return (l) -> {
+            Database db = spaceFunction.apply(l).getDatabase();
+            Filter filter = getLegacyFilterFromOp(op, l);
+            int upperBound = op.getAsRequiredFunction("upperbound", Integer.class).apply(l);
+
+            return new DataApiCollectionLegacyCountDocumentsOp(
+                db,
+                db.getCollection(targetFunction.apply(l)),
+                filter,
+                upperBound
+            );
+        };
+    }
+
+    @Override
+    public DataApiBaseOp getOp(long cycle) {
+        return opFunction.apply(cycle);
+    }
+}

nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/ops/DataApiCollectionCountDocumentsOp.java

@@ -22,12 +22,17 @@ import com.datastax.astra.client.collections.exceptions.TooManyDocumentsToCountE
 import com.datastax.astra.client.collections.definition.documents.Document;
 import com.datastax.astra.client.core.query.Filter;
 
+import java.util.Optional;
+
 public class DataApiCollectionCountDocumentsOp extends DataApiBaseOp {
     private final Collection<Document> collection;
     private final Filter filter;
-    private final int upperBound;
+    private final Optional<Integer> upperBound;
 
-    public DataApiCollectionCountDocumentsOp(Database db, Collection<Document> collection, Filter filter, int upperBound) {
+    private final int MAX_UPPER_BOUND = 1000;
+    private final int SILENT_FAILED_COUNT_RESULT = -1;
+
+    public DataApiCollectionCountDocumentsOp(Database db, Collection<Document> collection, Filter filter, Optional<Integer> upperBound) {
         super(db);
         this.collection = collection;
         this.filter = filter;
@@ -36,10 +41,21 @@ public class DataApiCollectionCountDocumentsOp extends DataApiBaseOp {
 
     @Override
     public Object apply(long value) {
+        // upper-bound being optional enables a pattern that ignores the "moreData"
+        // (getting rid of the 'upper bound' client-only feature). Still honored if passed (incl. throwing the error).
+        int upperBoundToUse = MAX_UPPER_BOUND;
+        boolean shouldThrowBoundExceeded = false;
+        if (upperBound.isPresent()) {
+            upperBoundToUse = upperBound.get();
+            shouldThrowBoundExceeded = true;
+        }
         try {
-            return collection.countDocuments(filter, upperBound);
+            return collection.countDocuments(filter, upperBoundToUse);
         } catch (TooManyDocumentsToCountException e) {
-            throw new RuntimeException(e);
+            if (shouldThrowBoundExceeded) {
+                throw new RuntimeException(e);
+            }
+            return SILENT_FAILED_COUNT_RESULT;
         }
     }
 }

nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/ops/DataApiCollectionLegacyCountDocumentsOp.java

@@ -0,0 +1,45 @@
+/*
+ * Copyright (c) nosqlbench
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.nosqlbench.adapter.dataapi.ops;
+
+import com.datastax.astra.client.collections.Collection;
+import com.datastax.astra.client.databases.Database;
+import com.datastax.astra.client.collections.exceptions.TooManyDocumentsToCountException;
+import com.datastax.astra.client.collections.definition.documents.Document;
+import com.datastax.astra.client.core.query.Filter;
+
+public class DataApiCollectionLegacyCountDocumentsOp extends DataApiBaseOp {
+    private final Collection<Document> collection;
+    private final Filter filter;
+    private final int upperBound;
+
+    public DataApiCollectionLegacyCountDocumentsOp(Database db, Collection<Document> collection, Filter filter, int upperBound) {
+        super(db);
+        this.collection = collection;
+        this.filter = filter;
+        this.upperBound = upperBound;
+    }
+
+    @Override
+    public Object apply(long value) {
+        try {
+            return collection.countDocuments(filter, upperBound);
+        } catch (TooManyDocumentsToCountException e) {
+            throw new RuntimeException(e);
+        }
+    }
+}

nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/ops/DataApiOpType.java

@@ -42,7 +42,6 @@ public enum DataApiOpType {
     find_one_and_delete,
     replace_one,
     find_one_and_replace,
-    count_documents,
     
     // LEGACY OPS
     // admin ops:
@@ -61,6 +60,7 @@ public enum DataApiOpType {
     find_one,
     find_one_and_update,
     estimated_document_count,
+    count_documents,
 
     // NEW-STYLE OPS
     // admin ops:
@@ -73,4 +73,5 @@ public enum DataApiOpType {
     collection_find_one,
     collection_find_one_and_update,
     collection_estimated_document_count,
+    collection_count_documents,
 }