com.tigrisdata.db.client.StandardTigrisAsyncClient Maven / Gradle / Ivy

Go to download
/*
 * Copyright 2022 Tigris Data, Inc.
 *
 * 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 com.tigrisdata.db.client;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.util.concurrent.ListenableFuture;
import com.tigrisdata.db.api.v1.grpc.Api;
import com.tigrisdata.db.api.v1.grpc.TigrisDBGrpc;
import static com.tigrisdata.db.client.Messages.CREATE_DB_FAILED;
import static com.tigrisdata.db.client.Messages.DROP_DB_FAILED;
import static com.tigrisdata.db.client.Messages.LIST_DBS_FAILED;
import static com.tigrisdata.db.client.TypeConverter.toCreateDatabaseRequest;
import static com.tigrisdata.db.client.TypeConverter.toDropDatabaseRequest;
import static com.tigrisdata.db.client.TypeConverter.toListDatabasesRequest;
import com.tigrisdata.db.client.auth.AuthorizationToken;
import com.tigrisdata.db.client.config.TigrisConfiguration;
import com.tigrisdata.db.client.error.TigrisException;
import com.tigrisdata.tools.schema.core.ModelToJsonSchema;
import com.tigrisdata.tools.schema.core.StandardModelToTigrisDBJsonSchema;
import io.grpc.ManagedChannel;
import io.grpc.ManagedChannelBuilder;
import io.grpc.Status;
import io.grpc.StatusRuntimeException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.List;
import java.util.Optional;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Executor;
import java.util.concurrent.Executors;
import java.util.function.BiConsumer;

/** Async client for TigrisDB */
public class StandardTigrisAsyncClient extends AbstractTigrisDBClient implements TigrisAsyncClient {

  private final TigrisDBGrpc.TigrisDBFutureStub stub;
  private final Executor executor;
  private static final Logger log = LoggerFactory.getLogger(StandardTigrisAsyncClient.class);

  private StandardTigrisAsyncClient(TigrisConfiguration clientConfiguration) {
    this(clientConfiguration, Executors.newCachedThreadPool());
  }

  StandardTigrisAsyncClient(TigrisConfiguration clientConfiguration, Executor executor) {
    // TODO: authorization token injection
    super(clientConfiguration, Optional.empty(), new StandardModelToTigrisDBJsonSchema());
    this.stub = TigrisDBGrpc.newFutureStub(channel);
    this.executor = executor;
  }

  @VisibleForTesting
  StandardTigrisAsyncClient(
      AuthorizationToken authorizationToken,
      TigrisConfiguration configuration,
      ManagedChannelBuilder managedChannelBuilder) {
    super(
        authorizationToken,
        configuration,
        managedChannelBuilder,
        new StandardModelToTigrisDBJsonSchema());
    this.stub = TigrisDBGrpc.newFutureStub(channel);
    this.executor = Executors.newFixedThreadPool(Runtime.getRuntime().availableProcessors());
  }

  /**
   * Creates a new instance of @{@link StandardTigrisAsyncClient} with the given inputs
   *
   * @param tigrisConfiguration configuration
   * @return a new instance of @{@link StandardTigrisAsyncClient}
   */
  public static StandardTigrisAsyncClient getInstance(TigrisConfiguration tigrisConfiguration) {
    return new StandardTigrisAsyncClient(tigrisConfiguration);
  }

  /**
   * Creates a new instance of @{@link StandardTigrisAsyncClient} with the given inputs
   *
   * @param tigrisConfiguration configuration
   * @param executor executor that executes the future translation
   * @return a new instance of @{@link StandardTigrisAsyncClient}
   */
  public static StandardTigrisAsyncClient getInstance(
      TigrisConfiguration tigrisConfiguration, Executor executor) {
    return new StandardTigrisAsyncClient(tigrisConfiguration, executor);
  }

  @Override
  public TigrisAsyncDatabase getDatabase(String databaseName) {
    return new StandardTigrisAsyncDatabase(
        databaseName, stub, channel, executor, objectMapper, modelToJsonSchema);
  }

  @Override
  public CompletableFuture> listDatabases(
      DatabaseOptions listDatabaseOptions) {

    ListenableFuture listenableFuture =
        stub.listDatabases(toListDatabasesRequest(listDatabaseOptions));
    return Utilities.transformFuture(
        listenableFuture,
        listDatabasesResponse -> {
          List tigrisAsyncDatabases = new ArrayList<>();
          for (Api.DatabaseInfo databaseInfo : listDatabasesResponse.getDatabasesList()) {
            tigrisAsyncDatabases.add(
                new StandardTigrisAsyncDatabase(
                    databaseInfo.getDb(),
                    stub,
                    channel,
                    executor,
                    objectMapper,
                    modelToJsonSchema));
          }
          return tigrisAsyncDatabases;
        },
        executor,
        LIST_DBS_FAILED);
  }

  @Override
  public CompletableFuture createDatabaseIfNotExists(String databaseName) {
    ListenableFuture createDatabaseResponse =
        stub.createDatabase(
            toCreateDatabaseRequest(databaseName, DatabaseOptions.DEFAULT_INSTANCE));
    return Utilities.transformFuture(
        createDatabaseResponse,
        response ->
            new StandardTigrisAsyncDatabase(
                databaseName, stub, channel, executor, objectMapper, modelToJsonSchema),
        executor,
        CREATE_DB_FAILED,
        Optional.of(
            new CreateDatabaseExceptionHandler(
                databaseName, stub, executor, channel, objectMapper, modelToJsonSchema)));
  }

  @Override
  public CompletableFuture dropDatabase(String databaseName) {
    ListenableFuture dropDatabaseResponse =
        stub.dropDatabase(toDropDatabaseRequest(databaseName, DatabaseOptions.DEFAULT_INSTANCE));
    return Utilities.transformFuture(
        dropDatabaseResponse,
        response -> new DropDatabaseResponse(response.getStatus(), response.getMessage()),
        executor,
        DROP_DB_FAILED);
  }

  @Override
  public void close() {
    channel.shutdown();
  }

  @VisibleForTesting
  ManagedChannel getChannel() {
    return channel;
  }

  /**
   * This is the exception handler for CreateDatabase operation, here it will swallow the exception
   * if the server says database already exists, it will pass the exception further otherwise.
   */
  static class CreateDatabaseExceptionHandler
      implements BiConsumer, Throwable> {
    private final String dbName;
    private final TigrisDBGrpc.TigrisDBFutureStub stub;
    private final Executor executor;
    private final ManagedChannel channel;
    private final ObjectMapper objectMapper;
    private final ModelToJsonSchema modelToJsonSchema;

    public CreateDatabaseExceptionHandler(
        String dbName,
        TigrisDBGrpc.TigrisDBFutureStub stub,
        Executor executor,
        ManagedChannel channel,
        ObjectMapper objectMapper,
        ModelToJsonSchema modelToJsonSchema) {
      this.dbName = dbName;
      this.stub = stub;
      this.executor = executor;
      this.channel = channel;
      this.objectMapper = objectMapper;
      this.modelToJsonSchema = modelToJsonSchema;
    }

    @Override
    public void accept(
        CompletableFuture completableFuture, Throwable throwable) {
      if (throwable instanceof StatusRuntimeException) {
        if (((StatusRuntimeException) throwable).getStatus().getCode()
            == Status.ALREADY_EXISTS.getCode()) {
          // swallow the already exists exception
          log.info("database already exists: {}", dbName);
          completableFuture.complete(
              new StandardTigrisAsyncDatabase(
                  dbName, stub, channel, executor, objectMapper, modelToJsonSchema));
          return;
        }
      }
      // pass on the error otherwise
      completableFuture.completeExceptionally(new TigrisException(CREATE_DB_FAILED, throwable));
    }
  }
}