IntenseFinance/nosqlbench
3
0
Fork 0
mirror of https://github.com/nosqlbench/nosqlbench.git synced 2025-02-25 18:55:28 -06:00
Code Issues Packages Projects Releases Wiki Activity

Spanner - Handle Drop DB more elegantly

Browse Source
This commit is contained in:
Madhavan Sridharan
2024-10-09 13:11:41 -04:00
parent 941b623f63
commit 18e432cffd
7 changed files with 43 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

36
nb-adapters/adapter-gcp-spanner/src/main/java/io/nosqlbench/adapter/gcpspanner/opdispensers/GCPSpannerDropDatabaseDdlOpDispenser.java
View File

@@ -17,7 +17,10 @@
package io.nosqlbench.adapter.gcpspanner.opdispensers;
import com.google.cloud.spanner.DatabaseAdminClient;
import com.google.cloud.spanner.Options;
import io.nosqlbench.adapter.gcpspanner.GCPSpannerDriverAdapter;
import io.nosqlbench.adapter.gcpspanner.GCPSpannerSpace;
import io.nosqlbench.adapter.gcpspanner.ops.GCPSpannerBaseOp;
import io.nosqlbench.adapter.gcpspanner.ops.GCPSpannerDropDatabaseDdlOp;
import io.nosqlbench.adapters.api.templating.ParsedOp;
@@ -45,7 +48,7 @@ public class GCPSpannerDropDatabaseDdlOpDispenser extends GCPSpannerBaseOpDispen
*/
public GCPSpannerDropDatabaseDdlOpDispenser(GCPSpannerDriverAdapter adapter, ParsedOp op, LongFunction<String> targetFunction) {
super(adapter, op, targetFunction);
this.opFunction = DropOpFunction(op);
this.opFunction = dropOpFunction(op);
}
/**
@@ -54,14 +57,27 @@ public class GCPSpannerDropDatabaseDdlOpDispenser extends GCPSpannerBaseOpDispen
* @param op the {@link ParsedOp} instance
* @return a {@link LongFunction} that generates {@link GCPSpannerDropDatabaseDdlOp} instances
*/
private LongFunction<GCPSpannerDropDatabaseDdlOp> DropOpFunction(ParsedOp op) {
return (l) -> new GCPSpannerDropDatabaseDdlOp(
spaceFunction.apply(l).getSpanner(),
l,
targetFunction.apply(l),
spaceFunction.apply(l).getDbAdminClient(),
spaceFunction.apply(l).getDbAdminClient().getDatabase(spaceFunction.apply(l).getInstanceId(), spaceFunction.apply(l).getDatabaseIdString())
);
private LongFunction<GCPSpannerDropDatabaseDdlOp> dropOpFunction(ParsedOp op) {
return l -> {
GCPSpannerSpace space = spaceFunction.apply(l);
if (space.getSpanner() != null && space.getDbAdminClient() != null && space.getDbClient() != null) {
String dbId = targetFunction.apply(l);
DatabaseAdminClient dbAdminClient = space.getDbAdminClient();
if (dbAdminClient.listDatabases(spaceFunction.apply(l).getInstanceId(), Options.filter(null))
.streamValues().anyMatch(d -> d.getId().getDatabase().equals(dbId))) {
return new GCPSpannerDropDatabaseDdlOp(
space.getSpanner(),
l,
dbId,
space.getDbAdminClient(),
space.getDbAdminClient().getDatabase(space.getInstanceId(), space.getDatabaseIdString())
);
}
}
logger.info("Database to be deleted not found [{}]", space.getDatabaseId().getName());
return null;
};
}
/**
@@ -72,6 +88,6 @@ public class GCPSpannerDropDatabaseDdlOpDispenser extends GCPSpannerBaseOpDispen
*/
@Override
public GCPSpannerBaseOp<?> getOp(long value) {
return opFunction.apply(value);
return opFunction != null ? opFunction.apply(value) : null;
}
}

11
nb-adapters/adapter-gcp-spanner/src/main/java/io/nosqlbench/adapter/gcpspanner/ops/GCPSpannerDropDatabaseDdlOp.java
View File

@@ -18,6 +18,7 @@ package io.nosqlbench.adapter.gcpspanner.ops;
import com.google.cloud.spanner.Database;
import com.google.cloud.spanner.DatabaseAdminClient;
import com.google.cloud.spanner.DatabaseNotFoundException;
import com.google.cloud.spanner.Spanner;
/**
@@ -56,11 +57,17 @@ public class GCPSpannerDropDatabaseDdlOp extends GCPSpannerBaseOp<Long> {
@Override
public Object applyOp(long value) {
try {
db.drop();
if (null != db && db.exists()) {
db.drop();
}
} catch (Exception e) {
logger.warn("Error dropping database using the Database object: {}. Will re-try using the DBAdminClient now...", e.getMessage());
try {
dbAdminClient.dropDatabase(databaseId, null);
if (null != dbAdminClient) {
dbAdminClient.dropDatabase(databaseId, null);
}
} catch (DatabaseNotFoundException noDB) {
logger.info("Database does not exist. {}", noDB.getMessage());
} catch (Exception ex) {
throw new RuntimeException(ex);
}