diff --git a/nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiInsertOneOpDispenser.java b/nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiInsertOneOpDispenser.java index ee5b118a6..2c43ca8ee 100644 --- a/nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiInsertOneOpDispenser.java +++ b/nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiInsertOneOpDispenser.java @@ -16,19 +16,38 @@ package io.nosqlbench.adapter.dataapi.opdispensers; +import com.datastax.astra.client.model.Document; import io.nosqlbench.adapter.dataapi.DataApiDriverAdapter; import io.nosqlbench.adapter.dataapi.ops.DataApiBaseOp; +import io.nosqlbench.adapter.dataapi.ops.DataApiInsertOneOp; 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 DataApiInsertOneOpDispenser extends DataApiOpDispenser { + private static final Logger logger = LogManager.getLogger(DataApiInsertOneOpDispenser.class); + private final LongFunction opFunction; + public DataApiInsertOneOpDispenser(DataApiDriverAdapter adapter, ParsedOp op, LongFunction targetFunction) { super(adapter, op, targetFunction); + this.opFunction = createOpFunction(op); + } + + private LongFunction createOpFunction(ParsedOp op) { + return (l) -> { + Document.parse(op.get("document", l)); + return new DataApiInsertOneOp( + spaceFunction.apply(l).getDatabase(), + targetFunction.apply(l), + Document.parse(op.get("document", l)) + ); + }; } @Override public DataApiBaseOp getOp(long value) { - return null; + return opFunction.apply(value); } } diff --git a/nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiInsertOneVectorOpDispenser.java b/nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiInsertOneVectorOpDispenser.java new file mode 100644 index 000000000..040dcbc01 --- /dev/null +++ b/nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiInsertOneVectorOpDispenser.java @@ -0,0 +1,56 @@ +/* + * 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.model.Document; +import io.nosqlbench.adapter.dataapi.DataApiDriverAdapter; +import io.nosqlbench.adapter.dataapi.ops.DataApiBaseOp; +import io.nosqlbench.adapter.dataapi.ops.DataApiInsertOneOp; +import io.nosqlbench.adapter.dataapi.ops.DataApiInsertOneVectorOp; +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 DataApiInsertOneVectorOpDispenser extends DataApiOpDispenser { + private static final Logger logger = LogManager.getLogger(DataApiInsertOneVectorOpDispenser.class); + private final LongFunction opFunction; + + public DataApiInsertOneVectorOpDispenser(DataApiDriverAdapter adapter, ParsedOp op, LongFunction targetFunction) { + super(adapter, op, targetFunction); + this.opFunction = createOpFunction(op); + } + + private LongFunction createOpFunction(ParsedOp op) { + return (l) -> { + Document.parse(op.get("document", l)); + float[] vector = op.get("vector", l); + return new DataApiInsertOneVectorOp( + spaceFunction.apply(l).getDatabase(), + targetFunction.apply(l), + Document.parse(op.get("document", l)), + vector + ); + }; + } + + @Override + public DataApiBaseOp getOp(long value) { + return opFunction.apply(value); + } +} diff --git a/nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/ops/DataApiInsertOneOp.java b/nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/ops/DataApiInsertOneOp.java new file mode 100644 index 000000000..13fc5893b --- /dev/null +++ b/nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/ops/DataApiInsertOneOp.java @@ -0,0 +1,20 @@ +package io.nosqlbench.adapter.dataapi.ops; + +import com.datastax.astra.client.Database; +import com.datastax.astra.client.model.Document; + +public class DataApiInsertOneOp extends DataApiBaseOp { + private final Document doc; + private final String collectionName; + + public DataApiInsertOneOp(Database db, String collectionName, Document doc) { + super(db); + this.collectionName = collectionName; + this.doc = doc; + } + + @Override + public Object apply(long value) { + return db.getCollection(collectionName).insertOne(doc); + } +} diff --git a/nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/ops/DataApiInsertOneVectorOp.java b/nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/ops/DataApiInsertOneVectorOp.java new file mode 100644 index 000000000..dcb1966aa --- /dev/null +++ b/nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/ops/DataApiInsertOneVectorOp.java @@ -0,0 +1,22 @@ +package io.nosqlbench.adapter.dataapi.ops; + +import com.datastax.astra.client.Database; +import com.datastax.astra.client.model.Document; + +public class DataApiInsertOneVectorOp extends DataApiBaseOp { + private final Document doc; + private final String collectionName; + private float[] vector; + + public DataApiInsertOneVectorOp(Database db, String collectionName, Document doc, float[] vector) { + super(db); + this.collectionName = collectionName; + this.doc = doc; + this.vector = vector; + } + + @Override + public Object apply(long value) { + return db.getCollection(collectionName).insertOne(doc, vector); + } +}