org.apache.beam.it.gcp.bigquery

Class BigQueryResourceManager

java.lang.Object

org.apache.beam.it.gcp.bigquery.BigQueryResourceManager

All Implemented Interfaces:

org.apache.beam.it.common.ResourceManager

public final class BigQueryResourceManager
extends java.lang.Object
implements org.apache.beam.it.common.ResourceManager

Client for managing BigQuery resources.

The class supports one dataset, and multiple tables per dataset object.

The class is thread-safe.

Nested Class Summary

Nested Classes

Modifier and Type	Class and Description
static class	BigQueryResourceManager.Builder Builder for BigQueryResourceManager.

Method Summary

Modifier and Type	Method and Description
static BigQueryResourceManager.Builder	builder(java.lang.String testId, java.lang.String projectId, com.google.auth.Credentials credentials)
void	cleanupAll() Deletes all created resources (dataset and tables) and cleans up the BigQuery client, making the manager object unusable.
java.lang.String	createDataset(java.lang.String region) Create a BigQuery dataset in which all tables will exist.
com.google.cloud.bigquery.TableId	createTable(java.lang.String tableName, com.google.cloud.bigquery.Schema schema) Creates a table within the current dataset given a table name and schema.
com.google.cloud.bigquery.TableId	createTable(java.lang.String tableName, com.google.cloud.bigquery.Schema schema, java.lang.Long expirationTimeMillis) Creates a table within the current dataset given a table name and schema.
java.lang.String	getDatasetId() Return the dataset ID this Resource Manager uses to create and manage tables in.
java.lang.String	getProjectId()
java.lang.Long	getRowCount(java.lang.String table) Gets the number of rows in the table.
com.google.cloud.bigquery.TableResult	readTable(java.lang.String tableName) Reads all the rows in a table and returns a TableResult containing a JSON string representation.
com.google.cloud.bigquery.TableResult	readTable(java.lang.String tableName, int numRows) Reads number of rows in a table and returns a TableResult containing a JSON string representation.
com.google.cloud.bigquery.TableResult	readTable(com.google.cloud.bigquery.TableId table) Reads all the rows in a table and returns a TableResult containing a JSON string representation.
com.google.cloud.bigquery.TableResult	readTable(com.google.cloud.bigquery.TableId table, int numRows) Reads number of rows in a table and returns a TableResult containing a JSON string representation.
com.google.cloud.bigquery.TableResult	runQuery(java.lang.String query) Runs the specified query.
void	write(java.lang.String tableName, com.google.cloud.bigquery.InsertAllRequest.RowToInsert row) Writes a given row into a table.
void	write(java.lang.String tableName, java.util.List<com.google.cloud.bigquery.InsertAllRequest.RowToInsert> rows) Writes a collection of table rows into a single table.

Methods inherited from class java.lang.Object

clone, equals, finalize, getClass, hashCode, notify, notifyAll, toString, wait, wait, wait

Method Detail

builder

public static BigQueryResourceManager.Builder builder(java.lang.String testId,
                                                      java.lang.String projectId,
                                                      com.google.auth.Credentials credentials)

getProjectId

public java.lang.String getProjectId()

getDatasetId

public java.lang.String getDatasetId()

Return the dataset ID this Resource Manager uses to create and manage tables in.

Returns:

the dataset ID.

createDataset

public java.lang.String createDataset(java.lang.String region)
                               throws BigQueryResourceManagerException

Create a BigQuery dataset in which all tables will exist.

Parameters:

region - the region to store the dataset in.

Returns:

Dataset id that was created

Throws:

BigQueryResourceManagerException - if there is an error creating the dataset in BigQuery.

createTable

public com.google.cloud.bigquery.TableId createTable(java.lang.String tableName,
                                                     com.google.cloud.bigquery.Schema schema)
                                              throws BigQueryResourceManagerException

Creates a table within the current dataset given a table name and schema.

This table will automatically expire 1 hour after creation if not cleaned up manually or by calling the cleanupAll() method.

Note: Implementations may do dataset creation here, if one does not already exist.

Parameters:

tableName - The name of the table.

schema - A schema object that defines the table.

Returns:

The TableId (reference) to the table

Throws:

BigQueryResourceManagerException - if there is an error creating the table in BigQuery.

createTable

public com.google.cloud.bigquery.TableId createTable(java.lang.String tableName,
                                                     com.google.cloud.bigquery.Schema schema,
                                                     java.lang.Long expirationTimeMillis)
                                              throws BigQueryResourceManagerException

Creates a table within the current dataset given a table name and schema.

This table will automatically expire at the time specified by expirationTime if not cleaned up manually or by calling the cleanupAll() method.

Note: Implementations may do dataset creation here, if one does not already exist.

Parameters:

tableName - The name of the table.

schema - A schema object that defines the table.

expirationTimeMillis - Sets the time when this table expires, in milliseconds since the epoch.

Returns:

The TableId (reference) to the table

Throws:

BigQueryResourceManagerException - if there is an error creating the table in BigQuery.

write

public void write(java.lang.String tableName,
                  com.google.cloud.bigquery.InsertAllRequest.RowToInsert row)
           throws BigQueryResourceManagerException

Writes a given row into a table. This method requires createTable(String, Schema) to be called for the target table beforehand.

Parameters:

tableName - The name of the table to insert the given row into.

row - A row object representing the table row.

Throws:

BigQueryResourceManagerException - if method is called after resources have been cleaned up, if the manager object has no dataset, if the table does not exist or if there is an Exception when attempting to insert the rows.

write

public void write(java.lang.String tableName,
                  java.util.List<com.google.cloud.bigquery.InsertAllRequest.RowToInsert> rows)
           throws BigQueryResourceManagerException

Writes a collection of table rows into a single table. This method requires createTable(String, Schema) to be called for the target table beforehand.

Parameters:

tableName - The name of the table to insert the given rows into.

rows - A collection of table rows.

Throws:

BigQueryResourceManagerException - if method is called after resources have been cleaned up, if the manager object has no dataset, if the table does not exist or if there is an Exception when attempting to insert the rows.

runQuery

public com.google.cloud.bigquery.TableResult runQuery(java.lang.String query)

Runs the specified query.

Parameters:

query - the query to execute

getRowCount

public java.lang.Long getRowCount(java.lang.String table)

Gets the number of rows in the table.

Parameters:

table - the name of the table

readTable

public com.google.cloud.bigquery.TableResult readTable(com.google.cloud.bigquery.TableId table)
                                                throws BigQueryResourceManagerException

Reads all the rows in a table and returns a TableResult containing a JSON string representation. This method requires createTable(String, Schema) to be called for the target table beforehand.

Parameters:

table - The table reference to read rows from.

Returns:

A TableResult containing all the rows in the table in JSON.

Throws:

BigQueryResourceManagerException - if method is called after resources have been cleaned up, if the manager object has no dataset, if the table does not exist or if there is an Exception when attempting to insert the rows.

readTable

public com.google.cloud.bigquery.TableResult readTable(java.lang.String tableName)
                                                throws BigQueryResourceManagerException

Reads all the rows in a table and returns a TableResult containing a JSON string representation. This method requires createTable(String, Schema) to be called for the target table beforehand.

Parameters:

tableName - The name of the table to read rows from.

Returns:

A TableResult containing all the rows in the table in JSON.

Throws:

BigQueryResourceManagerException - if method is called after resources have been cleaned up, if the manager object has no dataset, if the table does not exist or if there is an Exception when attempting to insert the rows.

readTable

public com.google.cloud.bigquery.TableResult readTable(com.google.cloud.bigquery.TableId table,
                                                       int numRows)
                                                throws BigQueryResourceManagerException

Reads number of rows in a table and returns a TableResult containing a JSON string representation. This method requires createTable(String, Schema) to be called for the target table beforehand.

Parameters:

table - The table reference to read rows from.

Returns:

A TableResult containing all the rows in the table in JSON.

Throws:

BigQueryResourceManagerException - if method is called after resources have been cleaned up, if the manager object has no dataset, if the table does not exist or if there is an Exception when attempting to insert the rows.

readTable

public com.google.cloud.bigquery.TableResult readTable(java.lang.String tableName,
                                                       int numRows)
                                                throws BigQueryResourceManagerException

Reads number of rows in a table and returns a TableResult containing a JSON string representation. This method requires createTable(String, Schema) to be called for the target table beforehand.

Parameters:

tableName - The name of the table to read rows from.

numRows - number of rows to read from the table

Returns:

A TableResult containing all the rows in the table in JSON.

Throws:

BigQueryResourceManagerException - if method is called after resources have been cleaned up, if the manager object has no dataset, if the table does not exist or if there is an Exception when attempting to insert the rows.

cleanupAll

public void cleanupAll()
                throws BigQueryResourceManagerException

Deletes all created resources (dataset and tables) and cleans up the BigQuery client, making the manager object unusable.

Specified by:

cleanupAll in interface org.apache.beam.it.common.ResourceManager

Throws:

BigQueryResourceManagerException - if there is an error deleting the tables or dataset in BigQuery.