src/main/java/com/googlesource/gerrit/plugins/spannerrefdb/DatabaseSchemaCreator.java

// Copyright (C) 2023 The Android Open Source Project
//
// 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.googlesource.gerrit.plugins.spannerrefdb;

import com.gerritforge.gerrit.globalrefdb.GlobalRefDbSystemError;
import com.google.api.gax.longrunning.OperationFuture;
import com.google.cloud.spanner.DatabaseAdminClient;
import com.google.cloud.spanner.DatabaseId;
import com.google.cloud.spanner.ErrorCode;
import com.google.cloud.spanner.SpannerException;
import com.google.common.flogger.FluentLogger;
import com.google.gerrit.extensions.events.LifecycleListener;
import com.google.inject.Inject;
import com.google.inject.Singleton;
import com.google.spanner.admin.database.v1.UpdateDatabaseDdlMetadata;
import java.util.Arrays;

@Singleton
class DatabaseSchemaCreator implements LifecycleListener {
  private static final FluentLogger logger = FluentLogger.forEnclosingClass();

  private DatabaseAdminClient dbAdminClient;
  private DatabaseId dbId;

  @Inject
  DatabaseSchemaCreator(DatabaseAdminClient dbAdminClient, DatabaseId dbId) {
    this.dbAdminClient = dbAdminClient;
    this.dbId = dbId;
  }

  @Override
  public void start() {
    createTable(
        "CREATE TABLE refs ("
            + "project  STRING(MAX) NOT NULL,"
            + "ref  STRING(MAX) NOT NULL,"
            + "value  STRING(MAX) NOT NULL"
            + ") PRIMARY KEY (project, ref)");
    createTable(
        "CREATE TABLE locks ("
            + "project  STRING(MAX) NOT NULL,"
            + "ref  STRING(MAX) NOT NULL,"
            + "heartbeat TIMESTAMP OPTIONS ("
            + "    allow_commit_timestamp = true"
            + "),"
            + "token TIMESTAMP OPTIONS ("
            + "    allow_commit_timestamp = true"
            + "),"
            + "owner STRING(MAX) NOT NULL"
            + ") PRIMARY KEY (project, ref)");
  }

  @Override
  public void stop() {}

  private void createTable(String statement) {
    try {
      OperationFuture<Void, UpdateDatabaseDdlMetadata> op =
          dbAdminClient.updateDatabaseDdl(
              dbId.getInstanceId().getInstance(),
              dbId.getDatabase(),
              Arrays.asList(statement),
              null);
      op.get();
      logger.atInfo().log("Created table in database [%s] with statement %s", dbId, statement);
    } catch (Exception e) {
      Throwable cause = e.getCause();
      if (cause instanceof SpannerException) {
        ErrorCode code = ((SpannerException) cause).getErrorCode();
        if (code != ErrorCode.FAILED_PRECONDITION) {
          throw new GlobalRefDbSystemError("Failed to create table", e);
        }
      }
    }
  }
}