Merge branch 'main' into snyk-upgrade-51d0a4e16d574c4403f48d70cffd520f

This commit is contained in:
Jonathan Shook 2024-10-01 15:29:53 -05:00 committed by GitHub
commit 9d8cde2f7e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
24 changed files with 1106 additions and 8 deletions

View File

@ -415,6 +415,12 @@
<version>1.13.2</version>
</dependency>
<dependency>
<groupId>com.google.cloud</groupId>
<artifactId>google-cloud-spanner</artifactId>
<version>6.76.0</version>
</dependency>
</dependencies>
</dependencyManagement>

View File

@ -0,0 +1,55 @@
<!--
~ Copyright (c) 2020-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.
-->
<!--<?xml version="1.0" encoding="UTF-8"?>-->
<project xmlns="http://maven.apache.org/POM/4.0.0"
xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 https://maven.apache.org/xsd/maven-4.0.0.xsd">
<modelVersion>4.0.0</modelVersion>
<artifactId>adapter-gcp-spanner</artifactId>
<packaging>jar</packaging>
<parent>
<artifactId>mvn-defaults</artifactId>
<groupId>io.nosqlbench</groupId>
<version>${revision}</version>
<relativePath>../../mvn-defaults</relativePath>
</parent>
<name>${project.artifactId}</name>
<description>
A NoSQLBbench driver adapter module for the Google Spanner database.
</description>
<dependencies>
<dependency>
<groupId>io.nosqlbench</groupId>
<artifactId>nb-annotations</artifactId>
<version>${revision}</version>
<scope>compile</scope>
</dependency>
<dependency>
<groupId>io.nosqlbench</groupId>
<artifactId>adapters-api</artifactId>
<version>${revision}</version>
<scope>compile</scope>
</dependency>
<dependency>
<groupId>com.google.cloud</groupId>
<artifactId>google-cloud-spanner</artifactId>
</dependency>
</dependencies>
</project>

View File

@ -0,0 +1,34 @@
/*
* Copyright (c) 2020-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;
import com.google.cloud.spanner.ResultSet;
import java.util.ArrayList;
import java.util.List;
public class GCPSpannerAdapterUtils {
public static final String SPANNER = "spanner";
public static int[] getKeyArrayFromResultSet(ResultSet rs) {
List<Integer> values = new ArrayList<>();
while(rs.next()) {
values.add(Integer.valueOf(rs.getString(0)));
}
return values.stream().mapToInt(i -> i).toArray();
}
}

View File

@ -0,0 +1,55 @@
/*
* Copyright (c) 2020-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;
import io.nosqlbench.adapter.gcpspanner.ops.GCPSpannerBaseOp;
import io.nosqlbench.adapters.api.activityimpl.OpMapper;
import io.nosqlbench.adapters.api.activityimpl.uniform.BaseDriverAdapter;
import io.nosqlbench.adapters.api.activityimpl.uniform.DriverAdapter;
import io.nosqlbench.nb.annotations.Service;
import io.nosqlbench.nb.api.components.core.NBComponent;
import io.nosqlbench.nb.api.config.standard.NBConfigModel;
import io.nosqlbench.nb.api.config.standard.NBConfiguration;
import io.nosqlbench.nb.api.labels.NBLabels;
import java.util.function.Function;
import static io.nosqlbench.adapter.gcpspanner.GCPSpannerAdapterUtils.SPANNER;
@Service(value = DriverAdapter.class, selector = SPANNER)
public class GCPSpannerDriverAdapter extends BaseDriverAdapter<GCPSpannerBaseOp<?>, GCPSpannerSpace> {
public GCPSpannerDriverAdapter(NBComponent parentComponent, NBLabels labels) {
super(parentComponent, labels);
}
@Override
public OpMapper<GCPSpannerBaseOp<?>> getOpMapper() {
return new GCPSpannerOpMapper(this);
}
@Override
public Function<String, ? extends GCPSpannerSpace> getSpaceInitializer(NBConfiguration cfg) {
return (s) -> new GCPSpannerSpace(s, cfg);
}
@Override
public NBConfigModel getConfigModel() {
return super.getConfigModel().add(GCPSpannerSpace.getConfigModel());
}
}

View File

@ -0,0 +1,32 @@
/*
* Copyright (c) 2020-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;
import io.nosqlbench.adapter.diag.DriverAdapterLoader;
import io.nosqlbench.nb.annotations.Service;
import io.nosqlbench.nb.api.components.core.NBComponent;
import io.nosqlbench.nb.api.labels.NBLabels;
import static io.nosqlbench.adapter.gcpspanner.GCPSpannerAdapterUtils.SPANNER;
@Service(value = DriverAdapterLoader.class, selector = SPANNER)
public class GCPSpannerDriverAdapterLoader implements DriverAdapterLoader {
@Override
public GCPSpannerDriverAdapter load(NBComponent parent, NBLabels childLabels) {
return new GCPSpannerDriverAdapter(parent, childLabels);
}
}

View File

@ -0,0 +1,66 @@
/*
* Copyright (c) 2020-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;
import io.nosqlbench.adapter.gcpspanner.opdispensers.*;
import io.nosqlbench.adapter.gcpspanner.ops.GCPSpannerBaseOp;
import io.nosqlbench.adapter.gcpspanner.types.GCPSpannerOpType;
import io.nosqlbench.adapters.api.activityimpl.OpDispenser;
import io.nosqlbench.adapters.api.activityimpl.OpMapper;
import io.nosqlbench.adapters.api.templating.ParsedOp;
import io.nosqlbench.engine.api.templating.TypeAndTarget;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
public class GCPSpannerOpMapper implements OpMapper<GCPSpannerBaseOp<?>> {
private static final Logger logger = LogManager.getLogger(GCPSpannerOpMapper.class);
private final GCPSpannerDriverAdapter adapter;
/**
* Create a new {@code GCPSpannerOpMapper} implementing the {@link OpMapper}.
* interface.
*
* @param adapter The associated {@link GCPSpannerDriverAdapter}
*/
public GCPSpannerOpMapper(GCPSpannerDriverAdapter adapter) {
this.adapter = adapter;
}
/**
* Given an instance of a {@link ParsedOp} returns the appropriate
* {@link GCPSpannerBaseOpDispenser} subclass.
*
* @param op The {@link ParsedOp} to be evaluated
* @return The correct {@link GCPSpannerBaseOpDispenser} subclass based on
* the op type
*/
@Override
public OpDispenser<? extends GCPSpannerBaseOp<?>> apply(ParsedOp op) {
TypeAndTarget<GCPSpannerOpType, String> typeAndTarget = op.getTypeAndTarget(GCPSpannerOpType.class,
String.class, "type", "target");
logger.info(() -> "Using '" + typeAndTarget.enumId + "' op type for op template '" + op.getName() + "'");
return switch (typeAndTarget.enumId) {
case update_database_ddl ->
new GCPSpannerUpdateDatabaseDdlOpDispenser(adapter, op, typeAndTarget.targetFunction);
case insert ->
new GCPSpannerInsertOpDispenser(adapter, op, typeAndTarget.targetFunction);
case execute_dml ->
new GCPSpannerExecuteDmlOpDispenser(adapter, op, typeAndTarget.targetFunction);
};
}
}

View File

@ -0,0 +1,117 @@
/*
* Copyright (c) 2020-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;
import com.google.cloud.spanner.*;
import io.nosqlbench.nb.api.config.standard.ConfigModel;
import io.nosqlbench.nb.api.config.standard.NBConfigModel;
import io.nosqlbench.nb.api.config.standard.NBConfiguration;
import io.nosqlbench.nb.api.config.standard.Param;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
/**
* @see <a href="https://cloud.google.com/spanner/docs/find-approximate-nearest-neighbors">ANN Docs</a>
* @see <a href="https://github.com/googleapis/java-spanner">Spanner Java Client</a>
* @see <a href="https://cloud.google.com/spanner/docs/getting-started/java">Getting started in Java</a>
* @see <a href="https://cloud.google.com/docs/authentication#getting_credentials_for_server-centric_flow">Authentication methods at Google</a>
* @see <a href="https://cloud.google.com/java/docs/reference/google-cloud-spanner/latest/overview">Library Reference Doc</a>
* @see <a href="https://cloud.google.com/spanner/docs/reference/standard-sql/dml-syntax">DML Syntax</a>
* @see <a href="https://cloud.google.com/spanner/docs/reference/rpc">spanner rpc api calls</a>
* @see <a href="https://cloud.google.com/spanner/docs/reference/standard-sql/data-definition-language#vector_index_statements">SQL functionality related to vector indices</a>
*/
public class GCPSpannerSpace implements AutoCloseable {
private final static Logger logger = LogManager.getLogger(GCPSpannerSpace.class);
private final String name;
private final NBConfiguration cfg;
protected Spanner spanner;
protected DatabaseAdminClient dbAdminClient;
protected DatabaseClient dbClient;
/**
* Create a new {@code GCPSpannerSpace} Object which stores all stateful
* contextual information needed to interact with the <b>Google Spanner</b>
* database instance.
*
* @param name The name of this space
* @param cfg The configuration ({@link NBConfiguration}) for this nb run
*/
public GCPSpannerSpace(String name, NBConfiguration cfg) {
this.name = name;
this.cfg = cfg;
}
public synchronized Spanner getSpanner() {
if (spanner == null) {
createSpanner();
}
return spanner;
}
public synchronized DatabaseAdminClient getDbAdminClient() {
return dbAdminClient;
}
public synchronized DatabaseClient getDbClient() {
return dbClient;
}
public DatabaseId getDatabaseId() {
return DatabaseId.of(cfg.get("project_id"), cfg.get("instance_id"), cfg.get("database_id"));
}
public String getInstanceId() {
return cfg.get("instance_id");
}
public String getDatabaseIdString() {
return cfg.get("database_id");
}
private void createSpanner() {
if (
cfg.getOptional("database_id").isEmpty() ||
cfg.getOptional("project_id").isEmpty() ||
cfg.getOptional("instance_id").isEmpty() ||
cfg.getOptional("service_account_file").isEmpty()) {
throw new RuntimeException("You must provide all 'service_account_file', 'project_id', 'instance_id' &" +
" 'database_id' to configure a Google Spanner client");
}
String projectId = cfg.get("project_id");
String instanceId = cfg.get("instance_id");
String databaseId = cfg.get("database_id");
spanner = SpannerOptions.newBuilder().setProjectId(projectId).build().getService();
dbAdminClient = spanner.getDatabaseAdminClient();
dbClient = spanner.getDatabaseClient(DatabaseId.of(projectId, instanceId, databaseId));
}
public static NBConfigModel getConfigModel() {
return ConfigModel.of(GCPSpannerSpace.class)
.add(Param.optional("service_account_file", String.class, "the file to load the api token/key from. See https://cloud.google.com/docs/authentication/provide-credentials-adc#service-account"))
.add(Param.optional("project_id", String.class,"Project ID containing the Spanner database. See https://cloud.google.com/resource-manager/docs/creating-managing-projects"))
.add(Param.optional("instance_id", String.class, "Spanner database's Instance ID containing. See https://cloud.google.com/spanner/docs/getting-started/java#create_an_instance"))
.add(Param.optional("database_id", String.class, "Spanner Database ID. See https://cloud.google.com/spanner/docs/getting-started/java#create_a_database"))
.asReadOnly();
}
@Override
public void close() {
if (spanner != null) {
spanner.close();
}
spanner = null;
}
}

View File

@ -0,0 +1,58 @@
/*
* Copyright (c) 2020-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 io.nosqlbench.adapters.api.activityimpl.BaseOpDispenser;
import io.nosqlbench.adapter.gcpspanner.ops.GCPSpannerBaseOp;
import io.nosqlbench.adapter.gcpspanner.GCPSpannerSpace;
import io.nosqlbench.adapters.api.activityimpl.uniform.DriverAdapter;
import io.nosqlbench.adapters.api.templating.ParsedOp;
import java.util.function.LongFunction;
/**
* Abstract base class for GCP Spanner operation dispensers.
* This class extends the BaseOpDispenser and provides common functionality
* for creating GCP Spanner operations.
*/
public abstract class GCPSpannerBaseOpDispenser extends BaseOpDispenser<GCPSpannerBaseOp<?>, GCPSpannerSpace> {
/**
* A function that provides the target string based on a long input.
*/
protected final LongFunction<String> targetFunction;
/**
* A function that provides the GCP Spanner space based on a long input.
*/
protected final LongFunction<GCPSpannerSpace> spaceFunction;
/**
* Constructs a new GCPSpannerBaseOpDispenser.
*
* @param adapter the driver adapter for GCP Spanner operations
* @param op the parsed operation
* @param targetFunction a function that provides the target string
*/
protected GCPSpannerBaseOpDispenser(DriverAdapter<? extends GCPSpannerBaseOp<?>, GCPSpannerSpace> adapter, ParsedOp op,
LongFunction<String> targetFunction) {
super(adapter, op);
this.targetFunction = targetFunction;
this.spaceFunction = adapter.getSpaceFunc(op);
}
}

View File

@ -0,0 +1,86 @@
/*
* Copyright (c) 2020-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;
/**
* Dispenser class for creating GCP Spanner Execute DML operations.
* This class extends the GCPSpannerBaseOpDispenser and provides functionality
* to create and configure GCPSpannerExecuteDmlOp instances.
*/
public class GCPSpannerExecuteDmlOpDispenser extends GCPSpannerBaseOpDispenser {
private static final Logger logger = LogManager.getLogger(GCPSpannerExecuteDmlOpDispenser.class);
private final LongFunction<GCPSpannerExecuteDmlOp> opFunction;
/**
* Constructs a new GCPSpannerExecuteDmlOpDispenser.
*
* @param adapter the driver adapter for GCP Spanner operations
* @param op the parsed operation
* @param targetFunction a function that provides the target string
*/
public GCPSpannerExecuteDmlOpDispenser(GCPSpannerDriverAdapter adapter, ParsedOp op, LongFunction<String> targetFunction) {
super(adapter, op, targetFunction);
this.opFunction = createOpFunction(op);
}
/**
* Creates a function that generates GCPSpannerExecuteDmlOp instances.
*
* @param op the parsed operation
* @return a function that generates GCPSpannerExecuteDmlOp instances
*/
private LongFunction<GCPSpannerExecuteDmlOp> createOpFunction(ParsedOp op) {
return (l) -> new GCPSpannerExecuteDmlOp(
spaceFunction.apply(l).getSpanner(),
l,
generateStatement(targetFunction.apply(l)),
spaceFunction.apply(l).getDbClient()
);
}
/**
* Generates a Spanner Statement from a DML string.
*
* @param dml the DML string
* @return the generated Statement
*/
private Statement generateStatement(String dml) {
return Statement.of(dml);
}
/**
* Retrieves the GCP Spanner operation for the given value.
*
* @param value the input value
* @return the GCP Spanner operation
*/
@Override
public GCPSpannerBaseOp<?> getOp(long value) {
return opFunction.apply(value);
}
}

View File

@ -0,0 +1,99 @@
/*
* Copyright (c) 2020-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.Mutation;
import com.google.cloud.spanner.Value;
import io.nosqlbench.adapter.gcpspanner.GCPSpannerDriverAdapter;
import io.nosqlbench.adapter.gcpspanner.ops.GCPSpannerBaseOp;
import io.nosqlbench.adapter.gcpspanner.ops.GCPSpannerInsertOp;
import io.nosqlbench.adapters.api.templating.ParsedOp;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import java.util.Collections;
import java.util.Map;
import java.util.function.LongFunction;
/**
* This class is responsible for dispensing GCP Spanner insert vector operations.
* It extends the GCPSpannerBaseOpDispenser and provides the necessary implementation
* to create and configure GCPSpannerInsertVectorOp instances.
*/
public class GCPSpannerInsertOpDispenser extends GCPSpannerBaseOpDispenser {
private static final Logger logger = LogManager.getLogger(GCPSpannerInsertOpDispenser.class);
private final LongFunction<Map> queryParamsFunction;
/**
* Constructs a new GCPSpannerInsertVectorOpDispenser.
*
* @param adapter the GCP Spanner driver adapter
* @param op the parsed operation
* @param targetFunction a function that provides the target table name based on a long value
*/
public GCPSpannerInsertOpDispenser(GCPSpannerDriverAdapter adapter, ParsedOp op, LongFunction<String> targetFunction) {
super(adapter, op, targetFunction);
this.queryParamsFunction = createParamsFunction(op);
}
/**
* Creates a function that provides query parameters based on a long value.
*
* @param op the parsed operation
* @return a function that provides query parameters
*/
private LongFunction<Map> createParamsFunction(ParsedOp op) {
return op.getAsOptionalFunction("query_params", Map.class)
.orElse(_ -> Collections.emptyMap());
}
/**
* Returns a GCPSpannerInsertVectorOp instance configured with the provided value.
*
* @param value the value used to configure the operation
* @return a configured GCPSpannerInsertVectorOp instance
*/
@Override
public GCPSpannerBaseOp<?> getOp(long value) {
Mutation.WriteBuilder builder = Mutation.newInsertBuilder(targetFunction.apply(value));
Map<String, Object> params = queryParamsFunction.apply(value);
for (Map.Entry<String, Object> entry : params.entrySet()) {
builder.set(entry.getKey()).to(convertToValue(entry));
}
return new GCPSpannerInsertOp(
spaceFunction.apply(value).getSpanner(),
value,
builder.build(),
spaceFunction.apply(value).getDbClient()
);
}
private Value convertToValue(Map.Entry<String, Object> entry) {
return switch(entry.getValue()) {
case String s -> Value.string(s);
case Integer i -> Value.int64(i);
case Long l -> Value.int64(l);
case Double d -> Value.float64(d);
case Float f -> Value.float32(f);
case long[] larr -> Value.int64Array(larr);
case float[] farr -> Value.float32Array(farr);
case double[] darr -> Value.float64Array(darr);
default -> throw new IllegalArgumentException("Unsupported value type: " + entry.getValue().getClass());
};
}
}

View File

@ -0,0 +1,74 @@
/*
* Copyright (c) 2020-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 io.nosqlbench.adapter.gcpspanner.GCPSpannerDriverAdapter;
import io.nosqlbench.adapter.gcpspanner.ops.GCPSpannerBaseOp;
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;
/**
* Dispenser class for creating instances of GCPSpannerUpdateDatabaseDdlOp.
*/
public class GCPSpannerUpdateDatabaseDdlOpDispenser extends GCPSpannerBaseOpDispenser {
private static final Logger logger = LogManager.getLogger(GCPSpannerUpdateDatabaseDdlOpDispenser.class);
private final LongFunction<GCPSpannerUpdateDatabaseDdlOp> opFunction;
/**
* Constructor for GCPSpannerUpdateDatabaseDdlOpDispenser.
*
* @param adapter the GCPSpannerDriverAdapter instance
* @param op the ParsedOp instance
* @param targetFunction a LongFunction that provides the target string
*/
public GCPSpannerUpdateDatabaseDdlOpDispenser(GCPSpannerDriverAdapter adapter, ParsedOp op, LongFunction<String> targetFunction) {
super(adapter, op, targetFunction);
this.opFunction = createOpFunction(op);
}
/**
* Creates a LongFunction that generates GCPSpannerUpdateDatabaseDdlOp instances.
*
* @param op the ParsedOp instance
* @return a LongFunction that generates GCPSpannerUpdateDatabaseDdlOp instances
*/
private LongFunction<GCPSpannerUpdateDatabaseDdlOp> createOpFunction(ParsedOp op) {
return (l) -> new GCPSpannerUpdateDatabaseDdlOp(
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())
);
}
/**
* Retrieves an operation instance based on the provided value.
*
* @param value the long value used to generate the operation
* @return a GCPSpannerBaseOp instance
*/
@Override
public GCPSpannerBaseOp<?> getOp(long value) {
return opFunction.apply(value);
}
}

View File

@ -0,0 +1,88 @@
/*
* Copyright (c) 2020-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.Spanner;
import io.nosqlbench.adapters.api.activityimpl.uniform.flowtypes.CycleOp;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import java.util.function.LongFunction;
/**
* Abstract base class for GCP Spanner operations.
* This class implements the CycleOp interface and provides a template for executing operations with a Spanner client.
*
* @param <T> the type of the request parameter
*/
public abstract class GCPSpannerBaseOp<T> implements CycleOp<Object> {
protected final static Logger logger = LogManager.getLogger(GCPSpannerBaseOp.class);
protected final Spanner spannerClient;
protected final T request;
protected final LongFunction<Object> apiCall;
/**
* Constructs a new GCPSpannerBaseOp with the specified Spanner client and request parameter.
*
* @param spannerClient the Spanner client to use for operations
* @param requestParam the request parameter for the operation
*/
public GCPSpannerBaseOp(Spanner spannerClient, T requestParam) {
this.spannerClient = spannerClient;
this.request = requestParam;
this.apiCall = this::applyOp;
}
/**
* Applies the operation for the given cycle value.
* This method logs the operation and handles any exceptions by throwing a RuntimeException.
*
* @param value the cycle value
* @return the result of the operation
*/
@Override
public final Object apply(long value) {
logger.trace(() -> "applying op: " + this);
try {
Object result = applyOp(value);
return result;
} catch (Exception rte) {
throw new RuntimeException(rte);
}
}
/**
* Abstract method to be implemented by subclasses to define the specific operation logic.
*
* @param value the cycle value
* @return the result of the operation
*/
public abstract Object applyOp(long value);
/**
* Returns a string representation of the GCPSpannerBaseOp.
*
* @return a string representation of the GCPSpannerBaseOp
*/
@Override
public String toString() {
return "GCPSpannerBaseOp(" + this.request.getClass().getSimpleName() + ")";
}
}

View File

@ -0,0 +1,57 @@
/*
* Copyright (c) 2020-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.*;
/**
* This class represents an operation to execute a DML statement on Google Cloud Spanner.
* It extends the GCPSpannerBaseOp class and overrides the applyOp method to execute the DML statement.
*/
public class GCPSpannerExecuteDmlOp extends GCPSpannerBaseOp<Long> {
private final Statement statement;
private final DatabaseClient dbClient;
/**
* Constructs a new GCPSpannerExecuteDmlOp.
*
* @param spanner the Spanner instance
* @param requestParam the request parameter
* @param statement the DML statement to execute
* @param dbClient the DatabaseClient to use for executing the statement
*/
public GCPSpannerExecuteDmlOp(Spanner spanner, Long requestParam, Statement statement,
DatabaseClient dbClient) {
super(spanner, requestParam);
this.statement = statement;
this.dbClient = dbClient;
}
/**
* Executes the DML statement using the provided value.
*
* @param value the value to use for the operation
* @return the result of the DML execution
*/
@Override
public Object applyOp(long value) {
try (ReadContext context = dbClient.singleUse()) {
return context.executeQuery(statement);
}
}
}

View File

@ -0,0 +1,58 @@
/*
* Copyright (c) 2020-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.Spanner;
import com.google.cloud.spanner.DatabaseClient;
import com.google.cloud.spanner.Mutation;
import java.util.Collections;
/**
* This class represents an operation to insert a vector into a Google Cloud Spanner database.
* It extends the GCPSpannerBaseOp class and provides the implementation for the applyOp method.
*/
public class GCPSpannerInsertOp extends GCPSpannerBaseOp<Long> {
private final Mutation mutation;
private final DatabaseClient dbClient;
/**
* Constructs a new GCPSpannerInsertVectorOp.
*
* @param searchIndexClient the Spanner client used to interact with the database
* @param requestParam the request parameter
* @param mutation the Mutation object representing the data to be inserted
* @param dbClient the DatabaseClient used to execute the mutation
*/
public GCPSpannerInsertOp(Spanner searchIndexClient, Long requestParam, Mutation mutation, DatabaseClient dbClient) {
super(searchIndexClient, requestParam);
this.mutation = mutation;
this.dbClient = dbClient;
}
/**
* Applies the insert operation using the provided mutation.
*
* @param value the value to be used in the operation
* @return the result of the write operation
*/
@Override
public Object applyOp(long value) {
return dbClient.write(Collections.singletonList(mutation));
}
}

View File

@ -0,0 +1,70 @@
/*
* Copyright (c) 2020-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.api.gax.longrunning.OperationFuture;
import com.google.cloud.spanner.*;
import com.google.common.collect.ImmutableList;
import com.google.spanner.admin.database.v1.UpdateDatabaseDdlMetadata;
/**
* This class represents an operation to update the database DDL (Data Definition Language) in Google Cloud Spanner.
* It extends the GCPSpannerBaseOp class and provides the implementation for applying the DDL update operation.
*/
public class GCPSpannerUpdateDatabaseDdlOp extends GCPSpannerBaseOp<Long> {
private final String createTableStatement;
private final DatabaseAdminClient dbAdminClient;
private final Database db;
/**
* Constructs a new GCPSpannerUpdateDatabaseDdlOp.
*
* @param searchIndexClient the Spanner client
* @param requestParam the request parameter
* @param createTableStatement the SQL statement to create the table
* @param dbAdminClient the DatabaseAdminClient to execute the DDL update
* @param db the Database object representing the target database
*/
public GCPSpannerUpdateDatabaseDdlOp(Spanner searchIndexClient, Long requestParam, String createTableStatement,
DatabaseAdminClient dbAdminClient, Database db) {
super(searchIndexClient, requestParam);
this.createTableStatement = createTableStatement;
this.dbAdminClient = dbAdminClient;
this.db = db;
}
/**
* Applies the DDL update operation.
*
* @param value the value to be used in the operation
* @return the result of the operation
* @throws RuntimeException if an error occurs during the operation
*/
@Override
public Object applyOp(long value) {
OperationFuture<Void, UpdateDatabaseDdlMetadata> operation = dbAdminClient.updateDatabaseDdl(
db,
ImmutableList.of(createTableStatement),
null);
try {
return operation.get();
} catch (Exception e) {
throw new RuntimeException(e);
}
}
}

View File

@ -0,0 +1,36 @@
/*
* Copyright (c) 2020-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.types;
/**
* 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 {
update_database_ddl,
insert,
execute_dml,
}

View File

@ -0,0 +1,13 @@
scenarios:
default:
execute_ddl: run driver=spanner tags==blocks:execute_ddl service_account_file=TEMPLATE(service_account_file)
project_id=TEMPLATE(project_id) instance_id=TEMPLATE(instance_id) database_id=TEMPLATE(database_id) cycles=1
# https://cloud.google.com/spanner/docs/reference/standard-sql/data-definition-language#vector_index_option_list
blocks:
execute_ddl:
ops:
op1:
update_database_ddl: |
CREATE VECTOR INDEX VectorsIndex ON vectors(value)
OPTIONS (distance_type = 'COSINE', tree_depth = 3, num_branches=1000, num_leaves = 1000000);

View File

@ -0,0 +1,18 @@
scenarios:
default:
execute_dml: run driver=spanner tags==blocks:execute_dml service_account_file=TEMPLATE(service_account_file)
project_id=TEMPLATE(project_id) instance_id=TEMPLATE(instance_id) database_id=TEMPLATE(database_id) cycles=TEMPLATE(cycles)
bindings:
rw_key: ToString();
test_vector_hdf5: HdfFileToFloatList("testdata/TEMPLATE(dataset).hdf5", "/test"); ToCqlVector()
validation_set_hdf5: HdfFileToIntArray("testdata/TEMPLATE(dataset).hdf5", "/neighbors")
blocks:
execute_dml:
ops:
op1:
execute_dml: |
SELECT * FROM vectors@{FORCE_INDEX=VectorsIndex} ORDER BY APPROX_COSINE_DISTANCE(ARRAY<FLOAT32>{test_vector_hdf5},
value, options => JSON '{"num_leaves_to_search": 10}') LIMIT 100

View File

@ -0,0 +1,11 @@
scenarios:
default:
execute_ddl: run driver=spanner tags==blocks:execute_ddl service_account_file=TEMPLATE(service_account_file)
project_id=TEMPLATE(project_id) instance_id=TEMPLATE(instance_id) database_id=TEMPLATE(database_id) cycles=1
blocks:
execute_ddl:
ops:
op1:
update_database_ddl: |
CREATE TABLE vectors (keycol STRING(100),value ARRAY<FLOAT32>(vector_length=>25) NOT NULL) PRIMARY KEY(keycol)

View File

@ -0,0 +1,18 @@
scenarios:
default:
insert_vector: >-
run driver=spanner tags==blocks:insert_vector service_account_file=TEMPLATE(service_account_file)
project_id=TEMPLATE(project_id) instance_id=TEMPLATE(instance_id) database_id=TEMPLATE(database_id) cycles=TEMPLATE(cycles)
bindings:
rw_key: ToString();
train_floatlist: HdfFileToFloatArray("glove-25-angular.hdf5", "/train");
blocks:
insert_vector:
ops:
op1:
insert_vector: "vectors"
query_params:
keycol: "{rw_key}"
value: "{train_floatlist}"

View File

@ -0,0 +1,30 @@
# Google Spanner driver adapter
The Google Cloud Spanner driver adapter is a NoSQLBench adapter for the `gcp_spanner` driver, a Java driver
for connecting to and performing operations on an instance of a Google Cloud Spanner database.
## Activity Parameters
The following parameters must be supplied to the adapter at runtime in order to successfully connect to an
instance of the [Google Cloud Spanner database](https://cloud.google.com/java/docs/reference/google-cloud-spanner/latest/overview):
* `service_account_file` - In order to connect to a Spanner database you must have a [IAM service account](https://cloud.google.com/docs/authentication/provide-credentials-adc#service-account)
defined with the appropriate permissions associated with the adapter. Once the service account is created you can download
a file from the gcp console in JSON format that contains the credentials for the service account. This file must be provided
to the adapter at runtime.
* `project_id` - Project ID containing the Spanner database. See [Creating a project](https://cloud.google.com/resource-manager/docs/creating-managing-projects).
* `instance_id` - Spanner database's Instance ID. See [Creating an instance](https://cloud.google.com/spanner/docs/getting-started/java#create_an_instance).
* `database_id` - Spanner database's Database ID. See [Creating a database](https://cloud.google.com/spanner/docs/getting-started/java#create_a_database).
* In addition to this the environment variable `GOOGLE_APPLICATION_CREDENTIALS` must be set to the path of the service account file.
## Op Templates
The Google Cloud Spanner adapter supports the following operations:
* `update_database_ddl` - Data Definition Language operations such as creating and dropping tables, indexes, etc.
* `execute_dml` - Data Manipulation Language operations. Read only operations are supported at this time, including queries
and vector queries.
* `insert` - Insert a single record, vector or non-vector, of data into the database.
## Examples
---

View File

@ -61,13 +61,6 @@
<version>5.0.1</version>
</dependency>
<!-- &lt;!&ndash; https://mvnrepository.com/artifact/org.apache.commons/commons-collections4 &ndash;&gt;-->
<!-- <dependency>-->
<!-- <groupId>org.apache.commons</groupId>-->
<!-- <artifactId>commons-collections4</artifactId>-->
<!-- <version>4.4</version>-->
<!-- </dependency>-->
</dependencies>
<build>

View File

@ -295,6 +295,20 @@
</dependencies>
</profile>
<profile>
<id>adapter-gcp-spanner-include</id>
<activation>
<activeByDefault>true</activeByDefault>
</activation>
<dependencies>
<dependency>
<groupId>io.nosqlbench</groupId>
<artifactId>adapter-gcp-spanner</artifactId>
<version>${revision}</version>
</dependency>
</dependencies>
</profile>
</profiles>
</project>

View File

@ -214,5 +214,15 @@
</modules>
</profile>
<profile>
<id>adapter-gcp-spanner-module</id>
<activation>
<activeByDefault>true</activeByDefault>
</activation>
<modules>
<module>adapter-gcp-spanner</module>
</modules>
</profile>
</profiles>
</project>