add ops for dml and ddl support

This commit is contained in:
Mark Wolters 2024-09-24 16:16:09 -04:00
parent 0cb84c27c3
commit 9d5824b650
10 changed files with 130 additions and 25 deletions

View File

@ -16,6 +16,12 @@
package io.nosqlbench.adapter.gcpspanner;
import com.google.cloud.spanner.ResultSet;
public class GCPSpannerAdapterUtils {
public static final String SPANNER = "gcp_spanner";
public static int[] getKeyArrayFromResultSet(ResultSet rs) {
return rs.getLongList(0).stream().mapToInt(Math::toIntExact).toArray();
}
}

View File

@ -55,10 +55,12 @@ public class GCPSpannerOpMapper implements OpMapper<GCPSpannerBaseOp<?>> {
logger.info(() -> "Using '" + typeAndTarget.enumId + "' op type for op template '" + op.getName() + "'");
return switch (typeAndTarget.enumId) {
case create_table ->
new GCPSpannerCreateTableOpDispenser(adapter, op, typeAndTarget.targetFunction);
case update_database_ddl ->
new GCPSpannerUpdateDatabaseDdlOpDispenser(adapter, op, typeAndTarget.targetFunction);
case insert_vector ->
new GCPSpannerInsertVectorOpDispenser(adapter, op, typeAndTarget.targetFunction);
case execute_dml ->
new GCPSpannerExecuteDmlOpDispenser(adapter, op, typeAndTarget.targetFunction);
};
}
}

View File

@ -113,7 +113,7 @@ public class GCPSpannerSpace implements AutoCloseable {
.asReadOnly();
}
@Override
public void close() throws Exception {
public void close() {
if (spanner != null) {
spanner.close();
}

View File

@ -0,0 +1,57 @@
/*
* 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.gcpspanner.opdispensers;
import com.google.cloud.spanner.Statement;
import io.nosqlbench.adapter.gcpspanner.GCPSpannerDriverAdapter;
import io.nosqlbench.adapter.gcpspanner.ops.GCPSpannerBaseOp;
import io.nosqlbench.adapter.gcpspanner.ops.GCPSpannerExecuteDmlOp;
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 GCPSpannerExecuteDmlOpDispenser extends GCPSpannerBaseOpDispenser {
private static final Logger logger = LogManager.getLogger(GCPSpannerExecuteDmlOpDispenser.class);
private final LongFunction<GCPSpannerExecuteDmlOp> opFunction;
public GCPSpannerExecuteDmlOpDispenser(GCPSpannerDriverAdapter adapter, ParsedOp op, LongFunction<String> targetFunction) {
super(adapter, op, targetFunction);
this.opFunction = createOpFunction(op);
}
private LongFunction<GCPSpannerExecuteDmlOp> createOpFunction(ParsedOp op) {
return (l) -> new GCPSpannerExecuteDmlOp(
spaceFunction.apply(l).getSpanner(),
l,
generateStatement(op.getAsRequiredFunction("DML", String.class).apply(l)),
spaceFunction.apply(l).getDbClient()
);
}
private Statement generateStatement(String dml) {
return Statement.of(dml);
}
@Override
public GCPSpannerBaseOp<?> getOp(long value) {
return opFunction.apply(value);
}
}

View File

@ -17,30 +17,27 @@
package io.nosqlbench.adapter.gcpspanner.opdispensers;
import com.google.cloud.spanner.Database;
import com.google.cloud.spanner.DatabaseAdminClient;
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.GCPSpannerCreateTableOp;
import io.nosqlbench.adapter.gcpspanner.ops.GCPSpannerUpdateDatabaseDdlOp;
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 GCPSpannerCreateTableOpDispenser extends GCPSpannerBaseOpDispenser {
private static final Logger logger = LogManager.getLogger(GCPSpannerCreateTableOpDispenser.class);
private final LongFunction<GCPSpannerCreateTableOp> opFunction;
public class GCPSpannerUpdateDatabaseDdlOpDispenser extends GCPSpannerBaseOpDispenser {
private static final Logger logger = LogManager.getLogger(GCPSpannerUpdateDatabaseDdlOpDispenser.class);
private final LongFunction<GCPSpannerUpdateDatabaseDdlOp> opFunction;
public GCPSpannerCreateTableOpDispenser(GCPSpannerDriverAdapter adapter, ParsedOp op, LongFunction<String> targetFunction) {
public GCPSpannerUpdateDatabaseDdlOpDispenser(GCPSpannerDriverAdapter adapter, ParsedOp op, LongFunction<String> targetFunction) {
super(adapter, op, targetFunction);
this.opFunction = createOpFunction(op);
}
private LongFunction<GCPSpannerCreateTableOp> createOpFunction(ParsedOp op) {
private LongFunction<GCPSpannerUpdateDatabaseDdlOp> createOpFunction(ParsedOp op) {
return (l) -> new GCPSpannerCreateTableOp(
return (l) -> new GCPSpannerUpdateDatabaseDdlOp(
spaceFunction.apply(l).getSpanner(),
l,
op.getAsRequiredFunction("DDL", String.class).apply(l),

View File

@ -37,12 +37,6 @@ public abstract class GCPSpannerBaseOp<T> implements CycleOp<Object> {
this.apiCall = this::applyOp;
}
public GCPSpannerBaseOp(Spanner spanner, T requestParam, LongFunction<Object> call) {
this.spannerClient = spanner;
this.request = requestParam;
this.apiCall = call;
}
@Override
public final Object apply(long value) {
logger.trace(() -> "applying op: " + this);

View File

@ -0,0 +1,39 @@
/*
* 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.gcpspanner.ops;
import com.google.cloud.spanner.*;
public class GCPSpannerExecuteDmlOp extends GCPSpannerBaseOp<Long> {
private final Statement statement;
private final DatabaseClient dbClient;
public GCPSpannerExecuteDmlOp(Spanner spanner, Long requestParam, Statement statement,
DatabaseClient dbClient) {
super(spanner, requestParam);
this.statement = statement;
this.dbClient = dbClient;
}
@Override
public Object applyOp(long value) {
try (ReadContext context = dbClient.singleUse()) {
return context.executeQuery(statement);
}
}
}

View File

@ -21,7 +21,6 @@ import com.google.cloud.spanner.Spanner;
import com.google.cloud.spanner.DatabaseClient;
import com.google.cloud.spanner.Mutation;
import java.util.Arrays;
import java.util.Collections;
public class GCPSpannerInsertVectorOp extends GCPSpannerBaseOp<Long> {

View File

@ -22,13 +22,13 @@ import com.google.cloud.spanner.*;
import com.google.common.collect.ImmutableList;
import com.google.spanner.admin.database.v1.UpdateDatabaseDdlMetadata;
public class GCPSpannerCreateTableOp extends GCPSpannerBaseOp<Long> {
public class GCPSpannerUpdateDatabaseDdlOp extends GCPSpannerBaseOp<Long> {
private final String createTableStatement;
private final DatabaseAdminClient dbAdminClient;
private final Database db;
public GCPSpannerCreateTableOp(Spanner searchIndexClient, Long requestParam, String createTableStatement,
DatabaseAdminClient dbAdminClient, Database db) {
public GCPSpannerUpdateDatabaseDdlOp(Spanner searchIndexClient, Long requestParam, String createTableStatement,
DatabaseAdminClient dbAdminClient, Database db) {
super(searchIndexClient, requestParam);
this.createTableStatement = createTableStatement;
this.dbAdminClient = dbAdminClient;

View File

@ -17,9 +17,20 @@
package io.nosqlbench.adapter.gcpspanner.types;
/**
* https://cloud.google.com/spanner/docs/reference/standard-sql/data-definition-language#vector_index_statements
* All the spanner rpc api calls are defined <a href="https://cloud.google.com/spanner/docs/reference/rpc">here</a>, representing a
* guide to the set of operations we should define if we want to implement full Spanner api support.
* <p>
* NOTE that the vector search functionality is still in pre-GA and is not available through rpc calls other than simply
* calling ExecuteSql. The SQL functionality related to vector indices is documented
* <a href="https://cloud.google.com/spanner/docs/reference/standard-sql/data-definition-language#vector_index_statements">here</a>
* <p>
* KNN and ANN search through Google SQL are documented respectively
* <a href="https://cloud.google.com/spanner/docs/find-k-nearest-neighbors">here</a>
* and
* <a href="https://cloud.google.com/spanner/docs/find-approximate-nearest-neighbors#query-vector-embeddings">here</a></a>
*/
public enum GCPSpannerOpType {
create_table,
update_database_ddl,
insert_vector,
execute_dml,
}