feat: add support for ExternalDatasetReference

google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/Dataset.java

@@ -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);

google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/DatasetInfo.java

@@ -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 {
@@ -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,
@@ -183,6 +191,7 @@ static final class BuilderImpl extends Builder {
     private EncryptionConfiguration defaultEncryptionConfiguration;
     private Long defaultPartitionExpirationMs;
     private String defaultCollation;
+    private ExternalDatasetReference externalDatasetReference;
 
     BuilderImpl() {}
 
@@ -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) {
@@ -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
@@ -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);
@@ -358,6 +378,7 @@ public DatasetInfo build() {
     defaultEncryptionConfiguration = builder.defaultEncryptionConfiguration;
     defaultPartitionExpirationMs = builder.defaultPartitionExpirationMs;
     defaultCollation = builder.defaultCollation;
+    externalDatasetReference = builder.externalDatasetReference;
   }
 
   /** Returns the dataset identity. */
@@ -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);
@@ -510,6 +539,7 @@ public String toString() {
         .add("defaultEncryptionConfiguration", defaultEncryptionConfiguration)
         .add("defaultPartitionExpirationMs", defaultPartitionExpirationMs)
         .add("defaultCollation", defaultCollation)
+        .add("externalDatasetReference", externalDatasetReference)
         .toString();
   }
 
@@ -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;
   }
 

google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ExternalDatasetReference.java

@@ -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();
+  }
+}

google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/DatasetInfoTest.java

@@ -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)
@@ -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() {
@@ -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());
@@ -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
@@ -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());
@@ -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()));
   }
@@ -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());
   }
 }

google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/DatasetTest.java

@@ -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;
 
@@ -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);
@@ -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());
   }
 }