find and insert ops

nosqlbench · May 9, 2024 · 464b6a8 · 464b6a8
1 parent 66ab79e
commit 464b6a8
Show file tree

Hide file tree

Showing 5 changed files with 168 additions and 6 deletions.
diff --git a/...n/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiFindOneAndUpdateOpDispenser.java b/...n/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiFindOneAndUpdateOpDispenser.java
@@ -16,19 +16,43 @@
 
 package io.nosqlbench.adapter.dataapi.opdispensers;
 
+import com.datastax.astra.client.Database;
+import com.datastax.astra.client.model.Filter;
+import com.datastax.astra.client.model.Update;
 import io.nosqlbench.adapter.dataapi.DataApiDriverAdapter;
 import io.nosqlbench.adapter.dataapi.ops.DataApiBaseOp;
+import io.nosqlbench.adapter.dataapi.ops.DataApiFindOneAndUpdateOp;
 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 DataApiFindOneAndUpdateOpDispenser extends DataApiOpDispenser {
+    private static final Logger logger = LogManager.getLogger(DataApiFindOneAndUpdateOpDispenser.class);
+    private final LongFunction<DataApiFindOneAndUpdateOp> opFunction;
     public DataApiFindOneAndUpdateOpDispenser(DataApiDriverAdapter adapter, ParsedOp op, LongFunction<String> targetFunction) {
         super(adapter, op, targetFunction);
+        this.opFunction = createOpFunction(op);
+    }
+
+    private LongFunction<DataApiFindOneAndUpdateOp> createOpFunction(ParsedOp op) {
+        return (l) -> {
+            Database db = spaceFunction.apply(l).getDatabase();
+            Filter filter = getFilterFromOp(op, l);
+            Update update = getUpdates(op, l);
+
+            return new DataApiFindOneAndUpdateOp(
+                db,
+                db.getCollection(targetFunction.apply(l)),
+                filter,
+                update
+            );
+        };
     }
 
     @Override
     public DataApiBaseOp getOp(long value) {
-        return null;
+        return opFunction.apply(value);
     }
 }
diff --git a/...rc/main/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiInsertManyOpDispenser.java b/...rc/main/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiInsertManyOpDispenser.java
@@ -16,19 +16,49 @@
 
 package io.nosqlbench.adapter.dataapi.opdispensers;
 
+import com.datastax.astra.client.model.Document;
+import com.datastax.astra.client.model.InsertManyOptions;
 import io.nosqlbench.adapter.dataapi.DataApiDriverAdapter;
 import io.nosqlbench.adapter.dataapi.ops.DataApiBaseOp;
+import io.nosqlbench.adapter.dataapi.ops.DataApiInsertManyOp;
 import io.nosqlbench.adapters.api.templating.ParsedOp;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
 
+import java.util.ArrayList;
+import java.util.List;
 import java.util.function.LongFunction;
 
 public class DataApiInsertManyOpDispenser extends DataApiOpDispenser {
+    private static final Logger logger = LogManager.getLogger(DataApiInsertManyOpDispenser.class);
+    private final LongFunction<DataApiInsertManyOp> opFunction;
+
     public DataApiInsertManyOpDispenser(DataApiDriverAdapter adapter, ParsedOp op, LongFunction<String> targetFunction) {
         super(adapter, op, targetFunction);
+        this.opFunction = createOpFunction(op);
+    }
+
+    private LongFunction<DataApiInsertManyOp> createOpFunction(ParsedOp op) {
+        return (l) -> {
+            List<Document> documents = new ArrayList<>();
+            op.getAsRequiredFunction("documents", List.class).apply(l).forEach(o -> documents.add(Document.parse(o.toString())));
+            return new DataApiInsertManyOp(
+                spaceFunction.apply(l).getDatabase(),
+                targetFunction.apply(l),
+                documents,
+                getInsertManyOptions(op, l)
+            );
+        };
+    }
+
+    private InsertManyOptions getInsertManyOptions(ParsedOp op, long l) {
+        InsertManyOptions options = new InsertManyOptions();
+
+        return options;
     }
 
     @Override
     public DataApiBaseOp getOp(long value) {
-        return null;
+        return opFunction.apply(value);
     }
 }
diff --git a/...-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiOpDispenser.java b/...-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiOpDispenser.java
@@ -16,10 +16,7 @@
 
 package io.nosqlbench.adapter.dataapi.opdispensers;
 
-import com.datastax.astra.client.model.Filter;
-import com.datastax.astra.client.model.Filters;
-import com.datastax.astra.client.model.Sort;
-import com.datastax.astra.client.model.Sorts;
+import com.datastax.astra.client.model.*;
 import io.nosqlbench.adapter.dataapi.DataApiSpace;
 import io.nosqlbench.adapter.dataapi.ops.DataApiBaseOp;
 import io.nosqlbench.adapters.api.activityimpl.BaseOpDispenser;
@@ -93,4 +90,34 @@ protected Filter getFilterFromOp(ParsedOp op, long l) {
         return filter;
     }
 
+    protected Update getUpdates(ParsedOp op, long l) {
+        Update update = new Update();
+        Optional<LongFunction<Map>> updatesFunction = op.getAsOptionalFunction("updates", Map.class);
+        if (updatesFunction.isPresent()) {
+            Map<String, Object> updates = updatesFunction.get().apply(l);
+            for (Map.Entry<String, Object> entry : updates.entrySet()) {
+                if (entry.getKey().equalsIgnoreCase("update")) {
+                    Map<String, Object> updateFields = (Map<String, Object>) entry.getValue();
+                    switch (updateFields.get("operation").toString()) {
+                        case "set" ->
+                            update = Updates.set(updateFields.get("field").toString(), updateFields.get("value"));
+                        case "inc" ->
+                            update = Updates.inc(updateFields.get("field").toString(), (double) updateFields.get("value"));
+                        case "unset" -> update = Updates.unset(updateFields.get("field").toString());
+                        case "addToSet" ->
+                            update = Updates.addToSet(updateFields.get("field").toString(), updateFields.get("value"));
+                        case "min" ->
+                            update = Updates.min(updateFields.get("field").toString(), (double) updateFields.get("value"));
+                        case "rename" ->
+                            update = Updates.rename(updateFields.get("field").toString(), updateFields.get("value").toString());
+                        default -> logger.error("Operation " + updateFields.get("operation") + " not supported");
+                    }
+                } else {
+                    logger.error("Filter " + entry.getKey() + " not supported");
+                }
+            }
+        }
+        return update;
+    }
+
 }
diff --git a/...er-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/ops/DataApiFindOneAndUpdateOp.java b/...er-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/ops/DataApiFindOneAndUpdateOp.java
@@ -0,0 +1,39 @@
+/*
+ * 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.*;
+
+public class DataApiFindOneAndUpdateOp extends DataApiBaseOp {
+    private final Collection<Document> collection;
+    private final Filter filter;
+    private final Update update;
+
+    public DataApiFindOneAndUpdateOp(Database db, Collection<Document> collection, Filter filter, Update update) {
+        super(db);
+        this.collection = collection;
+        this.filter = filter;
+        this.update = update;
+    }
+
+    @Override
+    public Object apply(long value) {
+        return collection.findOneAndUpdate(filter, update);
+    }
+}
diff --git a/.../adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/ops/DataApiInsertManyOp.java b/.../adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/ops/DataApiInsertManyOp.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.Database;
+import com.datastax.astra.client.model.Document;
+import com.datastax.astra.client.model.InsertManyOptions;
+
+import java.util.List;
+
+public class DataApiInsertManyOp extends DataApiBaseOp {
+    private final List<? extends Document> documents;
+    private final String collectionName;
+    private final InsertManyOptions options;
+
+
+    public DataApiInsertManyOp(Database db, String collectionName, List<? extends Document> documents, InsertManyOptions options) {
+        super(db);
+        this.collectionName = collectionName;
+        this.documents = documents;
+        this.options = options;
+    }
+
+    @Override
+    public Object apply(long value) {
+        return db.getCollection(collectionName).insertMany(documents, options);
+    }
+}