sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/TestBigQuery.java - beam - Git at Google

 /*
  * Licensed to the Apache Software Foundation (ASF) under one
  * or more contributor license agreements.  See the NOTICE file
  * distributed with this work for additional information
  * regarding copyright ownership.  The ASF licenses this file
  * to you 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 org.apache.beam.sdk.io.gcp.bigquery;

 import static org.joda.time.Seconds.secondsBetween;
 import static org.junit.Assert.assertThat;

 import com.google.api.client.http.HttpRequestInitializer;
 import com.google.api.services.bigquery.Bigquery;
 import com.google.api.services.bigquery.model.Table;
 import com.google.api.services.bigquery.model.TableReference;
 import com.google.api.services.bigquery.model.TableRow;
 import com.google.api.services.bigquery.model.TableSchema;
 import com.google.auth.Credentials;
 import com.google.auth.http.HttpCredentialsAdapter;
 import com.google.cloud.hadoop.util.ChainingHttpRequestInitializer;
 import java.io.IOException;
 import java.util.Collections;
 import java.util.List;
 import java.util.concurrent.ThreadLocalRandom;
 import java.util.stream.Collectors;
 import org.apache.beam.sdk.extensions.gcp.auth.NullCredentialInitializer;
 import org.apache.beam.sdk.extensions.gcp.util.RetryHttpRequestInitializer;
 import org.apache.beam.sdk.extensions.gcp.util.Transport;
 import org.apache.beam.sdk.schemas.Schema;
 import org.apache.beam.sdk.schemas.Schema.FieldType;
 import org.apache.beam.sdk.testing.TestPipeline;
 import org.apache.beam.sdk.testing.TestPipelineOptions;
 import org.apache.beam.sdk.values.Row;
 import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
 import org.hamcrest.Matcher;
 import org.joda.time.DateTime;
 import org.joda.time.Duration;
 import org.joda.time.Instant;
 import org.joda.time.format.DateTimeFormat;
 import org.joda.time.format.DateTimeFormatter;
 import org.junit.rules.TestRule;
 import org.junit.runner.Description;
 import org.junit.runners.model.Statement;

 /**
  * Test rule which creates a new table with specified schema, with randomized name and exposes few
  * APIs to work with it.
  *
  * <p>Deletes the table on test shutdown.
  */
 public class TestBigQuery implements TestRule {
   private static final DateTimeFormatter DATETIME_FORMAT =
       DateTimeFormat.forPattern("YYYY_MM_dd_HH_mm_ss_SSS");

   private TestBigQueryOptions pipelineOptions;
   private Schema schema;
   private Table table;
   private BigQueryServices.DatasetService datasetService;

   /**
    * Creates an instance of this rule.
    *
    * <p>Loads GCP configuration from {@link TestPipelineOptions}.
    */
   public static TestBigQuery create(Schema tableSchema) {
     return new TestBigQuery(
         TestPipeline.testingPipelineOptions().as(TestBigQueryOptions.class), tableSchema);
   }

   private TestBigQuery(TestBigQueryOptions pipelineOptions, Schema tableSchema) {
     this.pipelineOptions = pipelineOptions;
     this.schema = tableSchema;
   }

   @Override
   public Statement apply(Statement base, Description description) {
     return new Statement() {
       @Override
       public void evaluate() throws Throwable {
         if (TestBigQuery.this.datasetService != null) {
           throw new AssertionError(
               "BigQuery test was not shutdown previously. "
                   + "Table is'"
                   + table
                   + "'. "
                   + "Current test: "
                   + description.getDisplayName());
         }

         try {
           initializeBigQuery(description);
           base.evaluate();
         } finally {
           tearDown();
         }
       }
     };
   }

   private void initializeBigQuery(Description description)
       throws IOException, InterruptedException {

     this.datasetService = new BigQueryServicesImpl().getDatasetService(pipelineOptions);
     this.table = createTable(description);
   }

   private Table createTable(Description description) throws IOException, InterruptedException {
     TableReference tableReference =
         new TableReference()
             .setProjectId(pipelineOptions.getProject())
             .setDatasetId(pipelineOptions.getTargetDataset())
             .setTableId(createRandomizedName(description));

     table =
         new Table()
             .setTableReference(tableReference)
             .setSchema(BigQueryUtils.toTableSchema(schema))
             .setDescription(
                 "Table created for "
                     + description.getDisplayName()
                     + " by TestBigQueryRule. "
                     + "Should be automatically cleaned up after test completion.");

     if (datasetService.getTable(tableReference) != null) {
       throw new IllegalStateException(
           "Table '"
               + tableReference
               + "' already exists. "
               + "It should have been cleaned up by the test rule.");
     }

     datasetService.createTable(table);
     return table;
   }

   private void tearDown() throws IOException, InterruptedException {
     if (this.datasetService == null) {
       return;
     }

     try {
       if (table != null) {
         datasetService.deleteTable(table.getTableReference());
       }
     } finally {
       datasetService = null;
       table = null;
     }
   }

   static String createRandomizedName(Description description) throws IOException {
     StringBuilder topicName = new StringBuilder();

     if (description.getClassName() != null) {
       try {
         topicName.append(Class.forName(description.getClassName()).getSimpleName()).append("_");
       } catch (ClassNotFoundException e) {
         throw new RuntimeException(e);
       }
     }

     if (description.getMethodName() != null) {
       topicName.append(description.getMethodName()).append("_");
     }

     DATETIME_FORMAT.printTo(topicName, Instant.now());

     return topicName.toString()
         + "_"
         + String.valueOf(Math.abs(ThreadLocalRandom.current().nextLong()));
   }

   public String tableSpec() {
     return String.format(
         "%s:%s.%s",
         table.getTableReference().getProjectId(),
         table.getTableReference().getDatasetId(),
         table.getTableReference().getTableId());
   }

   public TableReference tableReference() {
     return table.getTableReference();
   }

   /**
    * Loads rows from BigQuery into {@link Row Rows} with given {@link Schema}.
    *
    * <p>Current implementation only supports flat {@link Row Rows} and target {@link Schema Schemas}
    * with {@link FieldType#STRING} fields only.
    */
   public List<Row> getFlatJsonRows(Schema rowSchema) {
     Bigquery bq = newBigQueryClient(pipelineOptions);
     return bqRowsToBeamRows(getSchema(bq), getTableRows(bq), rowSchema);
   }

   public RowsAssertion assertThatAllRows(Schema rowSchema) {
     return matcher -> duration -> pollAndAssert(rowSchema, matcher, duration);
   }

   private void pollAndAssert(
       Schema rowSchema, Matcher<Iterable<? extends Row>> matcher, Duration duration) {

     DateTime start = DateTime.now();
     while (true) {
       try {
         assertThat(getFlatJsonRows(rowSchema), matcher);
         break;
       } catch (AssertionError assertionError) {
         if (secondsBetween(start, DateTime.now()).isGreaterThan(duration.toStandardSeconds())) {
           throw assertionError;
         }
         sleep(15_000);
       }
     }
   }

   private List<Row> bqRowsToBeamRows(
       TableSchema bqSchema, List<TableRow> bqRows, Schema rowSchema) {
     if (bqRows == null) {
       return Collections.emptyList();
     }

     return bqRows.stream()
         .map(bqRow -> BigQueryUtils.toBeamRow(rowSchema, bqSchema, bqRow))
         .collect(Collectors.toList());
   }

   private TableSchema getSchema(Bigquery bq) {
     try {
       return bq.tables()
           .get(
               pipelineOptions.getProject(),
               pipelineOptions.getTargetDataset(),
               table.getTableReference().getTableId())
           .execute()
           .getSchema();
     } catch (IOException e) {
       throw new RuntimeException(e);
     }
   }

   private List<TableRow> getTableRows(Bigquery bq) {
     try {
       return bq.tabledata()
           .list(
               pipelineOptions.getProject(),
               pipelineOptions.getTargetDataset(),
               table.getTableReference().getTableId())
           .execute()
           .getRows();
     } catch (IOException e) {
       throw new RuntimeException(e);
     }
   }

   private static Bigquery newBigQueryClient(BigQueryOptions options) {
     return new Bigquery.Builder(
             Transport.getTransport(),
             Transport.getJsonFactory(),
             chainHttpRequestInitializer(
                 options.getGcpCredential(),
                 // Do not log 404. It clutters the output and is possibly even required by the
                 // caller.
                 new RetryHttpRequestInitializer(ImmutableList.of(404))))
         .setApplicationName(options.getAppName())
         .setGoogleClientRequestInitializer(options.getGoogleApiTrace())
         .build();
   }

   private static HttpRequestInitializer chainHttpRequestInitializer(
       Credentials credential, HttpRequestInitializer httpRequestInitializer) {
     if (credential == null) {
       return new ChainingHttpRequestInitializer(
           new NullCredentialInitializer(), httpRequestInitializer);
     } else {
       return new ChainingHttpRequestInitializer(
           new HttpCredentialsAdapter(credential), httpRequestInitializer);
     }
   }

   private void sleep(long l) {
     try {
       Thread.sleep(l);
     } catch (InterruptedException e) {
       throw new RuntimeException(e);
     }
   }

   /** Interface for creating a polling eventual assertion. */
   public interface RowsAssertion {
     PollingAssertion eventually(Matcher<Iterable<? extends Row>> matcher);
   }

   /** Interface to implement a polling assertion. */
   public interface PollingAssertion {
     void pollFor(Duration duration);
   }
 }
	/*
	* Licensed to the Apache Software Foundation (ASF) under one
	* or more contributor license agreements. See the NOTICE file
	* distributed with this work for additional information
	* regarding copyright ownership. The ASF licenses this file
	* to you 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 org.apache.beam.sdk.io.gcp.bigquery;

	import static org.joda.time.Seconds.secondsBetween;
	import static org.junit.Assert.assertThat;

	import com.google.api.client.http.HttpRequestInitializer;
	import com.google.api.services.bigquery.Bigquery;
	import com.google.api.services.bigquery.model.Table;
	import com.google.api.services.bigquery.model.TableReference;
	import com.google.api.services.bigquery.model.TableRow;
	import com.google.api.services.bigquery.model.TableSchema;
	import com.google.auth.Credentials;
	import com.google.auth.http.HttpCredentialsAdapter;
	import com.google.cloud.hadoop.util.ChainingHttpRequestInitializer;
	import java.io.IOException;
	import java.util.Collections;
	import java.util.List;
	import java.util.concurrent.ThreadLocalRandom;
	import java.util.stream.Collectors;
	import org.apache.beam.sdk.extensions.gcp.auth.NullCredentialInitializer;
	import org.apache.beam.sdk.extensions.gcp.util.RetryHttpRequestInitializer;
	import org.apache.beam.sdk.extensions.gcp.util.Transport;
	import org.apache.beam.sdk.schemas.Schema;
	import org.apache.beam.sdk.schemas.Schema.FieldType;
	import org.apache.beam.sdk.testing.TestPipeline;
	import org.apache.beam.sdk.testing.TestPipelineOptions;
	import org.apache.beam.sdk.values.Row;
	import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
	import org.hamcrest.Matcher;
	import org.joda.time.DateTime;
	import org.joda.time.Duration;
	import org.joda.time.Instant;
	import org.joda.time.format.DateTimeFormat;
	import org.joda.time.format.DateTimeFormatter;
	import org.junit.rules.TestRule;
	import org.junit.runner.Description;
	import org.junit.runners.model.Statement;

	/**
	* Test rule which creates a new table with specified schema, with randomized name and exposes few
	* APIs to work with it.
	*
	* <p>Deletes the table on test shutdown.
	*/
	public class TestBigQuery implements TestRule {
	private static final DateTimeFormatter DATETIME_FORMAT =
	DateTimeFormat.forPattern("YYYY_MM_dd_HH_mm_ss_SSS");

	private TestBigQueryOptions pipelineOptions;
	private Schema schema;
	private Table table;
	private BigQueryServices.DatasetService datasetService;

	/**
	* Creates an instance of this rule.
	*
	* <p>Loads GCP configuration from {@link TestPipelineOptions}.
	*/
	public static TestBigQuery create(Schema tableSchema) {
	return new TestBigQuery(
	TestPipeline.testingPipelineOptions().as(TestBigQueryOptions.class), tableSchema);
	}

	private TestBigQuery(TestBigQueryOptions pipelineOptions, Schema tableSchema) {
	this.pipelineOptions = pipelineOptions;
	this.schema = tableSchema;
	}

	@Override
	public Statement apply(Statement base, Description description) {
	return new Statement() {
	@Override
	public void evaluate() throws Throwable {
	if (TestBigQuery.this.datasetService != null) {
	throw new AssertionError(
	"BigQuery test was not shutdown previously. "
	+ "Table is'"
	+ table
	+ "'. "
	+ "Current test: "
	+ description.getDisplayName());
	}

	try {
	initializeBigQuery(description);
	base.evaluate();
	} finally {
	tearDown();
	}
	}
	};
	}

	private void initializeBigQuery(Description description)
	throws IOException, InterruptedException {

	this.datasetService = new BigQueryServicesImpl().getDatasetService(pipelineOptions);
	this.table = createTable(description);
	}

	private Table createTable(Description description) throws IOException, InterruptedException {
	TableReference tableReference =
	new TableReference()
	.setProjectId(pipelineOptions.getProject())
	.setDatasetId(pipelineOptions.getTargetDataset())
	.setTableId(createRandomizedName(description));

	table =
	new Table()
	.setTableReference(tableReference)
	.setSchema(BigQueryUtils.toTableSchema(schema))
	.setDescription(
	"Table created for "
	+ description.getDisplayName()
	+ " by TestBigQueryRule. "
	+ "Should be automatically cleaned up after test completion.");

	if (datasetService.getTable(tableReference) != null) {
	throw new IllegalStateException(
	"Table '"
	+ tableReference
	+ "' already exists. "
	+ "It should have been cleaned up by the test rule.");
	}

	datasetService.createTable(table);
	return table;
	}

	private void tearDown() throws IOException, InterruptedException {
	if (this.datasetService == null) {
	return;
	}

	try {
	if (table != null) {
	datasetService.deleteTable(table.getTableReference());
	}
	} finally {
	datasetService = null;
	table = null;
	}
	}

	static String createRandomizedName(Description description) throws IOException {
	StringBuilder topicName = new StringBuilder();

	if (description.getClassName() != null) {
	try {
	topicName.append(Class.forName(description.getClassName()).getSimpleName()).append("_");
	} catch (ClassNotFoundException e) {
	throw new RuntimeException(e);
	}
	}

	if (description.getMethodName() != null) {
	topicName.append(description.getMethodName()).append("_");
	}

	DATETIME_FORMAT.printTo(topicName, Instant.now());

	return topicName.toString()
	+ "_"
	+ String.valueOf(Math.abs(ThreadLocalRandom.current().nextLong()));
	}

	public String tableSpec() {
	return String.format(
	"%s:%s.%s",
	table.getTableReference().getProjectId(),
	table.getTableReference().getDatasetId(),
	table.getTableReference().getTableId());
	}

	public TableReference tableReference() {
	return table.getTableReference();
	}

	/**
	* Loads rows from BigQuery into {@link Row Rows} with given {@link Schema}.
	*
	* <p>Current implementation only supports flat {@link Row Rows} and target {@link Schema Schemas}
	* with {@link FieldType#STRING} fields only.
	*/
	public List<Row> getFlatJsonRows(Schema rowSchema) {
	Bigquery bq = newBigQueryClient(pipelineOptions);
	return bqRowsToBeamRows(getSchema(bq), getTableRows(bq), rowSchema);
	}

	public RowsAssertion assertThatAllRows(Schema rowSchema) {
	return matcher -> duration -> pollAndAssert(rowSchema, matcher, duration);
	}

	private void pollAndAssert(
	Schema rowSchema, Matcher<Iterable<? extends Row>> matcher, Duration duration) {

	DateTime start = DateTime.now();
	while (true) {
	try {
	assertThat(getFlatJsonRows(rowSchema), matcher);
	break;
	} catch (AssertionError assertionError) {
	if (secondsBetween(start, DateTime.now()).isGreaterThan(duration.toStandardSeconds())) {
	throw assertionError;
	}
	sleep(15_000);
	}
	}
	}

	private List<Row> bqRowsToBeamRows(
	TableSchema bqSchema, List<TableRow> bqRows, Schema rowSchema) {
	if (bqRows == null) {
	return Collections.emptyList();
	}

	return bqRows.stream()
	.map(bqRow -> BigQueryUtils.toBeamRow(rowSchema, bqSchema, bqRow))
	.collect(Collectors.toList());
	}

	private TableSchema getSchema(Bigquery bq) {
	try {
	return bq.tables()
	.get(
	pipelineOptions.getProject(),
	pipelineOptions.getTargetDataset(),
	table.getTableReference().getTableId())
	.execute()
	.getSchema();
	} catch (IOException e) {
	throw new RuntimeException(e);
	}
	}

	private List<TableRow> getTableRows(Bigquery bq) {
	try {
	return bq.tabledata()
	.list(
	pipelineOptions.getProject(),
	pipelineOptions.getTargetDataset(),
	table.getTableReference().getTableId())
	.execute()
	.getRows();
	} catch (IOException e) {
	throw new RuntimeException(e);
	}
	}

	private static Bigquery newBigQueryClient(BigQueryOptions options) {
	return new Bigquery.Builder(
	Transport.getTransport(),
	Transport.getJsonFactory(),
	chainHttpRequestInitializer(
	options.getGcpCredential(),
	// Do not log 404. It clutters the output and is possibly even required by the
	// caller.
	new RetryHttpRequestInitializer(ImmutableList.of(404))))
	.setApplicationName(options.getAppName())
	.setGoogleClientRequestInitializer(options.getGoogleApiTrace())
	.build();
	}

	private static HttpRequestInitializer chainHttpRequestInitializer(
	Credentials credential, HttpRequestInitializer httpRequestInitializer) {
	if (credential == null) {
	return new ChainingHttpRequestInitializer(
	new NullCredentialInitializer(), httpRequestInitializer);
	} else {
	return new ChainingHttpRequestInitializer(
	new HttpCredentialsAdapter(credential), httpRequestInitializer);
	}
	}

	private void sleep(long l) {
	try {
	Thread.sleep(l);
	} catch (InterruptedException e) {
	throw new RuntimeException(e);
	}
	}

	/** Interface for creating a polling eventual assertion. */
	public interface RowsAssertion {
	PollingAssertion eventually(Matcher<Iterable<? extends Row>> matcher);
	}

	/** Interface to implement a polling assertion. */
	public interface PollingAssertion {
	void pollFor(Duration duration);
	}
	}