Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: add support for ExternalDatasetReference #2871

Merged
merged 2 commits into from Sep 13, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
5 changes: 5 additions & 0 deletions google-cloud-bigquery/clirr-ignored-differences.xml
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>
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
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
@@ -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();
}
}
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());
}
}
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());
}
}