Merge pull request #1962 from nosqlbench/mwolters/dapi

Mwolters/dapi
This commit is contained in:
Jonathan Shook 2024-06-17 16:15:00 -05:00 committed by GitHub
commit 7c5613a42f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
43 changed files with 1748 additions and 68 deletions

View File

@ -64,6 +64,20 @@ public class DataApiOpMapper implements OpMapper<DataApiBaseOp> {
new DataApiListCollectionNamesOpDispenser(adapter, op, typeAndTarget.targetFunction);
case estimated_document_count ->
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);
case delete_all -> new DataApiDeleteAllOpDispenser(adapter, op, typeAndTarget.targetFunction);
case create_collection_with_class -> new DataApiCreateCollectionWithClassOpDispenser(adapter, op, typeAndTarget.targetFunction);
case create_database -> new DataApiCreateDatabaseOpDispenser(adapter, op, typeAndTarget.targetFunction);
case list_databases -> new DataApiListDatabasesOpDispenser(adapter, op, typeAndTarget.targetFunction);
case drop_database -> new DataApiDropDatabaseOpDispenser(adapter, op, typeAndTarget.targetFunction);
case get_database_info -> new DataApiGetDatabaseInfoOpDispenser(adapter, op, typeAndTarget.targetFunction);
case create_namespace -> new DataApiCreateNamespaceOpDispenser(adapter, op, typeAndTarget.targetFunction);
case list_namespaces -> new DataApiListNamespacesOpDispenser(adapter, op, typeAndTarget.targetFunction);
case drop_namespace -> new DataApiDropNamespaceOpDispenser(adapter, op, typeAndTarget.targetFunction);
};
}
}

View File

@ -18,6 +18,8 @@ package io.nosqlbench.adapter.dataapi;
import com.datastax.astra.client.DataAPIClient;
import com.datastax.astra.client.Database;
import com.datastax.astra.client.admin.AstraDBAdmin;
import com.datastax.astra.client.admin.DatabaseAdmin;
import io.nosqlbench.nb.api.config.standard.ConfigModel;
import io.nosqlbench.nb.api.config.standard.NBConfigModel;
import io.nosqlbench.nb.api.config.standard.NBConfiguration;
@ -42,10 +44,15 @@ public class DataApiSpace {
private Database database;
private String namespace;
private String superUserToken;
private AstraDBAdmin admin;
private DatabaseAdmin namespaceAdmin;
public DataApiSpace(String name, NBConfiguration cfg) {
this.config = cfg;
this.name = name;
setToken();
setSuperToken();
setApiEndpoint();
setNamespace();
createClient();
@ -59,6 +66,14 @@ public class DataApiSpace {
return database;
}
public AstraDBAdmin getAdmin() {
return admin;
}
public DatabaseAdmin getNamespaceAdmin() {
return namespaceAdmin;
}
private void createClient() {
this.dataAPIClient = new DataAPIClient(astraToken);
if (namespace != null) {
@ -66,6 +81,12 @@ public class DataApiSpace {
} else {
this.database = dataAPIClient.getDatabase(astraApiEndpoint);
}
if (superUserToken != null) {
this.admin = dataAPIClient.getAdmin(superUserToken);
} else {
this.admin = dataAPIClient.getAdmin();
}
this.namespaceAdmin = database.getDatabaseAdmin();
}
private void setApiEndpoint() {
@ -100,18 +121,33 @@ public class DataApiSpace {
String tokenFileContents = null;
Optional<String> tokenFilePath = config.getOptional("astraTokenFile");
if (tokenFilePath.isPresent()) {
Path path = Paths.get(tokenFilePath.get());
try {
tokenFileContents = Files.readAllLines(path).getFirst();
} catch (IOException e) {
String error = "Error while reading token from file:" + path;
logger.error(error, e);
throw new RuntimeException(e);
}
tokenFileContents = getTokenFileContents(tokenFilePath.get());
}
this.astraToken = (tokenFileContents != null) ? tokenFileContents : config.get("astraToken");
}
private String getTokenFileContents(String filePath) {
Path path = Paths.get(filePath);
try {
return Files.readAllLines(path).getFirst();
} catch (IOException e) {
String error = "Error while reading token from file:" + path;
logger.error(error, e);
throw new RuntimeException(e);
}
}
private void setSuperToken() {
String superTokenFileContents = null;
Optional<String> superTokenFilePath = config.getOptional("superTokenFile");
if (superTokenFilePath.isPresent()) {
superTokenFileContents = getTokenFileContents(superTokenFilePath.get());
}
Optional<String> maybeSuperToken = config.getOptional("superToken");
// It's fine if this is null
this.superUserToken = maybeSuperToken.orElse(superTokenFileContents);
}
public static NBConfigModel getConfigModel() {
return ConfigModel.of(DataApiSpace.class)
.add(
@ -136,6 +172,14 @@ public class DataApiSpace {
.setDescription("The Astra namespace to use")
)
.add(
Param.optional("superTokenFile", String.class)
.setDescription("optional file to load Astra admin user token from")
)
.add(
Param.optional("superToken", String.class)
.setDescription("optional Astra token used to connect as Admin user")
)
.asReadOnly();
}

View File

@ -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);
}
}

View File

@ -39,26 +39,12 @@ public class DataApiCreateCollectionOpDispenser extends DataApiOpDispenser {
private LongFunction<DataApiCreateCollectionOp> createOpFunction(ParsedOp op) {
return (l) -> {
CollectionOptions.CollectionOptionsBuilder optionsBldr = CollectionOptions.builder();
Optional<LongFunction<Integer>> dimFunc = op.getAsOptionalFunction("dimensions", Integer.class);
if (dimFunc.isPresent()) {
LongFunction<Integer> af = dimFunc.get();
optionsBldr.vectorDimension(af.apply(l));
}
// COSINE("cosine"),
// EUCLIDEAN("euclidean"),
// DOT_PRODUCT("dot_product");
Optional<LongFunction<String>> simFunc = op.getAsOptionalFunction("similarity", String.class);
if (simFunc.isPresent()) {
LongFunction<String> sf = simFunc.get();
optionsBldr.vectorSimilarity(SimilarityMetric.fromValue(sf.apply(l)));
}
DataApiCreateCollectionOp dataApiCreateCollectionOp =
new DataApiCreateCollectionOp(
spaceFunction.apply(l).getDatabase(),
targetFunction.apply(l),
optionsBldr.build());
this.getCollectionOptionsFromOp(op, l)
);
return dataApiCreateCollectionOp;
};

View File

@ -0,0 +1,61 @@
/*
* 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 io.nosqlbench.adapter.dataapi.DataApiDriverAdapter;
import io.nosqlbench.adapter.dataapi.ops.DataApiBaseOp;
import io.nosqlbench.adapter.dataapi.ops.DataApiCreateCollectionWithClassOp;
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 DataApiCreateCollectionWithClassOpDispenser extends DataApiOpDispenser {
private static final Logger logger = LogManager.getLogger(DataApiCreateCollectionWithClassOpDispenser.class);
private final LongFunction<DataApiCreateCollectionWithClassOp> opFunction;
public DataApiCreateCollectionWithClassOpDispenser(DataApiDriverAdapter adapter, ParsedOp op, LongFunction<String> targetFunction) {
super(adapter, op, targetFunction);
this.opFunction = createOpFunction(op);
}
private LongFunction<DataApiCreateCollectionWithClassOp> createOpFunction(ParsedOp op) {
return (l) -> new DataApiCreateCollectionWithClassOp(
spaceFunction.apply(l).getDatabase(),
targetFunction.apply(l),
this.getCollectionOptionsFromOp(op, l),
getCreateClass(op, l)
);
}
private Class<?> getCreateClass(ParsedOp op, long l) {
String className = op.getAsFunctionOr("createClass", "com.datastax.astra.client.model.Document").apply(l);
try {
return Class.forName(className);
} catch (ClassNotFoundException e) {
throw new RuntimeException(e);
}
}
@Override
public DataApiBaseOp getOp(long value) {
return opFunction.apply(value);
}
}

View File

@ -0,0 +1,67 @@
/*
* 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.dtsx.astra.sdk.db.domain.CloudProviderType;
import io.nosqlbench.adapter.dataapi.DataApiDriverAdapter;
import io.nosqlbench.adapter.dataapi.ops.DataApiBaseOp;
import io.nosqlbench.adapter.dataapi.ops.DataApiCreateDatabaseOp;
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 DataApiCreateDatabaseOpDispenser extends DataApiOpDispenser {
private static final Logger logger = LogManager.getLogger(DataApiCreateDatabaseOpDispenser.class);
private final LongFunction<DataApiCreateDatabaseOp> opFunction;
public DataApiCreateDatabaseOpDispenser(DataApiDriverAdapter adapter, ParsedOp op, LongFunction<String> targetFunction) {
super(adapter, op, targetFunction);
this.opFunction = createOpFunction(op);
}
private LongFunction<DataApiCreateDatabaseOp> createOpFunction(ParsedOp op) {
return (l) -> new DataApiCreateDatabaseOp(
spaceFunction.apply(l).getDatabase(),
spaceFunction.apply(l).getAdmin(),
targetFunction.apply(l),
getCloudProvider(op, l),
getRegion(op, l),
getWaitUntilActive(op, l)
);
}
/*
* These default to the same values used in the API if only the name is provided.
*/
private CloudProviderType getCloudProvider(ParsedOp op, long l) {
return CloudProviderType.valueOf(op.getAsFunctionOr("cloudProvider", "GCP").apply(l));
}
private String getRegion(ParsedOp op, long l) {
return op.getAsFunctionOr("region", "us-east1").apply(l);
}
private boolean getWaitUntilActive(ParsedOp op, long l) {
return op.getAsFunctionOr("waitUntilActive", "true").apply(l).equals("true");
}
@Override
public DataApiBaseOp getOp(long value) {
return opFunction.apply(value);
}
}

View File

@ -0,0 +1,49 @@
/*
* 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 io.nosqlbench.adapter.dataapi.DataApiDriverAdapter;
import io.nosqlbench.adapter.dataapi.ops.DataApiBaseOp;
import io.nosqlbench.adapter.dataapi.ops.DataApiCreateNamespaceOp;
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 DataApiCreateNamespaceOpDispenser extends DataApiOpDispenser {
private static final Logger logger = LogManager.getLogger(DataApiCreateNamespaceOpDispenser.class);
private final LongFunction<DataApiCreateNamespaceOp> opFunction;
public DataApiCreateNamespaceOpDispenser(DataApiDriverAdapter adapter, ParsedOp op, LongFunction<String> targetFunction) {
super(adapter, op, targetFunction);
this.opFunction = createOpFunction(op);
}
private LongFunction<DataApiCreateNamespaceOp> createOpFunction(ParsedOp op) {
return (l) -> new DataApiCreateNamespaceOp(
spaceFunction.apply(l).getDatabase(),
spaceFunction.apply(l).getNamespaceAdmin(),
targetFunction.apply(l)
);
}
@Override
public DataApiBaseOp getOp(long value) {
return opFunction.apply(value);
}
}

View File

@ -0,0 +1,53 @@
/*
* 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 io.nosqlbench.adapter.dataapi.DataApiDriverAdapter;
import io.nosqlbench.adapter.dataapi.ops.DataApiBaseOp;
import io.nosqlbench.adapter.dataapi.ops.DataApiDeleteAllOp;
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 DataApiDeleteAllOpDispenser extends DataApiOpDispenser {
private static final Logger logger = LogManager.getLogger(DataApiDeleteAllOpDispenser.class);
private final LongFunction<DataApiDeleteAllOp> opFunction;
public DataApiDeleteAllOpDispenser(DataApiDriverAdapter adapter, ParsedOp op, LongFunction<String> targetFunction) {
super(adapter, op, targetFunction);
this.opFunction = createOpFunction(op);
}
private LongFunction<DataApiDeleteAllOp> createOpFunction(ParsedOp op) {
return (l) -> {
Database db = spaceFunction.apply(l).getDatabase();
return new DataApiDeleteAllOp(
db,
db.getCollection(targetFunction.apply(l))
);
};
}
@Override
public DataApiBaseOp getOp(long value) {
return opFunction.apply(value);
}
}

View File

@ -0,0 +1,49 @@
/*
* 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 io.nosqlbench.adapter.dataapi.DataApiDriverAdapter;
import io.nosqlbench.adapter.dataapi.ops.DataApiBaseOp;
import io.nosqlbench.adapter.dataapi.ops.DataApiDropDatabaseOp;
import io.nosqlbench.adapter.dataapi.ops.DataApiListDatabasesOp;
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 DataApiDropDatabaseOpDispenser extends DataApiOpDispenser {
private static final Logger logger = LogManager.getLogger(DataApiDropDatabaseOpDispenser.class);
private final LongFunction<DataApiDropDatabaseOp> opFunction;
public DataApiDropDatabaseOpDispenser(DataApiDriverAdapter adapter, ParsedOp op, LongFunction<String> targetFunction) {
super(adapter, op, targetFunction);
this.opFunction = createOpFunction(op);
}
private LongFunction<DataApiDropDatabaseOp> createOpFunction(ParsedOp op) {
return (l) -> new DataApiDropDatabaseOp(
spaceFunction.apply(l).getDatabase(),
spaceFunction.apply(l).getAdmin(),
targetFunction.apply(l)
);
}
@Override
public DataApiBaseOp getOp(long value) {
return opFunction.apply(value);
}
}

View File

@ -0,0 +1,50 @@
/*
* 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 io.nosqlbench.adapter.dataapi.DataApiDriverAdapter;
import io.nosqlbench.adapter.dataapi.ops.DataApiBaseOp;
import io.nosqlbench.adapter.dataapi.ops.DataApiCreateNamespaceOp;
import io.nosqlbench.adapter.dataapi.ops.DataApiDropNamespaceOp;
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 DataApiDropNamespaceOpDispenser extends DataApiOpDispenser {
private static final Logger logger = LogManager.getLogger(DataApiDropNamespaceOpDispenser.class);
private final LongFunction<DataApiDropNamespaceOp> opFunction;
public DataApiDropNamespaceOpDispenser(DataApiDriverAdapter adapter, ParsedOp op, LongFunction<String> targetFunction) {
super(adapter, op, targetFunction);
this.opFunction = createOpFunction(op);
}
private LongFunction<DataApiDropNamespaceOp> createOpFunction(ParsedOp op) {
return (l) -> new DataApiDropNamespaceOp(
spaceFunction.apply(l).getDatabase(),
spaceFunction.apply(l).getNamespaceAdmin(),
targetFunction.apply(l)
);
}
@Override
public DataApiBaseOp getOp(long value) {
return opFunction.apply(value);
}
}

View File

@ -0,0 +1,62 @@
/*
* 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 com.datastax.astra.client.model.FindOneOptions;
import com.datastax.astra.client.model.Projection;
import com.datastax.astra.client.model.Sort;
import io.nosqlbench.adapter.dataapi.DataApiDriverAdapter;
import io.nosqlbench.adapter.dataapi.ops.DataApiBaseOp;
import io.nosqlbench.adapter.dataapi.ops.DataApiFindByIdOp;
import io.nosqlbench.adapter.dataapi.ops.DataApiFindOneOp;
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 DataApiFindByIdOpDispenser extends DataApiOpDispenser {
private static final Logger logger = LogManager.getLogger(DataApiFindByIdOpDispenser.class);
private final LongFunction<DataApiFindByIdOp> opFunction;
public DataApiFindByIdOpDispenser(DataApiDriverAdapter adapter, ParsedOp op, LongFunction<String> targetFunction) {
super(adapter, op, targetFunction);
this.opFunction = createOpFunction(op);
}
private LongFunction<DataApiFindByIdOp> createOpFunction(ParsedOp op) {
return (l) -> {
Database db = spaceFunction.apply(l).getDatabase();
Object id = getIdFromOp(op, l);
return new DataApiFindByIdOp(
db,
db.getCollection(targetFunction.apply(l)),
id
);
};
}
private Object getIdFromOp(ParsedOp op, long l) {
return op.getAsRequiredFunction("id", Object.class).apply(l);
}
@Override
public DataApiBaseOp getOp(long value) {
return opFunction.apply(value);
}
}

View File

@ -0,0 +1,67 @@
/*
* 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.DataApiFindDistinctOp;
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 DataApiFindDistinctOpDispenser extends DataApiOpDispenser {
private static final Logger logger = LogManager.getLogger(DataApiFindDistinctOpDispenser.class);
private final LongFunction<DataApiFindDistinctOp> opFunction;
public DataApiFindDistinctOpDispenser(DataApiDriverAdapter adapter, ParsedOp op, LongFunction<String> targetFunction) {
super(adapter, op, targetFunction);
this.opFunction = createOpFunction(op);
}
private LongFunction<DataApiFindDistinctOp> createOpFunction(ParsedOp op) {
return (l) -> {
Database db = spaceFunction.apply(l).getDatabase();
Filter filter = getFilterFromOp(op, l);
Class<?> targetClass = getTargetClass(op, l);
return new DataApiFindDistinctOp(
db,
db.getCollection(targetFunction.apply(l)),
op.getAsRequiredFunction("fieldName", String.class).apply(l),
filter,
targetClass
);
};
}
private Class<?> getTargetClass(ParsedOp op, long l) {
String className = op.getAsFunctionOr("resultClass", "java.lang.String").apply(l);
try {
return Class.forName(className);
} catch (ClassNotFoundException e) {
throw new RuntimeException(e);
}
}
@Override
public DataApiBaseOp getOp(long value) {
return opFunction.apply(value);
}
}

View File

@ -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);
}
}

View File

@ -28,6 +28,7 @@ import io.nosqlbench.adapters.api.templating.ParsedOp;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import java.util.Optional;
import java.util.function.LongFunction;
public class DataApiFindOpDispenser extends DataApiOpDispenser {
@ -67,7 +68,22 @@ public class DataApiFindOpDispenser extends DataApiOpDispenser {
if (projection != null) {
options = options.projection(projection);
}
options.setIncludeSimilarity(true);
Optional<LongFunction<Integer>> limitFunction = op.getAsOptionalFunction("limit", Integer.class);
if (limitFunction.isPresent()) {
options = options.limit(limitFunction.get().apply(l));
}
Optional<LongFunction<Integer>> skipFunction = op.getAsOptionalFunction("skip", Integer.class);
if (skipFunction.isPresent()) {
options = options.skip(skipFunction.get().apply(l));
}
Optional<LongFunction<Boolean>> includeSimilarityFunction = op.getAsOptionalFunction("includeSimilarity", Boolean.class);
if (includeSimilarityFunction.isPresent()) {
options.setIncludeSimilarity(includeSimilarityFunction.get().apply(l));
}
Optional<LongFunction<String>> pageStateFunction = op.getAsOptionalFunction("pageState", String.class);
if (pageStateFunction.isPresent()) {
options.setPageState(pageStateFunction.get().apply(l));
}
return options;
}

View File

@ -18,9 +18,6 @@ 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.FindOptions;
import com.datastax.astra.client.model.Projection;
import com.datastax.astra.client.model.Sort;
import io.nosqlbench.adapter.dataapi.DataApiDriverAdapter;
import io.nosqlbench.adapter.dataapi.ops.DataApiBaseOp;
import io.nosqlbench.adapter.dataapi.ops.DataApiFindVectorFilterOp;
@ -58,23 +55,6 @@ public class DataApiFindVectorFilterOpDispenser extends DataApiOpDispenser {
return op.getConfigOr("limit", 100, l);
}
private FindOptions getFindOptions(ParsedOp op, long l) {
FindOptions options = new FindOptions();
Sort sort = getSortFromOp(op, l);
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);
}
options.setIncludeSimilarity(true);
return options;
}
@Override
public DataApiBaseOp getOp(long value) {
return opFunction.apply(value);

View File

@ -55,23 +55,6 @@ public class DataApiFindVectorOpDispenser extends DataApiOpDispenser {
return op.getConfigOr("limit", 100, l);
}
private FindOptions getFindOptions(ParsedOp op, long l) {
FindOptions options = new FindOptions();
Sort sort = getSortFromOp(op, l);
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);
}
options.setIncludeSimilarity(true);
return options;
}
@Override
public DataApiBaseOp getOp(long value) {
return opFunction.apply(value);

View File

@ -0,0 +1,50 @@
/*
* 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 io.nosqlbench.adapter.dataapi.DataApiDriverAdapter;
import io.nosqlbench.adapter.dataapi.ops.DataApiBaseOp;
import io.nosqlbench.adapter.dataapi.ops.DataApiGetDatabaseInfoOp;
import io.nosqlbench.adapters.api.templating.ParsedOp;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import java.util.UUID;
import java.util.function.LongFunction;
public class DataApiGetDatabaseInfoOpDispenser extends DataApiOpDispenser {
private static final Logger logger = LogManager.getLogger(DataApiGetDatabaseInfoOpDispenser.class);
private final LongFunction<DataApiGetDatabaseInfoOp> opFunction;
public DataApiGetDatabaseInfoOpDispenser(DataApiDriverAdapter adapter, ParsedOp op, LongFunction<String> targetFunction) {
super(adapter, op, targetFunction);
this.opFunction = createOpFunction(op);
}
private LongFunction<DataApiGetDatabaseInfoOp> createOpFunction(ParsedOp op) {
return (l) -> new DataApiGetDatabaseInfoOp(
spaceFunction.apply(l).getDatabase(),
spaceFunction.apply(l).getAdmin(),
UUID.fromString(targetFunction.apply(l))
);
}
@Override
public DataApiBaseOp getOp(long value) {
return opFunction.apply(value);
}
}

View File

@ -66,6 +66,8 @@ public class DataApiInsertManyOpDispenser extends DataApiOpDispenser {
options = options.concurrency(Integer.parseInt(entry.getValue()));
case "ordered" ->
options = options.ordered(Boolean.parseBoolean(entry.getValue()));
case "timeout" ->
options = options.timeout(Integer.parseInt(entry.getValue()));
}
}
}

View File

@ -0,0 +1,47 @@
/*
* 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 io.nosqlbench.adapter.dataapi.DataApiDriverAdapter;
import io.nosqlbench.adapter.dataapi.ops.DataApiBaseOp;
import io.nosqlbench.adapter.dataapi.ops.DataApiListDatabasesOp;
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 DataApiListDatabasesOpDispenser extends DataApiOpDispenser {
private static final Logger logger = LogManager.getLogger(DataApiListDatabasesOpDispenser.class);
private final LongFunction<DataApiListDatabasesOp> opFunction;
public DataApiListDatabasesOpDispenser(DataApiDriverAdapter adapter, ParsedOp op, LongFunction<String> targetFunction) {
super(adapter, op, targetFunction);
this.opFunction = createOpFunction(op);
}
private LongFunction<DataApiListDatabasesOp> createOpFunction(ParsedOp op) {
return (l) -> new DataApiListDatabasesOp(
spaceFunction.apply(l).getDatabase(),
spaceFunction.apply(l).getAdmin()
);
}
@Override
public DataApiBaseOp getOp(long value) {
return opFunction.apply(value);
}
}

View File

@ -0,0 +1,48 @@
/*
* 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 io.nosqlbench.adapter.dataapi.DataApiDriverAdapter;
import io.nosqlbench.adapter.dataapi.ops.DataApiBaseOp;
import io.nosqlbench.adapter.dataapi.ops.DataApiListNamespacesOp;
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 DataApiListNamespacesOpDispenser extends DataApiOpDispenser {
private static final Logger logger = LogManager.getLogger(DataApiListNamespacesOpDispenser.class);
private final LongFunction<DataApiListNamespacesOp> opFunction;
public DataApiListNamespacesOpDispenser(DataApiDriverAdapter adapter, ParsedOp op, LongFunction<String> targetFunction) {
super(adapter, op, targetFunction);
this.opFunction = createOpFunction(op);
}
private LongFunction<DataApiListNamespacesOp> createOpFunction(ParsedOp op) {
return (l) -> new DataApiListNamespacesOp(
spaceFunction.apply(l).getDatabase(),
spaceFunction.apply(l).getNamespaceAdmin()
);
}
@Override
public DataApiBaseOp getOp(long value) {
return opFunction.apply(value);
}
}

View File

@ -23,10 +23,7 @@ import io.nosqlbench.adapters.api.activityimpl.BaseOpDispenser;
import io.nosqlbench.adapters.api.activityimpl.uniform.DriverAdapter;
import io.nosqlbench.adapters.api.templating.ParsedOp;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.*;
import java.util.function.LongFunction;
public abstract class DataApiOpDispenser extends BaseOpDispenser<DataApiBaseOp, DataApiSpace> {
@ -188,4 +185,45 @@ public abstract class DataApiOpDispenser extends BaseOpDispenser<DataApiBaseOp,
return projection;
}
protected CollectionOptions getCollectionOptionsFromOp(ParsedOp op, long l) {
CollectionOptions.CollectionOptionsBuilder optionsBldr = CollectionOptions.builder();
Optional<LongFunction<Integer>> dimFunc = op.getAsOptionalFunction("dimensions", Integer.class);
if (dimFunc.isPresent()) {
LongFunction<Integer> af = dimFunc.get();
optionsBldr = optionsBldr.vectorDimension(af.apply(l));
}
Optional<LongFunction<String>> simFunc = op.getAsOptionalFunction("similarity", String.class);
if (simFunc.isPresent()) {
LongFunction<String> sf = simFunc.get();
optionsBldr = optionsBldr.vectorSimilarity(SimilarityMetric.fromValue(sf.apply(l)));
}
Optional<LongFunction<String>> typeFunc = op.getAsOptionalFunction("collectionType", String.class);
if (typeFunc.isPresent()) {
LongFunction<String> tf = typeFunc.get();
optionsBldr = optionsBldr.defaultIdType(CollectionIdTypes.fromValue(tf.apply(l)));
}
Optional<LongFunction<String>> providerFunc = op.getAsOptionalFunction("serviceProvider", String.class);
Optional<LongFunction<String>> modeFunc = op.getAsOptionalFunction("serviceMode", String.class);
Optional<LongFunction<Map>> paramFunc = op.getAsOptionalFunction("serviceParameters", Map.class);
if (providerFunc.isPresent() && modeFunc.isPresent()) {
LongFunction<String> pf = providerFunc.get();
LongFunction<String> mf = modeFunc.get();
optionsBldr = paramFunc.isPresent() ?
optionsBldr.vectorize(pf.apply(l), mf.apply(l), paramFunc.get().apply(l)) :
optionsBldr.vectorize(pf.apply(l), mf.apply(l));
}
Optional<LongFunction<List>> allowFunc = op.getAsOptionalFunction("allowIndex", List.class);
if (allowFunc.isPresent()) {
LongFunction<List> af = allowFunc.get();
optionsBldr = optionsBldr.indexingAllow(Arrays.toString(af.apply(l).toArray(new String[0])));
}
Optional<LongFunction<List>> denyFunc = op.getAsOptionalFunction("denyIndex", List.class);
if (denyFunc.isPresent()) {
LongFunction<List> df = denyFunc.get();
optionsBldr = optionsBldr.indexingDeny(Arrays.toString(df.apply(l).toArray(new String[0])));
}
return optionsBldr.build();
}
}

View File

@ -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);
}
}

View File

@ -16,19 +16,57 @@
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.DataApiUpdateManyOp;
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 DataApiUpdateManyOpDispenser extends DataApiOpDispenser {
private static final Logger logger = LogManager.getLogger(DataApiUpdateManyOpDispenser.class);
private final LongFunction<DataApiUpdateManyOp> opFunction;
public DataApiUpdateManyOpDispenser(DataApiDriverAdapter adapter, ParsedOp op, LongFunction<String> targetFunction) {
super(adapter, op, targetFunction);
this.opFunction = createOpFunction(op);
}
private LongFunction<DataApiUpdateManyOp> createOpFunction(ParsedOp op) {
return (l) -> {
Database db = spaceFunction.apply(l).getDatabase();
Filter filter = getFilterFromOp(op, l);
UpdateManyOptions options = getUpdateManyOptions(op, l);
LongFunction<Map> docMapFunc = op.getAsRequiredFunction("updates", Map.class);
return new DataApiUpdateManyOp(
db,
db.getCollection(targetFunction.apply(l)),
filter,
new Update(docMapFunc.apply(l)),
options
);
};
}
private UpdateManyOptions getUpdateManyOptions(ParsedOp op, long l) {
UpdateManyOptions options = new UpdateManyOptions();
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 null;
return opFunction.apply(value);
}
}

View File

@ -16,19 +16,66 @@
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.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 DataApiUpdateOneOpDispenser extends DataApiOpDispenser {
private static final Logger logger = LogManager.getLogger(DataApiUpdateOneOpDispenser.class);
private final LongFunction<DataApiUpdateOneOp> opFunction;
public DataApiUpdateOneOpDispenser(DataApiDriverAdapter adapter, ParsedOp op, LongFunction<String> targetFunction) {
super(adapter, op, targetFunction);
this.opFunction = createOpFunction(op);
}
private LongFunction<DataApiUpdateOneOp> createOpFunction(ParsedOp op) {
return (l) -> {
Database db = spaceFunction.apply(l).getDatabase();
Filter filter = getFilterFromOp(op, l);
UpdateOneOptions options = getUpdateOneOptions(op, l);
LongFunction<Map> docMapFunc = op.getAsRequiredFunction("update", Map.class);
return new DataApiUpdateOneOp(
db,
db.getCollection(targetFunction.apply(l)),
filter,
new Update(docMapFunc.apply(l)),
options
);
};
}
private UpdateOneOptions getUpdateOneOptions(ParsedOp op, long l) {
UpdateOneOptions options = new UpdateOneOptions();
Sort sort = getSortFromOp(op, l);
float[] vector = getVectorFromOp(op, l);
Optional<LongFunction<Boolean>> upsertFunction = op.getAsOptionalFunction("upsert", Boolean.class);
if (upsertFunction.isPresent()) {
options = options.upsert(upsertFunction.get().apply(l));
}
if (sort != null) {
options = (vector != null) ? options.vector(vector, sort) : options.sort(sort);
}
return options;
}
private float[] getVectorFromOp(ParsedOp op, long l) {
return getVectorValues(op.get("vector", l));
}
@Override
public DataApiBaseOp getOp(long value) {
return null;
return opFunction.apply(value);
}
}

View File

@ -0,0 +1,33 @@
/*
* 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.admin.AstraDBAdmin;
import io.nosqlbench.adapters.api.activityimpl.uniform.flowtypes.CycleOp;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
public abstract class DataApiAdminOp extends DataApiBaseOp {
protected static final Logger logger = LogManager.getLogger(DataApiAdminOp.class);
protected final AstraDBAdmin admin;
public DataApiAdminOp(Database db, AstraDBAdmin admin) {
super(db);
this.admin = admin;
}
}

View File

@ -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);
}
}
}

View File

@ -0,0 +1,38 @@
/*
* 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.CollectionOptions;
public class DataApiCreateCollectionWithClassOp extends DataApiBaseOp {
private final String collectionName;
private final CollectionOptions options;
private final Class<?> clazz;
public DataApiCreateCollectionWithClassOp(Database db, String collectionName, CollectionOptions options, Class<?> clazz) {
super(db);
this.collectionName = collectionName;
this.options = options;
this.clazz = clazz;
}
@Override
public Object apply(long value) {
return db.createCollection(collectionName, options, clazz);
}
}

View File

@ -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.admin.AstraDBAdmin;
import com.dtsx.astra.sdk.db.domain.CloudProviderType;
public class DataApiCreateDatabaseOp extends DataApiAdminOp {
private final String name;
private final CloudProviderType cloud;
private final String cloudRegion;
private final boolean waitForDb;
public DataApiCreateDatabaseOp(Database db, AstraDBAdmin admin, String name, CloudProviderType cloud, String cloudRegion, boolean waitForDb) {
super(db, admin);
this.name = name;
this.cloud = cloud;
this.cloudRegion = cloudRegion;
this.waitForDb = waitForDb;
}
@Override
public Object apply(long value) {
return admin.createDatabase(name, cloud, cloudRegion, waitForDb);
}
}

View File

@ -0,0 +1,36 @@
/*
* 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.admin.AstraDBAdmin;
import com.datastax.astra.client.admin.DatabaseAdmin;
public class DataApiCreateNamespaceOp extends DataApiDBAdminOp {
private final String namespace;
public DataApiCreateNamespaceOp(Database db, DatabaseAdmin admin, String namespace) {
super(db, admin);
this.namespace = namespace;
}
@Override
public Object apply(long value) {
admin.createNamespace(namespace);
return true;
}
}

View File

@ -0,0 +1,33 @@
/*
* 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.admin.AstraDBAdmin;
import com.datastax.astra.client.admin.DatabaseAdmin;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
public abstract class DataApiDBAdminOp extends DataApiBaseOp {
protected static final Logger logger = LogManager.getLogger(DataApiDBAdminOp.class);
protected final DatabaseAdmin admin;
public DataApiDBAdminOp(Database db, DatabaseAdmin admin) {
super(db);
this.admin = admin;
}
}

View File

@ -0,0 +1,35 @@
/*
* 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;
public class DataApiDeleteAllOp extends DataApiBaseOp {
private final Collection collection;
public DataApiDeleteAllOp(Database db, Collection collection) {
super(db);
this.collection = collection;
}
@Override
public Object apply(long value) {
return collection.deleteAll();
}
}

View File

@ -0,0 +1,41 @@
/*
* 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.admin.AstraDBAdmin;
import java.util.UUID;
import java.util.regex.Pattern;
public class DataApiDropDatabaseOp extends DataApiAdminOp {
private static final Pattern UUID_PATTERN = Pattern.compile(
"^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}$");
private final String nameOrUUID;
private final boolean isUUID;
public DataApiDropDatabaseOp(Database db, AstraDBAdmin admin, String nameOrUUID) {
super(db, admin);
this.nameOrUUID = nameOrUUID;
this.isUUID = UUID_PATTERN.matcher(nameOrUUID).matches();
}
@Override
public Object apply(long value) {
return isUUID ? admin.dropDatabase(UUID.fromString(nameOrUUID)) : admin.dropDatabase(nameOrUUID);
}
}

View File

@ -0,0 +1,35 @@
/*
* 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.admin.DatabaseAdmin;
public class DataApiDropNamespaceOp extends DataApiDBAdminOp {
private final String namespace;
public DataApiDropNamespaceOp(Database db, DatabaseAdmin admin, String namespace) {
super(db, admin);
this.namespace = namespace;
}
@Override
public Object apply(long value) {
admin.dropNamespace(namespace);
return true;
}
}

View File

@ -0,0 +1,38 @@
/*
* 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.Document;
import com.datastax.astra.client.model.Filter;
public class DataApiFindByIdOp extends DataApiBaseOp {
private final Collection<Document> collection;
private final Object id;
public DataApiFindByIdOp(Database db, Collection<Document> collection, Object id) {
super(db);
this.collection = collection;
this.id = id;
}
@Override
public Object apply(long value) {
return collection.findById(id);
}
}

View File

@ -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.CollectionOptions;
import com.datastax.astra.client.model.Filter;
public class DataApiFindDistinctOp extends DataApiBaseOp {
private final Collection collection;
private final String fieldName;
private final Filter filter;
private final Class<?> resultClass;
public DataApiFindDistinctOp(Database db, Collection collection, String fieldName, Filter filter, Class<?> resultClass) {
super(db);
this.collection = collection;
this.fieldName = fieldName;
this.filter = filter;
this.resultClass = resultClass;
}
@Override
public Object apply(long value) {
return collection.distinct(fieldName, filter, resultClass);
}
}

View File

@ -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);
}
}

View File

@ -0,0 +1,37 @@
/*
* 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.admin.AstraDBAdmin;
import java.util.UUID;
import java.util.regex.Pattern;
public class DataApiGetDatabaseInfoOp extends DataApiAdminOp {
private final UUID uuid;
public DataApiGetDatabaseInfoOp(Database db, AstraDBAdmin admin, UUID uuid) {
super(db, admin);
this.uuid = uuid;
}
@Override
public Object apply(long value) {
return admin.getDatabaseInfo(uuid);
}
}

View File

@ -0,0 +1,33 @@
/*
* 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.admin.AstraDBAdmin;
import com.dtsx.astra.sdk.db.domain.CloudProviderType;
public class DataApiListDatabasesOp extends DataApiAdminOp {
public DataApiListDatabasesOp(Database db, AstraDBAdmin admin) {
super(db, admin);
}
@Override
public Object apply(long value) {
return admin.listDatabases();
}
}

View File

@ -0,0 +1,32 @@
/*
* 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.admin.DatabaseAdmin;
public class DataApiListNamespacesOp extends DataApiDBAdminOp {
public DataApiListNamespacesOp(Database db, DatabaseAdmin admin) {
super(db, admin);
}
@Override
public Object apply(long value) {
return admin.listNamespaceNames();
}
}

View File

@ -35,4 +35,18 @@ public enum DataApiOpType {
list_collections,
list_collection_names,
estimated_document_count,
find_by_id,
find_distinct,
count_documents,
replace_one,
find_one_and_replace,
delete_all,
create_collection_with_class,
create_database,
list_databases,
drop_database,
get_database_info,
create_namespace,
list_namespaces,
drop_namespace,
}

View File

@ -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);
}
}

View File

@ -0,0 +1,43 @@
/*
* 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.Update;
import com.datastax.astra.client.model.UpdateManyOptions;
public class DataApiUpdateManyOp extends DataApiBaseOp {
private final Collection collection;
private final Filter filter;
private final Update update;
private final UpdateManyOptions options;
public DataApiUpdateManyOp(Database db, Collection collection, Filter filter, Update update, UpdateManyOptions options) {
super(db);
this.collection = collection;
this.filter = filter;
this.update = update;
this.options = options;
}
@Override
public Object apply(long value) {
return collection.updateMany(filter, update, options);
}
}

View File

@ -0,0 +1,43 @@
/*
* 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.Update;
import com.datastax.astra.client.model.UpdateOneOptions;
public class DataApiUpdateOneOp extends DataApiBaseOp {
private final Collection collection;
private final Filter filter;
private final Update update;
private final UpdateOneOptions options;
public DataApiUpdateOneOp(Database db, Collection collection, Filter filter, Update update, UpdateOneOptions options) {
super(db);
this.collection = collection;
this.filter = filter;
this.update = update;
this.options = options;
}
@Override
public Object apply(long value) {
return collection.updateOne(filter, update, options);
}
}