/*
 * Copyright (C) 2020 Google 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.google.cloud.teleport.spanner;

import com.google.cloud.spanner.BatchClient;
import com.google.cloud.spanner.DatabaseAdminClient;
import com.google.cloud.spanner.DatabaseClient;
import com.google.cloud.spanner.DatabaseId;
import com.google.cloud.spanner.Spanner;
import com.google.cloud.spanner.SpannerException;
import com.google.cloud.spanner.SpannerOptions;
import com.google.cloud.spanner.TransactionContext;
import com.google.cloud.spanner.TransactionRunner;
import com.google.cloud.teleport.spanner.ddl.Ddl;
import com.google.cloud.teleport.spanner.ddl.RandomInsertMutationGenerator;
import java.util.Iterator;
import javax.annotation.Nullable;
import org.apache.beam.sdk.io.gcp.spanner.MutationGroup;
import org.apache.beam.sdk.io.gcp.spanner.SpannerConfig;
import org.apache.beam.sdk.options.ValueProvider;
import org.junit.rules.ExternalResource;

/** Facilitates setup and deletion of a Spanner database for integration tests. */
public class SpannerServerResource extends ExternalResource {
  // Modify the following parameters to match your Cloud Spanner instance.
  private final String projectId = "span-cloud-testing";
  private final String instanceId = "test-instance";
  private final String host = "https://spanner.googleapis.com";

  private Spanner client;
  private DatabaseAdminClient databaseAdminClient;

  @Override
  protected void before() {
    SpannerOptions spannerOptions =
        SpannerOptions.newBuilder().setProjectId(projectId).setHost(host).build();
    client = spannerOptions.getService();
    databaseAdminClient = client.getDatabaseAdminClient();
  }

  @Override
  protected void after() {
    client.close();
  }

  public void createDatabase(String dbName, Iterable<String> ddlStatements) throws Exception {
    // Waits for create database to complete.
    databaseAdminClient.createDatabase(instanceId, dbName, ddlStatements).get();
  }

  public void updateDatabase(String dbName, Iterable<String> ddlStatements) throws Exception  {
    databaseAdminClient.updateDatabaseDdl(instanceId, dbName, ddlStatements, null).get();
  }

  public void dropDatabase(String dbName) {
    try {
      databaseAdminClient.dropDatabase(instanceId, dbName);
    } catch (SpannerException e) {
      // Does not exist, ignore.
    }
  }

  public BatchClient getBatchClient(String dbName) {
    return client.getBatchClient(DatabaseId.of(projectId, instanceId, dbName));
  }

  public DatabaseClient getDbClient(String dbName) {
    return client.getDatabaseClient(DatabaseId.of(projectId, instanceId, dbName));
  }

  public SpannerConfig getSpannerConfig(String dbName) {
    return SpannerConfig.create()
        .withProjectId(projectId)
        .withInstanceId(instanceId)
        .withDatabaseId(dbName)
        .withHost(ValueProvider.StaticValueProvider.of(host));
  }

  public void populateRandomData(String db, Ddl ddl, int numBatches) throws Exception {

    final Iterator<MutationGroup> mutations = new RandomInsertMutationGenerator(ddl).stream()
        .iterator();

    for (int i = 0; i < numBatches; i++) {
      TransactionRunner transactionRunner = getDbClient(db).readWriteTransaction();
      transactionRunner.run(new TransactionRunner.TransactionCallable<Void>() {

        @Nullable
        @Override
        public Void run(TransactionContext transaction) {
          for (int i = 0; i < 10; i++) {
            MutationGroup m = mutations.next();
            transaction.buffer(m);
          }
          return null;
        }
      });
    }
  }

}

@@ -1678,16 +1678,102 @@
 stance.%0A
+  private static final String EMULATOR_HOST = System.getenv(%22SPANNER_EMULATOR_HOST%22);%0A
   privat
@@ -1821,16 +1821,16 @@
 sting%22;%0A
-
   privat
@@ -2094,18 +2094,65 @@
 ions
- =%0A       
+;%0A    if (EMULATOR_HOST == null) %7B%0A      spannerOptions =
  Spa
@@ -2219,24 +2219,213 @@
 t).build();%0A
+    %7D else %7B%0A      spannerOptions =%0A          SpannerOptions.newBuilder()%0A              .setProjectId(projectId)%0A              .setEmulatorHost(EMULATOR_HOST)%0A              .build();%0A    %7D%0A
     client =
@@ -3521,32 +3521,67 @@
 tring dbName) %7B%0A
+    if (EMULATOR_HOST == null) %7B%0A  
     return Spann
@@ -3602,24 +3602,26 @@
 e()%0A        
+ 
+ 
 .withProject
@@ -3626,32 +3626,34 @@
 ctId(projectId)%0A
+  
         .withIns
@@ -3668,24 +3668,26 @@
 instanceId)%0A
+  
         .wit
@@ -3710,24 +3710,26 @@
 me)%0A        
+  
 .withHost(Va
@@ -3767,24 +3767,269 @@
 .of(host));%0A
+    %7D else %7B%0A      return SpannerConfig.create()%0A          .withProjectId(projectId)%0A          .withInstanceId(instanceId)%0A          .withDatabaseId(dbName)%0A          .withEmulatorHost(ValueProvider.StaticValueProvider.of(EMULATOR_HOST));%0A    %7D%0A
   %7D%0A%0A  publi

Update to allow Cloud Spanner Emulator to use dataflow templates.