Skip to content

Commit

Permalink
feat: add support for ExternalDatasetReference (#2871)
Browse files Browse the repository at this point in the history 
API: https://cloud.google.com/bigquery/docs/reference/rest/v2/datasets#externaldatasetreference
  • Loading branch information
@obada-ab
obada-ab committed Sep 13, 2023
1 parent 7772c2a commit bbb86fd
Show file tree
Hide file tree
Showing 8 changed files with 264 additions and 1 deletion.
5 changes: 5 additions & 0 deletions google-cloud-bigquery/clirr-ignored-differences.xml
View file
Expand Up @@ -104,4 +104,9 @@
<className>com/google/cloud/bigquery/IndexUnusedReason*</className>
<method>*BaseTable(*)</method>
</difference>
<difference>
<differenceType>7013</differenceType>
<className>com/google/cloud/bigquery/DatasetInfo*</className>
<method>*setExternalDatasetReference(*)</method>
</difference>
</differences>
6 changes: 6 additions & 0 deletions google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/Dataset.java
View file
Expand Up @@ -152,6 +152,12 @@ public Builder setDefaultCollation(String defaultCollation) {
return this;
}

@Override
public Builder setExternalDatasetReference(ExternalDatasetReference externalDatasetReference) {
infoBuilder.setExternalDatasetReference(externalDatasetReference);
return this;
}

@Override
public Dataset build() {
return new Dataset(bigquery, infoBuilder);
Expand Down
33 changes: 33 additions & 0 deletions google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/DatasetInfo.java
View file
Expand Up @@ -73,6 +73,7 @@ public Dataset apply(DatasetInfo datasetInfo) {
private final EncryptionConfiguration defaultEncryptionConfiguration;
private final Long defaultPartitionExpirationMs;
private final String defaultCollation;
private final ExternalDatasetReference externalDatasetReference;

/** A builder for {@code DatasetInfo} objects. */
public abstract static class Builder {
Expand Down Expand Up @@ -127,6 +128,13 @@ public abstract static class Builder {

public abstract Builder setLabels(Map<String, String> labels);

/**
* Optional. Information about the external metadata storage where the dataset is defined.
* Filled out when the dataset type is EXTERNAL
*/
public abstract Builder setExternalDatasetReference(
ExternalDatasetReference externalDatasetReference);

/**
* The default encryption key for all tables in the dataset. Once this property is set, all
* newly-created partitioned tables in the dataset will have encryption key set to this value,
Expand Down Expand Up @@ -183,6 +191,7 @@ static final class BuilderImpl extends Builder {
private EncryptionConfiguration defaultEncryptionConfiguration;
private Long defaultPartitionExpirationMs;
private String defaultCollation;
private ExternalDatasetReference externalDatasetReference;

BuilderImpl() {}

Expand All @@ -202,6 +211,7 @@ static final class BuilderImpl extends Builder {
this.defaultEncryptionConfiguration = datasetInfo.defaultEncryptionConfiguration;
this.defaultPartitionExpirationMs = datasetInfo.defaultPartitionExpirationMs;
this.defaultCollation = datasetInfo.defaultCollation;
this.externalDatasetReference = datasetInfo.externalDatasetReference;
}

BuilderImpl(com.google.api.services.bigquery.model.Dataset datasetPb) {
Expand Down Expand Up @@ -236,6 +246,10 @@ public Acl apply(Dataset.Access accessPb) {
}
this.defaultPartitionExpirationMs = datasetPb.getDefaultPartitionExpirationMs();
this.defaultCollation = datasetPb.getDefaultCollation();
if (datasetPb.getExternalDatasetReference() != null) {
this.externalDatasetReference =
ExternalDatasetReference.fromPb(datasetPb.getExternalDatasetReference());
}
}

@Override
Expand Down Expand Up @@ -336,6 +350,12 @@ public Builder setDefaultCollation(String defaultCollation) {
return this;
}

@Override
public Builder setExternalDatasetReference(ExternalDatasetReference externalDatasetReference) {
this.externalDatasetReference = externalDatasetReference;
return this;
}

@Override
public DatasetInfo build() {
return new DatasetInfo(this);
Expand All @@ -358,6 +378,7 @@ public DatasetInfo build() {
defaultEncryptionConfiguration = builder.defaultEncryptionConfiguration;
defaultPartitionExpirationMs = builder.defaultPartitionExpirationMs;
defaultCollation = builder.defaultCollation;
externalDatasetReference = builder.externalDatasetReference;
}

/** Returns the dataset identity. */
Expand Down Expand Up @@ -487,6 +508,14 @@ public String getDefaultCollation() {
return defaultCollation;
}

/**
* Returns information about the external metadata storage where the dataset is defined. Filled
* out when the dataset type is EXTERNAL.
*/
public ExternalDatasetReference getExternalDatasetReference() {
return externalDatasetReference;
}

/** Returns a builder for the dataset object. */
public Builder toBuilder() {
return new BuilderImpl(this);
Expand All @@ -510,6 +539,7 @@ public String toString() {
.add("defaultEncryptionConfiguration", defaultEncryptionConfiguration)
.add("defaultPartitionExpirationMs", defaultPartitionExpirationMs)
.add("defaultCollation", defaultCollation)
.add("externalDatasetReference", externalDatasetReference)
.toString();
}

Expand Down Expand Up @@ -588,6 +618,9 @@ public Dataset.Access apply(Acl acl) {
if (defaultCollation != null) {
datasetPb.setDefaultCollation(defaultCollation);
}
if (externalDatasetReference != null) {
datasetPb.setExternalDatasetReference(externalDatasetReference.toPb());
}
return datasetPb;
}

Expand Down
79 changes: 79 additions & 0 deletions google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ExternalDatasetReference.java
View file
@@ -0,0 +1,79 @@
/*
* Copyright 2023 Google LLC
*
* 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.bigquery;

import com.google.auto.value.AutoValue;
import com.google.common.annotations.VisibleForTesting;
import java.io.Serializable;
import javax.annotation.Nullable;

/** Configures the access a dataset defined in an external metadata storage. */
@AutoValue
public abstract class ExternalDatasetReference implements Serializable {

public static ExternalDatasetReference.Builder newBuilder() {
return new AutoValue_ExternalDatasetReference.Builder();
}

static ExternalDatasetReference fromPb(
com.google.api.services.bigquery.model.ExternalDatasetReference externalDatasetReference) {
ExternalDatasetReference.Builder builder = newBuilder();

if (externalDatasetReference.getConnection() != null) {
builder.setConnection(externalDatasetReference.getConnection());
}
if (externalDatasetReference.getExternalSource() != null) {
builder.setExternalSource(externalDatasetReference.getExternalSource());
}

return builder.build();
}

public com.google.api.services.bigquery.model.ExternalDatasetReference toPb() {
com.google.api.services.bigquery.model.ExternalDatasetReference externalDatasetReference =
new com.google.api.services.bigquery.model.ExternalDatasetReference();

externalDatasetReference.setConnection(getConnection());
externalDatasetReference.setExternalSource(getExternalSource());
return externalDatasetReference;
}

@Nullable
public abstract String getConnection();

@Nullable
public abstract String getExternalSource();

/** Returns a builder for an ExternalDatasetReference. */
@VisibleForTesting
public abstract ExternalDatasetReference.Builder toBuilder();

@AutoValue.Builder
public abstract static class Builder {
/**
* The connection id that is used to access the external_source. Format:
* projects/{project_id}/locations/{location_id}/connections/{connection_id} *
*/
public abstract ExternalDatasetReference.Builder setConnection(String connection);

/** External source that backs this dataset * */
public abstract ExternalDatasetReference.Builder setExternalSource(String externalSource);

/** Creates a {@code ExternalDatasetReference} object. */
public abstract ExternalDatasetReference build();
}
}
39 changes: 39 additions & 0 deletions google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/DatasetInfoTest.java
View file
Expand Up @@ -58,6 +58,12 @@ public class DatasetInfoTest {
private static final DatasetId DATASET_ID_COMPLETE = DatasetId.of("project", "dataset");
private static final EncryptionConfiguration DATASET_ENCRYPTION_CONFIGURATION =
EncryptionConfiguration.newBuilder().setKmsKeyName("KMS_KEY_1").build();

private static final ExternalDatasetReference EXTERNAL_DATASET_REFERENCE =
ExternalDatasetReference.newBuilder()
.setExternalSource("source")
.setConnection("connection")
.build();
private static final DatasetInfo DATASET_INFO =
DatasetInfo.newBuilder(DATASET_ID)
.setAcl(ACCESS_RULES)
Expand All @@ -82,6 +88,8 @@ public class DatasetInfoTest {
.build();
private static final DatasetInfo DATASET_INFO_COMPLETE_WITH_IAM_MEMBER =
DATASET_INFO.toBuilder().setAcl(ACCESS_RULES_IAM_MEMBER).build();
private static final DatasetInfo DATASET_INFO_COMPLETE_WITH_EXTERNAL_DATASET_REFERENCE =
DATASET_INFO.toBuilder().setExternalDatasetReference(EXTERNAL_DATASET_REFERENCE).build();

@Test
public void testToBuilder() {
Expand All @@ -108,6 +116,28 @@ public void testToBuilderIncomplete() {
assertEquals(datasetInfo, datasetInfo.toBuilder().build());
}

@Test
public void testToBuilderWithExternalDatasetReference() {
compareDatasets(
DATASET_INFO_COMPLETE_WITH_EXTERNAL_DATASET_REFERENCE,
DATASET_INFO_COMPLETE_WITH_EXTERNAL_DATASET_REFERENCE.toBuilder().build());

ExternalDatasetReference externalDatasetReference =
ExternalDatasetReference.newBuilder()
.setExternalSource("source2")
.setConnection("connection2")
.build();
DatasetInfo datasetInfo =
DATASET_INFO_COMPLETE_WITH_EXTERNAL_DATASET_REFERENCE
.toBuilder()
.setExternalDatasetReference(externalDatasetReference)
.build();
assertEquals(externalDatasetReference, datasetInfo.getExternalDatasetReference());
datasetInfo =
datasetInfo.toBuilder().setExternalDatasetReference(EXTERNAL_DATASET_REFERENCE).build();
compareDatasets(DATASET_INFO_COMPLETE_WITH_EXTERNAL_DATASET_REFERENCE, datasetInfo);
}

@Test
public void testBuilder() {
assertNull(DATASET_INFO.getDatasetId().getProject());
Expand Down Expand Up @@ -137,6 +167,9 @@ public void testBuilder() {
assertEquals(LOCATION, DATASET_INFO_COMPLETE.getLocation());
assertEquals(SELF_LINK, DATASET_INFO_COMPLETE.getSelfLink());
assertEquals(LABELS, DATASET_INFO_COMPLETE.getLabels());
assertEquals(
EXTERNAL_DATASET_REFERENCE,
DATASET_INFO_COMPLETE_WITH_EXTERNAL_DATASET_REFERENCE.getExternalDatasetReference());
}

@Test
Expand All @@ -156,6 +189,7 @@ public void testOf() {
assertNull(datasetInfo.getDefaultEncryptionConfiguration());
assertNull(datasetInfo.getDefaultPartitionExpirationMs());
assertTrue(datasetInfo.getLabels().isEmpty());
assertNull(datasetInfo.getExternalDatasetReference());

datasetInfo = DatasetInfo.of(DATASET_ID);
assertEquals(DATASET_ID, datasetInfo.getDatasetId());
Expand All @@ -172,11 +206,15 @@ public void testOf() {
assertNull(datasetInfo.getDefaultEncryptionConfiguration());
assertNull(datasetInfo.getDefaultPartitionExpirationMs());
assertTrue(datasetInfo.getLabels().isEmpty());
assertNull(datasetInfo.getExternalDatasetReference());
}

@Test
public void testToPbAndFromPb() {
compareDatasets(DATASET_INFO_COMPLETE, DatasetInfo.fromPb(DATASET_INFO_COMPLETE.toPb()));
compareDatasets(
DATASET_INFO_COMPLETE_WITH_EXTERNAL_DATASET_REFERENCE,
DatasetInfo.fromPb(DATASET_INFO_COMPLETE_WITH_EXTERNAL_DATASET_REFERENCE.toPb()));
DatasetInfo datasetInfo = DatasetInfo.newBuilder("project", "dataset").build();
compareDatasets(datasetInfo, DatasetInfo.fromPb(datasetInfo.toPb()));
}
Expand Down Expand Up @@ -204,5 +242,6 @@ private void compareDatasets(DatasetInfo expected, DatasetInfo value) {
expected.getDefaultEncryptionConfiguration(), value.getDefaultEncryptionConfiguration());
assertEquals(
expected.getDefaultPartitionExpirationMs(), value.getDefaultPartitionExpirationMs());
assertEquals(expected.getExternalDatasetReference(), value.getExternalDatasetReference());
}
}
31 changes: 31 additions & 0 deletions google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/DatasetTest.java
View file
Expand Up @@ -83,6 +83,11 @@ public class DatasetTest {
TableInfo.newBuilder(
TableId.of(NEW_PROJECT_ID, "dataset", "table3"), EXTERNAL_TABLE_DEFINITION)
.build();
private static final ExternalDatasetReference EXTERNAL_DATASET_REFERENCE =
ExternalDatasetReference.newBuilder()
.setExternalSource("source")
.setConnection("connection")
.build();

@Rule public MockitoRule rule;

Expand Down Expand Up @@ -319,6 +324,31 @@ public void testToAndFromPb() {
compareDataset(expectedDataset, Dataset.fromPb(bigquery, expectedDataset.toPb()));
}

@Test
public void testExternalDatasetReference() {
Dataset datasetWithExternalDatasetReference =
new Dataset.Builder(bigquery, DATASET_ID)
.setAcl(ACCESS_RULES)
.setCreationTime(CREATION_TIME)
.setDefaultTableLifetime(DEFAULT_TABLE_EXPIRATION)
.setDescription(DESCRIPTION)
.setEtag(ETAG)
.setFriendlyName(FRIENDLY_NAME)
.setGeneratedId(GENERATED_ID)
.setLastModified(LAST_MODIFIED)
.setLocation(LOCATION)
.setSelfLink(SELF_LINK)
.setLabels(LABELS)
.setExternalDatasetReference(EXTERNAL_DATASET_REFERENCE)
.build();
assertEquals(
EXTERNAL_DATASET_REFERENCE,
datasetWithExternalDatasetReference.getExternalDatasetReference());
compareDataset(
datasetWithExternalDatasetReference,
datasetWithExternalDatasetReference.toBuilder().build());
}

private void compareDataset(Dataset expected, Dataset value) {
assertEquals(expected, value);
compareDatasetInfo(expected, value);
Expand All @@ -338,5 +368,6 @@ private void compareDatasetInfo(DatasetInfo expected, DatasetInfo value) {
assertEquals(expected.getCreationTime(), value.getCreationTime());
assertEquals(expected.getDefaultTableLifetime(), value.getDefaultTableLifetime());
assertEquals(expected.getLastModified(), value.getLastModified());
assertEquals(expected.getExternalDatasetReference(), value.getExternalDatasetReference());
}
}

0 comments on commit bbb86fd

Please sign in to comment.