adding administrative function ops

nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/DataApiOpMapper.java

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

nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/DataApiSpace.java

@@ -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());
+            tokenFileContents = getTokenFileContents(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);
-            }
         }
         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(

nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiCreateDatabaseOpDispenser.java

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

nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiDropDatabaseOpDispenser.java

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

nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiListDatabasesOpDispenser.java

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

nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/ops/DataApiAdminOp.java

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

nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/ops/DataApiCreateDatabaseOp.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.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);
+    }
+}

nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/ops/DataApiDropDatabaseOp.java

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

nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/ops/DataApiListDatabasesOp.java

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

nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/ops/DataApiOpType.java

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