You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@beam.apache.org by cc...@apache.org on 2019/02/09 02:20:33 UTC
[beam] branch master updated: [BEAM-5959] KMS support for BigQuery
This is an automated email from the ASF dual-hosted git repository.
ccy pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new d4e05a1 [BEAM-5959] KMS support for BigQuery
new 49403c0 Merge pull request #7744 from udim/cmek-bq
d4e05a1 is described below
commit d4e05a153a7db0958e7174abf8927c8625e172bf
Author: Udi Meiri <eh...@google.com>
AuthorDate: Mon Feb 4 17:30:03 2019 -0800
[BEAM-5959] KMS support for BigQuery
Adds withKmsKey() methods to reads and writes. Affects tables created
tables by: streaming writes, batched writes, query temporary tables,
table copies (for very large batched writes).
---
sdks/java/io/google-cloud-platform/build.gradle | 2 -
.../beam/sdk/io/gcp/bigquery/BatchLoads.java | 17 ++-
.../beam/sdk/io/gcp/bigquery/BigQueryIO.java | 30 +++++-
.../sdk/io/gcp/bigquery/BigQueryQuerySource.java | 22 +++-
.../beam/sdk/io/gcp/bigquery/CreateTables.java | 22 +++-
.../beam/sdk/io/gcp/bigquery/StreamingInserts.java | 38 +++++--
.../beam/sdk/io/gcp/bigquery/WriteRename.java | 16 ++-
.../beam/sdk/io/gcp/bigquery/WriteTables.java | 12 ++-
.../beam/sdk/io/gcp/testing/BigqueryClient.java | 30 ++++++
.../sdk/io/gcp/bigquery/BigQueryIOReadTest.java | 22 ++++
.../sdk/io/gcp/bigquery/BigQueryIOWriteTest.java | 7 +-
.../beam/sdk/io/gcp/bigquery/BigQueryKmsKeyIT.java | 119 +++++++++++++++++++++
.../beam/sdk/io/gcp/bigquery/FakeJobService.java | 3 +-
13 files changed, 303 insertions(+), 37 deletions(-)
diff --git a/sdks/java/io/google-cloud-platform/build.gradle b/sdks/java/io/google-cloud-platform/build.gradle
index 159ae73..532d3b8 100644
--- a/sdks/java/io/google-cloud-platform/build.gradle
+++ b/sdks/java/io/google-cloud-platform/build.gradle
@@ -111,13 +111,11 @@ task integrationTestKms(type: Test) {
def gcpProject = project.findProperty('gcpProject') ?: 'apache-beam-testing'
def gcpTempRoot = project.findProperty('gcpTempRoot') ?: 'gs://temp-storage-for-end-to-end-tests'
def dataflowKmsKey = project.findProperty('dataflowKmsKey') ?: "projects/apache-beam-testing/locations/global/keyRings/beam-it/cryptoKeys/test"
- def kmsKey = project.findProperty('kmsKey') ?: dataflowKmsKey
systemProperty "beamTestPipelineOptions", JsonOutput.toJson([
"--runner=DirectRunner",
"--project=${gcpProject}",
"--tempRoot=${gcpTempRoot}",
"--dataflowKmsKey=${dataflowKmsKey}",
- "--kmsKey=${kmsKey}",
])
// Disable Gradle cache: these ITs interact with live service that should always be considered "out of date"
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchLoads.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchLoads.java
index abbcd7b..a849ca3 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchLoads.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchLoads.java
@@ -129,6 +129,7 @@ class BatchLoads<DestinationT>
private Duration triggeringFrequency;
private ValueProvider<String> customGcsTempLocation;
private ValueProvider<String> loadJobProjectId;
+ private String kmsKey;
// The maximum number of times to retry failed load or copy jobs.
private int maxRetryJobs = DEFAULT_MAX_RETRY_JOBS;
@@ -141,7 +142,8 @@ class BatchLoads<DestinationT>
Coder<DestinationT> destinationCoder,
ValueProvider<String> customGcsTempLocation,
@Nullable ValueProvider<String> loadJobProjectId,
- boolean ignoreUnknownValues) {
+ boolean ignoreUnknownValues,
+ @Nullable String kmsKey) {
bigQueryServices = new BigQueryServicesImpl();
this.writeDisposition = writeDisposition;
this.createDisposition = createDisposition;
@@ -157,6 +159,7 @@ class BatchLoads<DestinationT>
this.customGcsTempLocation = customGcsTempLocation;
this.loadJobProjectId = loadJobProjectId;
this.ignoreUnknownValues = ignoreUnknownValues;
+ this.kmsKey = kmsKey;
}
void setTestServices(BigQueryServices bigQueryServices) {
@@ -319,7 +322,8 @@ class BatchLoads<DestinationT>
loadJobIdPrefixView,
writeDisposition,
createDisposition,
- maxRetryJobs))
+ maxRetryJobs,
+ kmsKey))
.withSideInputs(loadJobIdPrefixView));
writeSinglePartition(partitions.get(singlePartitionTag), loadJobIdPrefixView);
return writeResult(p);
@@ -381,7 +385,8 @@ class BatchLoads<DestinationT>
loadJobIdPrefixView,
writeDisposition,
createDisposition,
- maxRetryJobs))
+ maxRetryJobs,
+ kmsKey))
.withSideInputs(loadJobIdPrefixView));
writeSinglePartition(partitions.get(singlePartitionTag), loadJobIdPrefixView);
return writeResult(p);
@@ -554,7 +559,8 @@ class BatchLoads<DestinationT>
dynamicDestinations,
loadJobProjectId,
maxRetryJobs,
- ignoreUnknownValues));
+ ignoreUnknownValues,
+ kmsKey));
}
// In the case where the files fit into a single load job, there's no need to write temporary
@@ -586,7 +592,8 @@ class BatchLoads<DestinationT>
dynamicDestinations,
loadJobProjectId,
maxRetryJobs,
- ignoreUnknownValues));
+ ignoreUnknownValues,
+ kmsKey));
}
private WriteResult writeResult(Pipeline p) {
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java
index 35edb37..00012d5 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java
@@ -555,6 +555,8 @@ public class BigQueryIO {
abstract Builder<T> setParseFn(SerializableFunction<SchemaAndRecord, T> parseFn);
abstract Builder<T> setCoder(Coder<T> coder);
+
+ abstract Builder<T> setKmsKey(String kmsKey);
}
@Nullable
@@ -586,6 +588,9 @@ public class BigQueryIO {
@Nullable
abstract Coder<T> getCoder();
+ @Nullable
+ abstract String getKmsKey();
+
/**
* An enumeration type for the priority of a query.
*
@@ -643,7 +648,8 @@ public class BigQueryIO {
coder,
getParseFn(),
MoreObjects.firstNonNull(getQueryPriority(), QueryPriority.BATCH),
- getQueryLocation());
+ getQueryLocation(),
+ getKmsKey());
}
return source;
}
@@ -904,6 +910,11 @@ public class BigQueryIO {
return toBuilder().setCoder(coder).build();
}
+ /** For query sources, use this Cloud KMS key to encrypt any temporary tables created. */
+ public TypedRead<T> withKmsKey(String kmsKey) {
+ return toBuilder().setKmsKey(kmsKey).build();
+ }
+
/** See {@link Read#from(String)}. */
public TypedRead<T> from(String tableSpec) {
return from(StaticValueProvider.of(tableSpec));
@@ -1173,6 +1184,9 @@ public class BigQueryIO {
abstract Boolean getIgnoreUnknownValues();
+ @Nullable
+ abstract String getKmsKey();
+
abstract Builder<T> toBuilder();
@AutoValue.Builder
@@ -1228,6 +1242,8 @@ public class BigQueryIO {
abstract Builder<T> setIgnoreUnknownValues(Boolean ignoreUnknownValues);
+ abstract Builder<T> setKmsKey(String kmsKey);
+
abstract Write<T> build();
}
@@ -1430,7 +1446,7 @@ public class BigQueryIO {
}
/**
- * Specfies a policy for handling fPailed inserts.
+ * Specfies a policy for handling failed inserts.
*
* <p>Currently this only is allowed when writing an unbounded collection to BigQuery. Bounded
* collections are written using batch load jobs, so we don't get per-element failures.
@@ -1534,6 +1550,10 @@ public class BigQueryIO {
return toBuilder().setIgnoreUnknownValues(true).build();
}
+ Write<T> withKmsKey(String kmsKey) {
+ return toBuilder().setKmsKey(kmsKey).build();
+ }
+
@VisibleForTesting
/** This method is for test usage only */
public Write<T> withTestServices(BigQueryServices testServices) {
@@ -1737,7 +1757,8 @@ public class BigQueryIO {
.withTestServices(getBigQueryServices())
.withExtendedErrorInfo(getExtendedErrorInfo())
.withSkipInvalidRows(getSkipInvalidRows())
- .withIgnoreUnknownValues(getIgnoreUnknownValues());
+ .withIgnoreUnknownValues(getIgnoreUnknownValues())
+ .withKmsKey(getKmsKey());
return rowsWithDestination.apply(streamingInserts);
} else {
checkArgument(
@@ -1753,7 +1774,8 @@ public class BigQueryIO {
destinationCoder,
getCustomGcsTempLocation(),
getLoadJobProjectId(),
- getIgnoreUnknownValues());
+ getIgnoreUnknownValues(),
+ getKmsKey());
batchLoads.setTestServices(getBigQueryServices());
if (getMaxFilesPerBundle() != null) {
batchLoads.setMaxNumWritersPerBundle(getMaxFilesPerBundle());
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryQuerySource.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryQuerySource.java
index e367d34..a937959 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryQuerySource.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryQuerySource.java
@@ -21,6 +21,7 @@ import static org.apache.beam.sdk.io.gcp.bigquery.BigQueryHelpers.createJobIdTok
import static org.apache.beam.sdk.io.gcp.bigquery.BigQueryHelpers.createTempTableReference;
import static org.apache.beam.vendor.guava.v20_0.com.google.common.base.Preconditions.checkNotNull;
+import com.google.api.services.bigquery.model.EncryptionConfiguration;
import com.google.api.services.bigquery.model.Job;
import com.google.api.services.bigquery.model.JobConfigurationQuery;
import com.google.api.services.bigquery.model.JobReference;
@@ -57,7 +58,8 @@ class BigQueryQuerySource<T> extends BigQuerySourceBase<T> {
Coder<T> coder,
SerializableFunction<SchemaAndRecord, T> parseFn,
QueryPriority priority,
- String location) {
+ String location,
+ String kmsKey) {
return new BigQueryQuerySource<>(
stepUuid,
query,
@@ -67,7 +69,8 @@ class BigQueryQuerySource<T> extends BigQuerySourceBase<T> {
coder,
parseFn,
priority,
- location);
+ location,
+ kmsKey);
}
private final ValueProvider<String> query;
@@ -76,6 +79,7 @@ class BigQueryQuerySource<T> extends BigQuerySourceBase<T> {
private transient AtomicReference<JobStatistics> dryRunJobStats;
private final QueryPriority priority;
private final String location;
+ private final String kmsKey;
private BigQueryQuerySource(
String stepUuid,
@@ -86,7 +90,8 @@ class BigQueryQuerySource<T> extends BigQuerySourceBase<T> {
Coder<T> coder,
SerializableFunction<SchemaAndRecord, T> parseFn,
QueryPriority priority,
- String location) {
+ String location,
+ String kmsKey) {
super(stepUuid, bqServices, coder, parseFn);
this.query = checkNotNull(query, "query");
this.flattenResults = checkNotNull(flattenResults, "flattenResults");
@@ -94,6 +99,7 @@ class BigQueryQuerySource<T> extends BigQuerySourceBase<T> {
this.dryRunJobStats = new AtomicReference<>();
this.priority = priority;
this.location = location;
+ this.kmsKey = kmsKey;
}
@Override
@@ -142,7 +148,8 @@ class BigQueryQuerySource<T> extends BigQuerySourceBase<T> {
bqOptions.getProject(),
tableToExtract,
bqServices.getJobService(bqOptions),
- location);
+ location,
+ kmsKey);
return tableToExtract;
}
@@ -183,7 +190,8 @@ class BigQueryQuerySource<T> extends BigQuerySourceBase<T> {
String executingProject,
TableReference destinationTable,
JobService jobService,
- String bqLocation)
+ String bqLocation,
+ String kmsKey)
throws IOException, InterruptedException {
// Generate a transient (random) query job ID, because this code may be retried after the
// temporary dataset and table have already been deleted by a previous attempt -
@@ -211,6 +219,10 @@ class BigQueryQuerySource<T> extends BigQuerySourceBase<T> {
// Overwrite contents of the temporary table - it can only already exist if this
// is a retry of the splitting task, in which case we must not produce duplicate data.
.setWriteDisposition("WRITE_TRUNCATE");
+ if (kmsKey != null) {
+ queryConfig.setDestinationEncryptionConfiguration(
+ new EncryptionConfiguration().setKmsKeyName(kmsKey));
+ }
jobService.startQueryJob(jobRef, queryConfig);
Job job = jobService.pollJob(jobRef, JOB_POLL_MAX_RETRIES);
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/CreateTables.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/CreateTables.java
index 9d6268a..2616e55 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/CreateTables.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/CreateTables.java
@@ -19,6 +19,7 @@ package org.apache.beam.sdk.io.gcp.bigquery;
import static org.apache.beam.vendor.guava.v20_0.com.google.common.base.Preconditions.checkArgument;
+import com.google.api.services.bigquery.model.EncryptionConfiguration;
import com.google.api.services.bigquery.model.Table;
import com.google.api.services.bigquery.model.TableReference;
import com.google.api.services.bigquery.model.TableRow;
@@ -51,6 +52,7 @@ public class CreateTables<DestinationT>
private final CreateDisposition createDisposition;
private final BigQueryServices bqServices;
private final DynamicDestinations<?, DestinationT> dynamicDestinations;
+ private final String kmsKey;
/**
* The list of tables created so far, so we don't try the creation each time.
@@ -63,20 +65,26 @@ public class CreateTables<DestinationT>
public CreateTables(
CreateDisposition createDisposition,
DynamicDestinations<?, DestinationT> dynamicDestinations) {
- this(createDisposition, new BigQueryServicesImpl(), dynamicDestinations);
+ this(createDisposition, new BigQueryServicesImpl(), dynamicDestinations, null);
}
private CreateTables(
CreateDisposition createDisposition,
BigQueryServices bqServices,
- DynamicDestinations<?, DestinationT> dynamicDestinations) {
+ DynamicDestinations<?, DestinationT> dynamicDestinations,
+ String kmsKey) {
this.createDisposition = createDisposition;
this.bqServices = bqServices;
this.dynamicDestinations = dynamicDestinations;
+ this.kmsKey = kmsKey;
}
CreateTables<DestinationT> withTestServices(BigQueryServices bqServices) {
- return new CreateTables<>(createDisposition, bqServices, dynamicDestinations);
+ return new CreateTables<>(createDisposition, bqServices, dynamicDestinations, kmsKey);
+ }
+
+ CreateTables<DestinationT> withKmsKey(String kmsKey) {
+ return new CreateTables<>(createDisposition, bqServices, dynamicDestinations, kmsKey);
}
@Override
@@ -140,7 +148,7 @@ public class CreateTables<DestinationT>
// every thread from attempting a create and overwhelming our BigQuery quota.
synchronized (createdTables) {
if (!createdTables.contains(tableSpec)) {
- tryCreateTable(context, destination, tableDestination, tableSpec);
+ tryCreateTable(context, destination, tableDestination, tableSpec, kmsKey);
}
}
}
@@ -151,7 +159,8 @@ public class CreateTables<DestinationT>
ProcessContext context,
DestinationT destination,
TableDestination tableDestination,
- String tableSpec) {
+ String tableSpec,
+ String kmsKey) {
DatasetService datasetService =
bqServices.getDatasetService(context.getPipelineOptions().as(BigQueryOptions.class));
TableReference tableReference = tableDestination.getTableReference().clone();
@@ -178,6 +187,9 @@ public class CreateTables<DestinationT>
if (tableDestination.getTimePartitioning() != null) {
table.setTimePartitioning(tableDestination.getTimePartitioning());
}
+ if (kmsKey != null) {
+ table.setEncryptionConfiguration(new EncryptionConfiguration().setKmsKeyName(kmsKey));
+ }
datasetService.createTable(table);
}
} catch (Exception e) {
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StreamingInserts.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StreamingInserts.java
index b22cb23..87adb24 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StreamingInserts.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StreamingInserts.java
@@ -36,6 +36,7 @@ public class StreamingInserts<DestinationT>
private boolean extendedErrorInfo;
private final boolean skipInvalidRows;
private final boolean ignoreUnknownValues;
+ private final String kmsKey;
/** Constructor. */
public StreamingInserts(
@@ -48,7 +49,8 @@ public class StreamingInserts<DestinationT>
InsertRetryPolicy.alwaysRetry(),
false,
false,
- false);
+ false,
+ null);
}
/** Constructor. */
@@ -59,7 +61,8 @@ public class StreamingInserts<DestinationT>
InsertRetryPolicy retryPolicy,
boolean extendedErrorInfo,
boolean skipInvalidRows,
- boolean ignoreUnknownValues) {
+ boolean ignoreUnknownValues,
+ String kmsKey) {
this.createDisposition = createDisposition;
this.dynamicDestinations = dynamicDestinations;
this.bigQueryServices = bigQueryServices;
@@ -67,6 +70,7 @@ public class StreamingInserts<DestinationT>
this.extendedErrorInfo = extendedErrorInfo;
this.skipInvalidRows = skipInvalidRows;
this.ignoreUnknownValues = ignoreUnknownValues;
+ this.kmsKey = kmsKey;
}
/** Specify a retry policy for failed inserts. */
@@ -78,7 +82,8 @@ public class StreamingInserts<DestinationT>
retryPolicy,
extendedErrorInfo,
skipInvalidRows,
- ignoreUnknownValues);
+ ignoreUnknownValues,
+ kmsKey);
}
/** Specify whether to use extended error info or not. */
@@ -90,7 +95,8 @@ public class StreamingInserts<DestinationT>
retryPolicy,
extendedErrorInfo,
skipInvalidRows,
- ignoreUnknownValues);
+ ignoreUnknownValues,
+ kmsKey);
}
StreamingInserts<DestinationT> withSkipInvalidRows(boolean skipInvalidRows) {
@@ -101,7 +107,8 @@ public class StreamingInserts<DestinationT>
retryPolicy,
extendedErrorInfo,
skipInvalidRows,
- ignoreUnknownValues);
+ ignoreUnknownValues,
+ kmsKey);
}
StreamingInserts<DestinationT> withIgnoreUnknownValues(boolean ignoreUnknownValues) {
@@ -112,7 +119,20 @@ public class StreamingInserts<DestinationT>
retryPolicy,
extendedErrorInfo,
skipInvalidRows,
- ignoreUnknownValues);
+ ignoreUnknownValues,
+ kmsKey);
+ }
+
+ StreamingInserts<DestinationT> withKmsKey(String kmsKey) {
+ return new StreamingInserts<>(
+ createDisposition,
+ dynamicDestinations,
+ bigQueryServices,
+ retryPolicy,
+ extendedErrorInfo,
+ skipInvalidRows,
+ ignoreUnknownValues,
+ kmsKey);
}
StreamingInserts<DestinationT> withTestServices(BigQueryServices bigQueryServices) {
@@ -123,7 +143,8 @@ public class StreamingInserts<DestinationT>
retryPolicy,
extendedErrorInfo,
skipInvalidRows,
- ignoreUnknownValues);
+ ignoreUnknownValues,
+ kmsKey);
}
@Override
@@ -132,7 +153,8 @@ public class StreamingInserts<DestinationT>
input.apply(
"CreateTables",
new CreateTables<>(createDisposition, dynamicDestinations)
- .withTestServices(bigQueryServices));
+ .withTestServices(bigQueryServices)
+ .withKmsKey(kmsKey));
return writes.apply(
new StreamingWriteTables()
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteRename.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteRename.java
index c9ed9ed..4baf8fb 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteRename.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteRename.java
@@ -17,6 +17,7 @@
*/
package org.apache.beam.sdk.io.gcp.bigquery;
+import com.google.api.services.bigquery.model.EncryptionConfiguration;
import com.google.api.services.bigquery.model.JobConfigurationTableCopy;
import com.google.api.services.bigquery.model.JobReference;
import com.google.api.services.bigquery.model.TableReference;
@@ -57,6 +58,7 @@ class WriteRename extends DoFn<Iterable<KV<TableDestination, String>>, Void> {
private final WriteDisposition firstPaneWriteDisposition;
private final CreateDisposition firstPaneCreateDisposition;
private final int maxRetryJobs;
+ private final String kmsKey;
private static class PendingJobData {
final BigQueryHelpers.PendingJob retryJob;
@@ -80,12 +82,14 @@ class WriteRename extends DoFn<Iterable<KV<TableDestination, String>>, Void> {
PCollectionView<String> jobIdToken,
WriteDisposition writeDisposition,
CreateDisposition createDisposition,
- int maxRetryJobs) {
+ int maxRetryJobs,
+ String kmsKey) {
this.bqServices = bqServices;
this.jobIdToken = jobIdToken;
this.firstPaneWriteDisposition = writeDisposition;
this.firstPaneCreateDisposition = createDisposition;
this.maxRetryJobs = maxRetryJobs;
+ this.kmsKey = kmsKey;
}
@StartBundle
@@ -161,7 +165,8 @@ class WriteRename extends DoFn<Iterable<KV<TableDestination, String>>, Void> {
finalTableDestination.getTableReference(),
tempTables,
writeDisposition,
- createDisposition);
+ createDisposition,
+ kmsKey);
return new PendingJobData(retryJob, finalTableDestination, tempTables);
}
@@ -172,13 +177,18 @@ class WriteRename extends DoFn<Iterable<KV<TableDestination, String>>, Void> {
TableReference ref,
List<TableReference> tempTables,
WriteDisposition writeDisposition,
- CreateDisposition createDisposition) {
+ CreateDisposition createDisposition,
+ String kmsKey) {
JobConfigurationTableCopy copyConfig =
new JobConfigurationTableCopy()
.setSourceTables(tempTables)
.setDestinationTable(ref)
.setWriteDisposition(writeDisposition.name())
.setCreateDisposition(createDisposition.name());
+ if (kmsKey != null) {
+ copyConfig.setDestinationEncryptionConfiguration(
+ new EncryptionConfiguration().setKmsKeyName(kmsKey));
+ }
String bqLocation =
BigQueryHelpers.getDatasetLocation(datasetService, ref.getProjectId(), ref.getDatasetId());
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteTables.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteTables.java
index 3d94857..8a20eb7 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteTables.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteTables.java
@@ -19,6 +19,7 @@ package org.apache.beam.sdk.io.gcp.bigquery;
import static org.apache.beam.vendor.guava.v20_0.com.google.common.base.Preconditions.checkArgument;
+import com.google.api.services.bigquery.model.EncryptionConfiguration;
import com.google.api.services.bigquery.model.JobConfigurationLoad;
import com.google.api.services.bigquery.model.JobReference;
import com.google.api.services.bigquery.model.TableReference;
@@ -68,7 +69,7 @@ import org.slf4j.LoggerFactory;
/**
* Writes partitions to BigQuery tables.
*
- * <p>The input is a list of files corresponding to each partition of a table. loadThese files are
+ * <p>The input is a list of files corresponding to each partition of a table. These files are
* loaded into a temporary table (or into the final table if there is only one partition). The
* output is a {@link KV} mapping each final table to a list of the temporary tables containing its
* data.
@@ -95,6 +96,7 @@ class WriteTables<DestinationT>
private final ValueProvider<String> loadJobProjectId;
private final int maxRetryJobs;
private final boolean ignoreUnknownValues;
+ @Nullable private final String kmsKey;
private class WriteTablesDoFn
extends DoFn<KV<ShardedKey<DestinationT>, List<String>>, KV<TableDestination, String>> {
@@ -271,7 +273,8 @@ class WriteTables<DestinationT>
DynamicDestinations<?, DestinationT> dynamicDestinations,
@Nullable ValueProvider<String> loadJobProjectId,
int maxRetryJobs,
- boolean ignoreUnknownValues) {
+ boolean ignoreUnknownValues,
+ String kmsKey) {
this.tempTable = tempTable;
this.bqServices = bqServices;
this.loadJobIdPrefixView = loadJobIdPrefixView;
@@ -284,6 +287,7 @@ class WriteTables<DestinationT>
this.loadJobProjectId = loadJobProjectId;
this.maxRetryJobs = maxRetryJobs;
this.ignoreUnknownValues = ignoreUnknownValues;
+ this.kmsKey = kmsKey;
}
@Override
@@ -339,6 +343,10 @@ class WriteTables<DestinationT>
if (timePartitioning != null) {
loadConfig.setTimePartitioning(timePartitioning);
}
+ if (kmsKey != null) {
+ loadConfig.setDestinationEncryptionConfiguration(
+ new EncryptionConfiguration().setKmsKeyName(kmsKey));
+ }
String projectId = loadJobProjectId == null ? ref.getProjectId() : loadJobProjectId.get();
String bqLocation =
BigQueryHelpers.getDatasetLocation(datasetService, ref.getProjectId(), ref.getDatasetId());
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/testing/BigqueryClient.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/testing/BigqueryClient.java
index 7f94ff9..abf8a2e 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/testing/BigqueryClient.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/testing/BigqueryClient.java
@@ -500,4 +500,34 @@ public class BigqueryClient {
MAX_QUERY_RETRIES, tableName),
lastException);
}
+
+ public Table getTableResource(String projectId, String datasetId, String tableId)
+ throws IOException, InterruptedException {
+ Sleeper sleeper = Sleeper.DEFAULT;
+ BackOff backoff = BackOffAdapter.toGcpBackOff(BACKOFF_FACTORY.backoff());
+ IOException lastException = null;
+ do {
+ if (lastException != null) {
+ LOG.warn("Retrying tables.get ({}) after exception", tableId, lastException);
+ }
+ try {
+ Table response = this.bqClient.tables().get(projectId, datasetId, tableId).execute();
+ if (response != null) {
+ return response;
+ } else {
+ lastException =
+ new IOException("Expected valid response from tables.get, but received null.");
+ }
+ } catch (IOException e) {
+ // ignore and retry
+ lastException = e;
+ }
+ } while (BackOffUtils.next(sleeper, backoff));
+
+ throw new RuntimeException(
+ String.format(
+ "Unable to get BigQuery response after retrying %d times for tables.get (%s)",
+ MAX_QUERY_RETRIES, tableId),
+ lastException);
+ }
}
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOReadTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOReadTest.java
index eb132ab..db8d134 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOReadTest.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOReadTest.java
@@ -133,6 +133,10 @@ public class BigQueryIOReadTest implements Serializable {
checkReadQueryObjectWithValidate(read, query, true);
}
+ private void checkTypedReadQueryObject(BigQueryIO.TypedRead read, String query, String kmsKey) {
+ checkTypedReadQueryObjectWithValidate(read, query, kmsKey, true);
+ }
+
private void checkReadTableObjectWithValidate(
BigQueryIO.Read read, String project, String dataset, String table, boolean validate) {
assertEquals(project, read.getTable().getProjectId());
@@ -149,6 +153,14 @@ public class BigQueryIOReadTest implements Serializable {
assertEquals(validate, read.getValidate());
}
+ private void checkTypedReadQueryObjectWithValidate(
+ BigQueryIO.TypedRead read, String query, String kmsKey, boolean validate) {
+ assertNull(read.getTable());
+ assertEquals(query, read.getQuery().get());
+ assertEquals(kmsKey, read.getKmsKey());
+ assertEquals(validate, read.getValidate());
+ }
+
@Before
public void setUp() throws IOException, InterruptedException {
FakeDatasetService.setUp();
@@ -203,6 +215,13 @@ public class BigQueryIOReadTest implements Serializable {
}
@Test
+ public void testBuildQueryBasedTypedReadSource() {
+ BigQueryIO.TypedRead read =
+ BigQueryIO.readTableRows().fromQuery("foo_query").withKmsKey("kms_key");
+ checkTypedReadQueryObject(read, "foo_query", "kms_key");
+ }
+
+ @Test
public void testValidateReadSetsDefaultProject() throws Exception {
String tableId = "sometable";
TableReference tableReference =
@@ -620,6 +639,7 @@ public class BigQueryIOReadTest implements Serializable {
TableRowJsonCoder.of(),
BigQueryIO.TableRowParser.INSTANCE,
QueryPriority.BATCH,
+ null,
null);
fakeJobService.expectDryRunQuery(
@@ -688,6 +708,7 @@ public class BigQueryIOReadTest implements Serializable {
TableRowJsonCoder.of(),
BigQueryIO.TableRowParser.INSTANCE,
QueryPriority.BATCH,
+ null,
null);
options.setTempLocation(testFolder.getRoot().getAbsolutePath());
@@ -778,6 +799,7 @@ public class BigQueryIOReadTest implements Serializable {
TableRowJsonCoder.of(),
BigQueryIO.TableRowParser.INSTANCE,
QueryPriority.BATCH,
+ null,
null);
options.setTempLocation(testFolder.getRoot().getAbsolutePath());
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOWriteTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOWriteTest.java
index 47297e2..df18467 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOWriteTest.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOWriteTest.java
@@ -1235,7 +1235,8 @@ public class BigQueryIOWriteTest implements Serializable {
new IdentityDynamicTables(),
null,
4,
- false);
+ false,
+ null);
PCollection<KV<TableDestination, String>> writeTablesOutput =
writeTablesInput.apply(writeTables);
@@ -1317,7 +1318,8 @@ public class BigQueryIOWriteTest implements Serializable {
jobIdTokenView,
BigQueryIO.Write.WriteDisposition.WRITE_EMPTY,
BigQueryIO.Write.CreateDisposition.CREATE_IF_NEEDED,
- 3);
+ 3,
+ "kms_key");
DoFnTester<Iterable<KV<TableDestination, String>>, Void> tester = DoFnTester.of(writeRename);
tester.setSideInput(jobIdTokenView, GlobalWindow.INSTANCE, jobIdToken);
@@ -1329,6 +1331,7 @@ public class BigQueryIOWriteTest implements Serializable {
TableReference tableReference = tableDestination.getTableReference();
Table table = checkNotNull(fakeDatasetService.getTable(tableReference));
assertEquals(tableReference.getTableId() + "_desc", tableDestination.getTableDescription());
+ assertEquals("kms_key", table.getEncryptionConfiguration().getKmsKeyName());
Collection<TableRow> expectedRows = expectedRowsPerTable.get(tableDestination);
assertThat(
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryKmsKeyIT.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryKmsKeyIT.java
new file mode 100644
index 0000000..38ac61a
--- /dev/null
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryKmsKeyIT.java
@@ -0,0 +1,119 @@
+/*
+ * 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.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+
+import com.google.api.services.bigquery.model.Table;
+import com.google.api.services.bigquery.model.TableFieldSchema;
+import com.google.api.services.bigquery.model.TableSchema;
+import java.security.SecureRandom;
+import org.apache.beam.sdk.Pipeline;
+import org.apache.beam.sdk.extensions.gcp.options.GcpOptions;
+import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.Write.Method;
+import org.apache.beam.sdk.io.gcp.testing.BigqueryClient;
+import org.apache.beam.sdk.testing.TestPipeline;
+import org.apache.beam.sdk.testing.TestPipelineOptions;
+import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ImmutableList;
+import org.junit.AfterClass;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Integration tests for BigQuery operations that can use KMS keys, for use with DirectRunner.
+ *
+ * <p>Verification of KMS key usage is done on outputs, but not on any temporary files or tables
+ * used.
+ */
+@RunWith(JUnit4.class)
+public class BigQueryKmsKeyIT {
+
+ private static final Logger LOG = LoggerFactory.getLogger(BigQueryKmsKeyIT.class);
+
+ private static final BigqueryClient BQ_CLIENT = new BigqueryClient("BigQueryKmsKeyIT");
+ private static final String BIG_QUERY_DATASET_ID =
+ "bq_query_to_table_" + System.currentTimeMillis() + "_" + (new SecureRandom().nextInt(32));
+ private static final TableSchema OUTPUT_SCHEMA =
+ new TableSchema()
+ .setFields(ImmutableList.of(new TableFieldSchema().setName("fruit").setType("STRING")));
+
+ private static TestPipelineOptions options;
+ private static String project;
+ private static String kmsKey =
+ "projects/apache-beam-testing/locations/global/keyRings/beam-it/cryptoKeys/test";
+
+ @BeforeClass
+ public static void setupTestEnvironment() throws Exception {
+ options = TestPipeline.testingPipelineOptions().as(TestPipelineOptions.class);
+ project = options.as(GcpOptions.class).getProject();
+ BQ_CLIENT.createNewDataset(project, BIG_QUERY_DATASET_ID);
+ }
+
+ @AfterClass
+ public static void cleanup() {
+ LOG.info("Start to clean up tables and datasets.");
+ BQ_CLIENT.deleteDataset(project, BIG_QUERY_DATASET_ID);
+ }
+
+ /**
+ * Tests query job and table creation with KMS key settings.
+ *
+ * <p>Verifies table creation with KMS key.
+ */
+ private void testQueryAndWrite(Method method) throws Exception {
+ String outputTableId = "testQueryAndWrite_" + method.name();
+ String outputTableSpec = project + ":" + BIG_QUERY_DATASET_ID + "." + outputTableId;
+
+ options.setTempLocation(options.getTempRoot() + "/bq_it_temp");
+ Pipeline p = Pipeline.create(options);
+ // Reading triggers BQ query and extract jobs. Writing triggers either a load job or performs a
+ // streaming insert (depending on method).
+ p.apply(
+ BigQueryIO.readTableRows()
+ .fromQuery("SELECT * FROM (SELECT \"foo\" as fruit)")
+ .withKmsKey(kmsKey))
+ .apply(
+ BigQueryIO.writeTableRows()
+ .to(outputTableSpec)
+ .withSchema(OUTPUT_SCHEMA)
+ .withMethod(method)
+ .withKmsKey(kmsKey));
+ p.run().waitUntilFinish();
+
+ Table table = BQ_CLIENT.getTableResource(project, BIG_QUERY_DATASET_ID, outputTableId);
+ assertNotNull(String.format("table not found: %s", outputTableId), table);
+ assertNotNull(
+ "output table has no EncryptionConfiguration", table.getEncryptionConfiguration());
+ assertEquals(table.getEncryptionConfiguration().getKmsKeyName(), kmsKey);
+ }
+
+ @Test
+ public void testWithFileLoads() throws Exception {
+ testQueryAndWrite(Method.FILE_LOADS);
+ }
+
+ @Test
+ public void testWithStreamingInserts() throws Exception {
+ testQueryAndWrite(Method.STREAMING_INSERTS);
+ }
+}
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/FakeJobService.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/FakeJobService.java
index 26bea47..d03033e 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/FakeJobService.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/FakeJobService.java
@@ -407,7 +407,8 @@ public class FakeJobService implements JobService, Serializable {
new Table()
.setTableReference(destination)
.setSchema(schema)
- .setTimePartitioning(partitioning));
+ .setTimePartitioning(partitioning)
+ .setEncryptionConfiguration(copy.getDestinationEncryptionConfiguration()));
datasetService.insertAll(destination, allRows, null);
return new JobStatus().setState("DONE");
}