io.stargate.sgv2.api.common.schema.SchemaManager Maven / Gradle / Ivy
/*
* Copyright The Stargate Authors
*
* 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.stargate.sgv2.api.common.schema;
import com.google.protobuf.BytesValue;
import com.google.protobuf.Int32Value;
import io.grpc.Status;
import io.grpc.StatusRuntimeException;
import io.opentelemetry.extension.annotations.WithSpan;
import io.quarkus.cache.Cache;
import io.quarkus.cache.CacheInvalidate;
import io.quarkus.cache.CacheKey;
import io.quarkus.cache.CacheName;
import io.quarkus.cache.CacheResult;
import io.quarkus.cache.CaffeineCache;
import io.quarkus.cache.CompositeCacheKey;
import io.smallrye.mutiny.Multi;
import io.smallrye.mutiny.Uni;
import io.smallrye.mutiny.tuples.Tuple2;
import io.stargate.bridge.proto.QueryOuterClass;
import io.stargate.bridge.proto.Schema;
import io.stargate.bridge.proto.StargateBridge;
import io.stargate.sgv2.api.common.StargateRequestInfo;
import io.stargate.sgv2.api.common.grpc.UnauthorizedKeyspaceException;
import io.stargate.sgv2.api.common.grpc.UnauthorizedTableException;
import io.stargate.sgv2.api.common.grpc.proto.SchemaReads;
import jakarta.enterprise.context.ApplicationScoped;
import jakarta.inject.Inject;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Objects;
import java.util.Optional;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Function;
import java.util.function.Supplier;
import java.util.stream.Collectors;
@ApplicationScoped
public class SchemaManager {
@Inject
@CacheName("keyspace-cache")
Cache keyspaceCache;
@Inject StargateRequestInfo requestInfo;
/**
* Get the keyspace from the bridge. Note that this method is not doing any authorization. The
* check that the keyspace has correct hash on the bridge will be done.
*
* @param keyspace Keyspace name
* @return Uni containing Schema.CqlKeyspaceDescribe or null item in case keyspace
* does not exist.
*/
@WithSpan
public Uni getKeyspace(String keyspace) {
return getKeyspace(keyspace, true);
}
/**
* Get the keyspace from the bridge. Note that this method is not doing any authorization.
*
* @param keyspace Keyspace name
* @param validateHash If hash validation should be done for the keyspace. If false
* and the keyspace is already cached, then no calls to the bridge are executed.
* @return Uni containing Schema.CqlKeyspaceDescribe or null item in case keyspace
* does not exist.
*/
@WithSpan
public Uni getKeyspace(String keyspace, boolean validateHash) {
StargateBridge bridge = requestInfo.getStargateBridge();
return getKeyspaceInternal(bridge, keyspace, validateHash);
}
/**
* Get all keyspace from the bridge. Note that this method is not doing any authorization. The
* check that each keyspace has correct hash on the bridge will be done.
*
* @return Multi containing Schema.CqlKeyspaceDescribe
*/
@WithSpan
public Multi getKeyspaces() {
StargateBridge bridge = requestInfo.getStargateBridge();
// get all names
return getKeyspaceNames(bridge)
// then fetch each keyspace
.onItem()
.transformToUniAndMerge(keyspace -> getKeyspaceInternal(bridge, keyspace, true));
}
/**
* Get the table from the bridge. Note that this method is not doing any authorization. The check
* that the keyspace has correct hash on the bridge will be done.
*
* @param keyspace Keyspace name
* @param table Table name
* @param missingKeyspace Function of the keyspace in case it's not existing. Usually there to
* provide a failure.
* @return Uni containing Schema.CqlTable or null item in case the table does not
* exist.
*/
@WithSpan
public Uni getTable(
String keyspace,
String table,
Function> missingKeyspace) {
StargateBridge bridge = requestInfo.getStargateBridge();
return getTableInternal(bridge, keyspace, table, missingKeyspace);
}
/**
* Get all tables of a keyspace from the bridge. The check that the keyspace has correct hash on
* the bridge will be done.
*
* @param keyspace Keyspace name
* @param missingKeyspace Function of the keyspace in case it's not existing. Usually there to
* provide a failure.
* @return Multi of Schema.CqlTable
*/
@WithSpan
public Multi getTables(
String keyspace,
Function> missingKeyspace) {
StargateBridge bridge = requestInfo.getStargateBridge();
// get keyspace
return getKeyspaceInternal(bridge, keyspace, true)
// if not there, switch to function
.onItem()
.ifNull()
.switchTo(() -> missingKeyspace.apply(keyspace))
// otherwise get all tables
.onItem()
.ifNotNull()
.transformToMulti(k -> Multi.createFrom().iterable(k.getTablesList()));
}
/**
* Get the keyspace from the bridge. Prior to getting the keyspace it will execute the schema
* authorization request. The check that the keyspace has correct hash on the bridge will be done.
*
* Emits a failure in case:
*
*
* - Not authorized, with {@link UnauthorizedKeyspaceException}
*
*
* @param keyspace Keyspace name
* @return Uni containing Schema.CqlKeyspaceDescribe or null item in case keyspace
* does not exist.
*/
@WithSpan
public Uni getKeyspaceAuthorized(String keyspace) {
return getKeyspaceAuthorized(keyspace, true);
}
/**
* Get the keyspace from the bridge. Prior to getting the keyspace it will execute the schema
* authorization request.
*
* Emits a failure in case:
*
*
* - Not authorized, with {@link UnauthorizedKeyspaceException}
*
*
* @param keyspace Keyspace name
* @param validateHash If hash validation should be done for the keyspace. If false
* and the keyspace is already cached, then no calls to the bridge are executed.
* @return Uni containing Schema.CqlKeyspaceDescribe or null item in case keyspace
* does not exist.
*/
@WithSpan
public Uni getKeyspaceAuthorized(
String keyspace, boolean validateHash) {
StargateBridge bridge = requestInfo.getStargateBridge();
// first authorize read, then fetch
return authorizeKeyspaceInternal(bridge, keyspace)
// on result
.onItem()
.transformToUni(
authorized -> {
// if authorized, go fetch keyspace
// otherwise throw correct exception
if (authorized) {
return getKeyspaceInternal(bridge, keyspace, validateHash);
} else {
RuntimeException unauthorized = new UnauthorizedKeyspaceException(keyspace);
return Uni.createFrom().failure(unauthorized);
}
});
}
/**
* Get all keyspace from the bridge. Prior to getting each keyspace it will execute the schema
* authorization request (single request for all available keyspace). The check that each keyspace
* has correct hash on the bridge will be done.
*
* @return Multi containing Schema.CqlKeyspaceDescribe
*/
@WithSpan
public Multi getKeyspacesAuthorized() {
StargateBridge bridge = requestInfo.getStargateBridge();
// get all keyspace names
return getKeyspaceNames(bridge)
// collect list
.collect()
.asList()
// then go for the schema read
.onItem()
.transformToMulti(
keyspaceNames -> {
// if we have no keyspace return immediately
if (null == keyspaceNames || keyspaceNames.isEmpty()) {
return Multi.createFrom().empty();
}
// create schema reads for all keyspaces
List reads =
keyspaceNames.stream()
.map(n -> SchemaReads.keyspace(n))
.collect(Collectors.toList());
Schema.AuthorizeSchemaReadsRequest request =
Schema.AuthorizeSchemaReadsRequest.newBuilder().addAllSchemaReads(reads).build();
// execute request
return bridge
.authorizeSchemaReads(request)
// on response filter out
.onItem()
.ifNotNull()
.transformToMulti(
response -> {
List authorizedKeyspaces = new ArrayList<>(keyspaceNames.size());
List authorizedList = response.getAuthorizedList();
for (int i = 0; i < authorizedList.size(); i++) {
if (authorizedList.get(i)) {
authorizedKeyspaces.add(keyspaceNames.get(i));
}
}
// and return all authorized tables
return Multi.createFrom().iterable(authorizedKeyspaces);
});
})
// then fetch each authorized keyspace
.onItem()
.transformToUniAndMerge(keyspace -> getKeyspaceInternal(bridge, keyspace, true));
}
/**
* Get the table from the bridge. Prior to getting the keyspace it will execute the schema
* authorization request. The check that the keyspace has correct hash on the bridge will be done.
*
* Emits a failure in case:
*
*
* - Not authorized, with {@link UnauthorizedTableException}
*
*
* @param keyspace Keyspace name
* @param table Table name
* @param missingKeyspace Function of the keyspace in case it's not existing. Usually there to
* provide a failure.
* @return Uni containing Schema.CqlTable or null item in case the table does not
* exist.
*/
@WithSpan
public Uni getTableAuthorized(
String keyspace,
String table,
Function> missingKeyspace) {
StargateBridge bridge = requestInfo.getStargateBridge();
// first authorize read, then fetch
return authorizeTableInternal(bridge, keyspace, table)
// on result
.onItem()
.transformToUni(
authorized -> {
// if authorized, go fetch keyspace
// otherwise throw correct exception
if (authorized) {
return getTableInternal(bridge, keyspace, table, missingKeyspace);
} else {
RuntimeException unauthorized = new UnauthorizedTableException(keyspace, table);
return Uni.createFrom().failure(unauthorized);
}
});
}
/**
* Get all authorized tables from the bridge. The check that the keyspace has correct hash on the
* bridge will be done.
*
* Emits a failure in case:
*
*
* - Not authorized, with {@link UnauthorizedTableException}
*
*
* @param keyspace Keyspace name
* @param missingKeyspace Function of the keyspace in case it's not existing. Usually there to
* provide a failure.
* @return Multi of Schema.CqlTable
*/
@WithSpan
public Multi getTablesAuthorized(
String keyspace,
Function> missingKeyspace) {
StargateBridge bridge = requestInfo.getStargateBridge();
// get keyspace
return getKeyspaceInternal(bridge, keyspace, true)
// if keyspace not found switch to function
.onItem()
.ifNull()
.switchTo(() -> missingKeyspace.apply(keyspace))
// if it exists go forward
.onItem()
.ifNotNull()
.transformToMulti(
keyspaceDescribe -> {
// create schema reads for all tables
List tables = keyspaceDescribe.getTablesList();
// if empty break immediately
if (tables.isEmpty()) {
return Multi.createFrom().empty();
}
List reads =
tables.stream()
.map(t -> SchemaReads.table(keyspace, t.getName()))
.collect(Collectors.toList());
Schema.AuthorizeSchemaReadsRequest request =
Schema.AuthorizeSchemaReadsRequest.newBuilder().addAllSchemaReads(reads).build();
// execute request
return bridge
.authorizeSchemaReads(request)
// on response filter out
.onItem()
.ifNotNull()
.transformToMulti(
response -> {
List authorizedTables = new ArrayList<>(tables.size());
List authorizedList = response.getAuthorizedList();
for (int i = 0; i < authorizedList.size(); i++) {
if (authorizedList.get(i)) {
authorizedTables.add(tables.get(i));
}
}
// and return all authorized tables
return Multi.createFrom().iterable(authorizedTables);
});
});
}
/**
* Executes the optimistic query, by fetching the keyspace without authorization from the @{@link
* SchemaManager}.
*
* @param keyspace Keyspace name.
* @param table Table name.
* @param missingKeyspace Function in case the keyspace in case it's not existing. Usually there
* to * provide a failure.
* @param queryFunction Function that creates the query from the {@link Schema.CqlTable}. Note
* that the table can be null.
* @return Response when optimistic query is executed correctly.
*/
public Uni queryWithSchema(
String keyspace,
String table,
Function> missingKeyspace,
Function> queryFunction) {
// get table from the schema manager without the hash validation
Uni keyspaceUni = getKeyspace(keyspace, false);
Optional tenantId = requestInfo.getTenantId();
return queryWithSchema(
keyspaceUni,
keyspace,
table,
tenantId,
() -> missingKeyspace.apply(keyspace),
queryFunction,
true);
}
/**
* Executes the optimistic query, by fetching the keyspace with authorization from the @{@link
* SchemaManager}.
*
* @param keyspace Keyspace name.
* @param table Table name.
* @param missingKeyspace Function in case the keyspace in case it's not existing. Usually there
* to provide a failure.
* @param queryFunction Function that creates the query from the {@link Schema.CqlTable}. Note
* that the table can be null.
* @return Response when optimistic query is executed correctly.
*/
public Uni queryWithSchemaAuthorized(
String keyspace,
String table,
Function> missingKeyspace,
Function> queryFunction) {
// get table from the schema manager authorized without the hash validation
Uni keyspaceUni = getKeyspaceAuthorized(keyspace, false);
Optional tenantId = requestInfo.getTenantId();
return queryWithSchema(
keyspaceUni,
keyspace,
table,
tenantId,
() -> missingKeyspace.apply(keyspace),
queryFunction,
true);
}
// IMPORTANT this method must stay private
// it's internal query with schema implementation
// if authorization is needed, it must be done before this method is applied
private Uni queryWithSchema(
Uni keyspaceUni,
String keyspace,
String table,
Optional tenantId,
Supplier> missingKeyspace,
Function> queryFunction,
boolean revalidateOnTableMiss) {
// get from cql keyspace
return keyspaceUni
// if keyspace is found, execute optimistic query with that keyspace
.onItem()
.ifNotNull()
.transformToUni(
cqlKeyspace -> {
// try to find the table
Schema.CqlTable cqlTable = findTable(cqlKeyspace, table);
// having the table or revalidate false, continue
if (null != cqlTable || !revalidateOnTableMiss) {
return queryWithSchemaOnKeyspaceTable(cqlKeyspace, cqlTable, queryFunction)
// once we have a result pipe to our handler
.onItem()
.transformToUni(
queryWithSchemaHandler(
keyspace, table, tenantId, missingKeyspace, queryFunction));
} else {
// IMPORTANT
// we need to revalidate the keyspace, cause table could be added
// and we call getKeyspace no matter if this sequence was used for authorized use or
// not
// because authorization must always come before
Uni validatedKeyspace = getKeyspace(keyspace);
return queryWithSchema(
validatedKeyspace,
keyspace,
table,
tenantId,
missingKeyspace,
queryFunction,
false);
}
})
// if keyspace not found at first place, use mapper
.onItem()
.ifNull()
.switchTo(missingKeyspace);
}
// handles the QueryWithSchemaResponse
private Function>
queryWithSchemaHandler(
String keyspace,
String table,
Optional tenantId,
Supplier> missingKeyspace,
Function> queryFunction) {
return response -> {
// * if no keyspace, calls the missing key space function
// * if there is new keyspace, execute again with that keyspace, use same handler
// * if everything is fine, map to result
if (response.hasNoKeyspace()) {
return missingKeyspace.get();
} else if (response.hasNewKeyspace()) {
// first invalidate existing keyspace
return invalidateKeyspace(keyspace, tenantId)
// then cache the update
.flatMap(v -> cacheKeyspace(keyspace, tenantId, response.getNewKeyspace()))
// then query again and handler with same handler
// note that this is endlessly trying to execute the query in case there is always a
// changes
// keyspace
.flatMap(
updatedKeyspace -> {
Schema.CqlTable cqlTable = findTable(updatedKeyspace, table);
return queryWithSchemaOnKeyspaceTable(updatedKeyspace, cqlTable, queryFunction);
})
.onItem()
.transformToUni(
queryWithSchemaHandler(keyspace, table, tenantId, missingKeyspace, queryFunction));
} else {
return Uni.createFrom().item(response.getResponse());
}
};
}
// executes the optimistic query against the keyspace
private Uni queryWithSchemaOnKeyspaceTable(
Schema.CqlKeyspaceDescribe cqlKeyspace,
Schema.CqlTable cqlTable,
Function> queryFunction) {
// start sequence from table
return Uni.createFrom()
.item(cqlTable)
// query function should handle not found table,
// so just hit
.flatMap(queryFunction::apply)
.flatMap(
query -> {
// construct request
Schema.QueryWithSchema request =
Schema.QueryWithSchema.newBuilder()
.setQuery(query)
.setKeyspaceName(cqlKeyspace.getCqlKeyspace().getName())
.setKeyspaceHash(cqlKeyspace.getHash().getValue())
.build();
// fire call
return requestInfo.getStargateBridge().executeQueryWithSchema(request);
});
}
// authorizes a keyspace by provided name
private Uni authorizeKeyspaceInternal(StargateBridge bridge, String keyspaceName) {
Schema.SchemaRead schemaRead = SchemaReads.keyspace(keyspaceName);
return authorizeInternal(bridge, schemaRead);
}
// authorizes a table by provided name and keyspace
private Uni authorizeTableInternal(
StargateBridge bridge, String keyspaceName, String tableName) {
Schema.SchemaRead schemaRead = SchemaReads.table(keyspaceName, tableName);
return authorizeInternal(bridge, schemaRead);
}
// authorizes a single schema read
private Uni authorizeInternal(StargateBridge bridge, Schema.SchemaRead schemaRead) {
Schema.AuthorizeSchemaReadsRequest request =
Schema.AuthorizeSchemaReadsRequest.newBuilder().addSchemaReads(schemaRead).build();
// call bridge to authorize
return bridge
.authorizeSchemaReads(request)
// on result
.map(
response -> {
// we have only one schema read request
return response.getAuthorizedList().iterator().next();
});
}
// gets a keyspace by provided name
// if validate hash is false, cached keyspaces are not validated to have correct hash
private Uni getKeyspaceInternal(
StargateBridge bridge, String keyspaceName, boolean validateHash) {
Optional tenantId = requestInfo.getTenantId();
// check if cached
return Uni.createFrom()
.deferred(
() -> {
CompositeCacheKey cacheKey = new CompositeCacheKey(keyspaceName, tenantId);
CompletableFuture © 2015 - 2024 Weber Informatics LLC | Privacy Policy