adding administrative function ops

This commit is contained in:
Mark Wolters 2024-06-14 15:17:11 -04:00
parent f772af9d0e
commit f1ef3a24c1
10 changed files with 355 additions and 8 deletions

View File

@ -71,6 +71,9 @@ public class DataApiOpMapper implements OpMapper<DataApiBaseOp> {
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);
};
}
}

View File

@ -18,6 +18,7 @@ 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 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 +43,14 @@ public class DataApiSpace {
private Database database;
private String namespace;
private String superUserToken;
private AstraDBAdmin admin;
public DataApiSpace(String name, NBConfiguration cfg) {
this.config = cfg;
this.name = name;
setToken();
setSuperToken();
setApiEndpoint();
setNamespace();
createClient();
@ -59,6 +64,10 @@ public class DataApiSpace {
return database;
}
public AstraDBAdmin getAdmin() {
return admin;
}
private void createClient() {
this.dataAPIClient = new DataAPIClient(astraToken);
if (namespace != null) {
@ -66,6 +75,11 @@ public class DataApiSpace {
} else {
this.database = dataAPIClient.getDatabase(astraApiEndpoint);
}
if (superUserToken != null) {
this.admin = dataAPIClient.getAdmin(superUserToken);
} else {
this.admin = dataAPIClient.getAdmin();
}
}
private void setApiEndpoint() {
@ -100,18 +114,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(

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.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,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,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,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,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,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

@ -42,4 +42,7 @@ public enum DataApiOpType {
find_one_and_replace,
delete_all,
create_collection_with_class,
create_database,
list_databases,
drop_database,
}