diff --git a/nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/DataApiOpMapper.java b/nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/DataApiOpMapper.java index 2703f8aad..c4539be8d 100644 --- a/nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/DataApiOpMapper.java +++ b/nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/DataApiOpMapper.java @@ -66,6 +66,9 @@ public class DataApiOpMapper implements OpMapper { 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); }; } } diff --git a/nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiCountDocumentsOpDispenser.java b/nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiCountDocumentsOpDispenser.java new file mode 100644 index 000000000..09459c2c3 --- /dev/null +++ b/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 opFunction; + + public DataApiCountDocumentsOpDispenser(DataApiDriverAdapter adapter, ParsedOp op, LongFunction targetFunction) { + super(adapter, op, targetFunction); + this.opFunction = createOpFunction(op); + } + + private LongFunction 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); + } +} diff --git a/nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiFindOneAndReplaceOpDispenser.java b/nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiFindOneAndReplaceOpDispenser.java new file mode 100644 index 000000000..7dd835a4e --- /dev/null +++ b/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 opFunction; + + public DataApiFindOneAndReplaceOpDispenser(DataApiDriverAdapter adapter, ParsedOp op, LongFunction targetFunction) { + super(adapter, op, targetFunction); + this.opFunction = createOpFunction(op); + } + + private LongFunction createOpFunction(ParsedOp op) { + return (l) -> { + Database db = spaceFunction.apply(l).getDatabase(); + Filter filter = getFilterFromOp(op, l); + FindOneAndReplaceOptions options = getFindOneAndReplaceOptions(op, l); + LongFunction docMapFunc = op.getAsRequiredFunction("document", Map.class); + LongFunction 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> 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); + } +} diff --git a/nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiReplaceOneOpDispenser.java b/nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiReplaceOneOpDispenser.java new file mode 100644 index 000000000..144969125 --- /dev/null +++ b/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 opFunction; + + public DataApiReplaceOneOpDispenser(DataApiDriverAdapter adapter, ParsedOp op, LongFunction targetFunction) { + super(adapter, op, targetFunction); + this.opFunction = createOpFunction(op); + } + + private LongFunction createOpFunction(ParsedOp op) { + return (l) -> { + Database db = spaceFunction.apply(l).getDatabase(); + Filter filter = getFilterFromOp(op, l); + ReplaceOneOptions options = getReplaceOneOptions(op, l); + LongFunction docMapFunc = op.getAsRequiredFunction("document", Map.class); + LongFunction 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> 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); + } +} diff --git a/nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/ops/DataApiCountDocumentsOp.java b/nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/ops/DataApiCountDocumentsOp.java new file mode 100644 index 000000000..dfb27bba8 --- /dev/null +++ b/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 collection; + private final Filter filter; + private final int upperBound; + + public DataApiCountDocumentsOp(Database db, Collection 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); + } + } +} diff --git a/nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/ops/DataApiFindOneAndReplaceOp.java b/nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/ops/DataApiFindOneAndReplaceOp.java new file mode 100644 index 000000000..903fc3b94 --- /dev/null +++ b/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); + } +} diff --git a/nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/ops/DataApiOpType.java b/nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/ops/DataApiOpType.java index caf50cac7..fdaf746ae 100644 --- a/nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/ops/DataApiOpType.java +++ b/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, } diff --git a/nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/ops/DataApiReplaceOneOp.java b/nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/ops/DataApiReplaceOneOp.java new file mode 100644 index 000000000..b35f4338e --- /dev/null +++ b/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); + } +}