impl of replace ops

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

@@ -66,6 +66,9 @@ public class DataApiOpMapper implements OpMapper<DataApiBaseOp> {
                 new DataApiEstimatedDocumentCountOpDispenser(adapter, op, typeAndTarget.targetFunction);
             case find_by_id -> new DataApiFindByIdOpDispenser(adapter, op, typeAndTarget.targetFunction);
             case find_distinct -> new DataApiFindDistinctOpDispenser(adapter, op, typeAndTarget.targetFunction);
+            case count_documents -> new DataApiCountDocumentsOpDispenser(adapter, op, typeAndTarget.targetFunction);
+            case replace_one -> new DataApiReplaceOneOpDispenser(adapter, op, typeAndTarget.targetFunction);
+            case find_one_and_replace -> new DataApiFindOneAndReplaceOpDispenser(adapter, op, typeAndTarget.targetFunction);
         };
     }
 }

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

@@ -0,0 +1,58 @@
+/*
+ * Copyright (c) 2024 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.Database;
+import com.datastax.astra.client.model.Filter;
+import io.nosqlbench.adapter.dataapi.DataApiDriverAdapter;
+import io.nosqlbench.adapter.dataapi.ops.DataApiBaseOp;
+import io.nosqlbench.adapter.dataapi.ops.DataApiCountDocumentsOp;
+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 DataApiCountDocumentsOpDispenser extends DataApiOpDispenser {
+    private static final Logger logger = LogManager.getLogger(DataApiCountDocumentsOpDispenser.class);
+    private final LongFunction<DataApiCountDocumentsOp> opFunction;
+
+    public DataApiCountDocumentsOpDispenser(DataApiDriverAdapter adapter, ParsedOp op, LongFunction<String> targetFunction) {
+        super(adapter, op, targetFunction);
+        this.opFunction = createOpFunction(op);
+    }
+
+    private LongFunction<DataApiCountDocumentsOp> createOpFunction(ParsedOp op) {
+        return (l) -> {
+            Database db = spaceFunction.apply(l).getDatabase();
+            Filter filter = getFilterFromOp(op, l);
+            int upperBound = op.getAsRequiredFunction("upperbound", Integer.class).apply(l);
+
+            return new DataApiCountDocumentsOp(
+                db,
+                db.getCollection(targetFunction.apply(l)),
+                filter,
+                upperBound
+            );
+        };
+    }
+
+    @Override
+    public DataApiBaseOp getOp(long value) {
+        return opFunction.apply(value);
+    }
+}

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

@@ -0,0 +1,92 @@
+/*
+ * Copyright (c) 2024 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.Database;
+import com.datastax.astra.client.model.*;
+import io.nosqlbench.adapter.dataapi.DataApiDriverAdapter;
+import io.nosqlbench.adapter.dataapi.ops.DataApiBaseOp;
+import io.nosqlbench.adapter.dataapi.ops.DataApiFindOneAndReplaceOp;
+import io.nosqlbench.adapters.api.templating.ParsedOp;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+
+import java.util.Map;
+import java.util.Optional;
+import java.util.function.LongFunction;
+
+public class DataApiFindOneAndReplaceOpDispenser extends DataApiOpDispenser {
+    private static final Logger logger = LogManager.getLogger(DataApiFindOneAndReplaceOpDispenser.class);
+    private final LongFunction<DataApiFindOneAndReplaceOp> opFunction;
+
+    public DataApiFindOneAndReplaceOpDispenser(DataApiDriverAdapter adapter, ParsedOp op, LongFunction<String> targetFunction) {
+        super(adapter, op, targetFunction);
+        this.opFunction = createOpFunction(op);
+    }
+
+    private LongFunction<DataApiFindOneAndReplaceOp> createOpFunction(ParsedOp op) {
+        return (l) -> {
+            Database db = spaceFunction.apply(l).getDatabase();
+            Filter filter = getFilterFromOp(op, l);
+            FindOneAndReplaceOptions options = getFindOneAndReplaceOptions(op, l);
+            LongFunction<Map> docMapFunc = op.getAsRequiredFunction("document", Map.class);
+            LongFunction<Document> docFunc = (long m) -> new Document(docMapFunc.apply(m));
+
+            return new DataApiFindOneAndReplaceOp(
+                db,
+                db.getCollection(targetFunction.apply(l)),
+                filter,
+                docFunc.apply(l),
+                options
+            );
+        };
+    }
+
+    private FindOneAndReplaceOptions getFindOneAndReplaceOptions(ParsedOp op, long l) {
+        FindOneAndReplaceOptions options = new FindOneAndReplaceOptions();
+        Sort sort = getSortFromOp(op, l);
+        if (op.isDefined("vector")) {
+            float[] vector = getVectorValues(op, l);
+            if (sort != null) {
+                options = vector != null ? options.sort(vector, sort) : options.sort(sort);
+            } else if (vector != null) {
+                options = options.sort(vector);
+            }
+        }
+        Projection[] projection = getProjectionFromOp(op, l);
+        if (projection != null) {
+            options = options.projection(projection);
+        }
+        Optional<LongFunction<Boolean>> upsertFunction = op.getAsOptionalFunction("upsert", Boolean.class);
+        if (upsertFunction.isPresent()) {
+            options = options.upsert(upsertFunction.get().apply(l));
+        }
+        if (op.isDefined("returnDocument")) {
+            options = switch ((String) op.get("returnDocument", l)) {
+                case "after" -> options.returnDocumentAfter();
+                case "before" -> options.returnDocumentBefore();
+                default -> throw new RuntimeException("Invalid returnDocument value: " + op.get("returnDocument", l));
+            };
+        }
+        return options;
+    }
+
+    @Override
+    public DataApiBaseOp getOp(long value) {
+        return opFunction.apply(value);
+    }
+}

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

@@ -0,0 +1,74 @@
+/*
+ * Copyright (c) 2024 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.Database;
+import com.datastax.astra.client.model.*;
+import io.nosqlbench.adapter.dataapi.DataApiDriverAdapter;
+import io.nosqlbench.adapter.dataapi.ops.DataApiBaseOp;
+import io.nosqlbench.adapter.dataapi.ops.DataApiReplaceOneOp;
+import io.nosqlbench.adapter.dataapi.ops.DataApiUpdateOneOp;
+import io.nosqlbench.adapters.api.templating.ParsedOp;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+
+import java.util.Map;
+import java.util.Optional;
+import java.util.function.LongFunction;
+
+public class DataApiReplaceOneOpDispenser extends DataApiOpDispenser {
+    private static final Logger logger = LogManager.getLogger(DataApiReplaceOneOpDispenser.class);
+    private final LongFunction<DataApiReplaceOneOp> opFunction;
+
+    public DataApiReplaceOneOpDispenser(DataApiDriverAdapter adapter, ParsedOp op, LongFunction<String> targetFunction) {
+        super(adapter, op, targetFunction);
+        this.opFunction = createOpFunction(op);
+    }
+
+    private LongFunction<DataApiReplaceOneOp> createOpFunction(ParsedOp op) {
+        return (l) -> {
+            Database db = spaceFunction.apply(l).getDatabase();
+            Filter filter = getFilterFromOp(op, l);
+            ReplaceOneOptions options = getReplaceOneOptions(op, l);
+            LongFunction<Map> docMapFunc = op.getAsRequiredFunction("document", Map.class);
+            LongFunction<Document> docFunc = (long m) -> new Document(docMapFunc.apply(m));
+
+            return new DataApiReplaceOneOp(
+                db,
+                db.getCollection(targetFunction.apply(l)),
+                filter,
+                docFunc.apply(l),
+                options
+            );
+        };
+    }
+
+    private ReplaceOneOptions getReplaceOneOptions(ParsedOp op, long l) {
+        ReplaceOneOptions options = new ReplaceOneOptions();
+
+        Optional<LongFunction<Boolean>> upsertFunction = op.getAsOptionalFunction("upsert", Boolean.class);
+        if (upsertFunction.isPresent()) {
+            options = options.upsert(upsertFunction.get().apply(l));
+        }
+        return options;
+    }
+
+    @Override
+    public DataApiBaseOp getOp(long value) {
+        return opFunction.apply(value);
+    }
+}

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

@@ -0,0 +1,46 @@
+/*
+ * Copyright (c) 2024 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.Collection;
+import com.datastax.astra.client.Database;
+import com.datastax.astra.client.exception.TooManyDocumentsToCountException;
+import com.datastax.astra.client.model.Document;
+import com.datastax.astra.client.model.Filter;
+import com.datastax.astra.client.model.FindOneAndDeleteOptions;
+
+public class DataApiCountDocumentsOp extends DataApiBaseOp {
+    private final Collection<Document> collection;
+    private final Filter filter;
+    private final int upperBound;
+
+    public DataApiCountDocumentsOp(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/DataApiFindOneAndReplaceOp.java

@@ -0,0 +1,42 @@
+/*
+ * Copyright (c) 2024 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.Collection;
+import com.datastax.astra.client.Database;
+import com.datastax.astra.client.model.Filter;
+import com.datastax.astra.client.model.FindOneAndReplaceOptions;
+
+public class DataApiFindOneAndReplaceOp extends DataApiBaseOp {
+    private final Collection collection;
+    private final Filter filter;
+    private final Object replacement;
+    private final FindOneAndReplaceOptions options;
+
+    public DataApiFindOneAndReplaceOp(Database db, Collection collection, Filter filter, Object replacement, FindOneAndReplaceOptions options) {
+        super(db);
+        this.collection = collection;
+        this.filter = filter;
+        this.replacement = replacement;
+        this.options = options;
+    }
+
+    @Override
+    public Object apply(long value) {
+        return collection.findOneAndReplace(filter, replacement, options);
+    }
+}

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

@@ -37,4 +37,7 @@ public enum DataApiOpType {
     estimated_document_count,
     find_by_id,
     find_distinct,
+    count_documents,
+    replace_one,
+    find_one_and_replace,
 }

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

@@ -0,0 +1,42 @@
+/*
+ * Copyright (c) 2024 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.Collection;
+import com.datastax.astra.client.Database;
+import com.datastax.astra.client.model.Filter;
+import com.datastax.astra.client.model.ReplaceOneOptions;
+
+public class DataApiReplaceOneOp extends DataApiBaseOp {
+    private final Collection collection;
+    private final Filter filter;
+    private final Object replacement;
+    private final ReplaceOneOptions options;
+
+    public DataApiReplaceOneOp(Database db, Collection collection, Filter filter, Object replacement, ReplaceOneOptions options) {
+        super(db);
+        this.collection = collection;
+        this.filter = filter;
+        this.replacement = replacement;
+        this.options = options;
+    }
+
+    @Override
+    public Object apply(long value) {
+        return collection.replaceOne(filter, replacement, options);
+    }
+}