feat(bigquery): support IAM conditions in datasets in Java client. (#3602)

whuffman36 · web-flow · commit 6696a9c7d429 · 2025-01-02T13:19:49.000-08:00
* feat(bigquery): support IAM conditions in datasets in Java client.

* Fix formatting

* Account for possible null condition field in Acl.

* Add toString() method to Acl.Expr object.

Use service account in integration test instead of hardcoded
personal account.

Change Database API calls to only have one branch, toggling only
the access policy version in a conditional.

* Change Acl.User to be default google credentials in IT test

* fix formatting

* Add Acl.Expr builder. Fix review nits.

* fix formatting
diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/Acl.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/Acl.java
@@ -21,6 +21,7 @@
 import com.google.api.core.ApiFunction;
 import com.google.api.services.bigquery.model.Dataset.Access;
 import com.google.api.services.bigquery.model.DatasetAccessEntry;
+import com.google.api.services.bigquery.model.Expr;
 import com.google.cloud.StringEnumType;
 import com.google.cloud.StringEnumValue;
 import java.io.Serializable;
@@ -41,6 +42,7 @@ public final class Acl implements Serializable {
 
   private final Entity entity;
   private final Role role;
+  private final Expr condition;
 
   /**
    * Dataset roles supported by BigQuery.
@@ -568,9 +570,147 @@ Access toPb() {
     }
   }
 
+  /** Expr represents the conditional information related to dataset access policies. */
+  public static final class Expr implements Serializable {
+    // Textual representation of an expression in Common Expression Language syntax.
+    private final String expression;
+    /**
+     * Optional. Title for the expression, i.e. a short string describing its purpose. This can be
+     * used e.g. in UIs which allow to enter the expression.
+     */
+    private final String title;
+    /**
+     * Optional. Description of the expression. This is a longer text which describes the
+     * expression, e.g. when hovered over it in a UI.
+     */
+    private final String description;
+    /**
+     * Optional. String indicating the location of the expression for error reporting, e.g. a file
+     * name and a position in the file.
+     */
+    private final String location;
+
+    private static final long serialVersionUID = 7358264726377291156L;
+
+    static final class Builder {
+      private String expression;
+      private String title;
+      private String description;
+      private String location;
+
+      Builder() {}
+
+      Builder(Expr expr) {
+        this.expression = expr.expression;
+        this.title = expr.title;
+        this.description = expr.description;
+        this.location = expr.location;
+      }
+
+      Builder(com.google.api.services.bigquery.model.Expr bqExpr) {
+        this.expression = bqExpr.getExpression();
+        if (bqExpr.getTitle() != null) {
+          this.title = bqExpr.getTitle();
+        }
+        if (bqExpr.getDescription() != null) {
+          this.description = bqExpr.getDescription();
+        }
+        if (bqExpr.getLocation() != null) {
+          this.location = bqExpr.getLocation();
+        }
+      }
+
+      public Builder setExpression(String expression) {
+        this.expression = expression;
+        return this;
+      }
+
+      public Builder setTitle(String title) {
+        this.title = title;
+        return this;
+      }
+
+      public Builder setDescription(String description) {
+        this.description = description;
+        return this;
+      }
+
+      public Builder setLocation(String location) {
+        this.location = location;
+        return this;
+      }
+
+      public Expr build() {
+        return new Expr(this);
+      }
+    }
+
+    public Expr(Builder builder) {
+      this.expression = builder.expression;
+      this.title = builder.title;
+      this.description = builder.description;
+      this.location = builder.location;
+    }
+
+    public Expr(String expression, String title, String description, String location) {
+      this.expression = expression;
+      this.title = title;
+      this.description = description;
+      this.location = location;
+    }
+
+    com.google.api.services.bigquery.model.Expr toPb() {
+      com.google.api.services.bigquery.model.Expr bqExpr =
+          new com.google.api.services.bigquery.model.Expr();
+      bqExpr.setExpression(this.expression);
+      bqExpr.setTitle(this.title);
+      bqExpr.setDescription(this.description);
+      bqExpr.setLocation(this.location);
+      return bqExpr;
+    }
+
+    static Expr fromPb(com.google.api.services.bigquery.model.Expr bqExpr) {
+      return new Builder(bqExpr).build();
+    }
+
+    public Builder toBuilder() {
+      return new Builder(this);
+    }
+
+    @Override
+    public int hashCode() {
+      return Objects.hash(expression, title, description, location);
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+      if (this == obj) {
+        return true;
+      }
+      if (obj == null || getClass() != obj.getClass()) {
+        return false;
+      }
+      final Expr other = (Expr) obj;
+      return Objects.equals(this.expression, other.expression)
+          && Objects.equals(this.title, other.title)
+          && Objects.equals(this.description, other.description)
+          && Objects.equals(this.location, other.location);
+    }
+
+    @Override
+    public String toString() {
+      return toPb().toString();
+    }
+  }
+
   private Acl(Entity entity, Role role) {
+    this(entity, role, null);
+  }
+
+  private Acl(Entity entity, Role role, Expr condition) {
     this.entity = checkNotNull(entity);
     this.role = role;
+    this.condition = condition;
   }
 
   /** @return Returns the entity for this ACL. */
@@ -582,6 +722,10 @@ public Entity getEntity() {
   public Role getRole() {
     return role;
   }
+  /** @return Returns the condition specified by this ACL. */
+  public Expr getCondition() {
+    return condition;
+  }
 
   /**
    * @return Returns an Acl object.
@@ -592,6 +736,10 @@ public static Acl of(Entity entity, Role role) {
     return new Acl(entity, role);
   }
 
+  public static Acl of(Entity entity, Role role, Expr condition) {
+    return new Acl(entity, role, condition);
+  }
+
   /**
    * @param datasetAclEntity
    * @return Returns an Acl object for a datasetAclEntity.
@@ -618,7 +766,7 @@ public static Acl of(Routine routine) {
 
   @Override
   public int hashCode() {
-    return Objects.hash(entity, role);
+    return Objects.hash(entity, role, condition);
   }
 
   @Override
@@ -635,19 +783,26 @@ public boolean equals(Object obj) {
       return false;
     }
     final Acl other = (Acl) obj;
-    return Objects.equals(this.entity, other.entity) && Objects.equals(this.role, other.role);
+    return Objects.equals(this.entity, other.entity)
+        && Objects.equals(this.role, other.role)
+        && Objects.equals(this.condition, other.condition);
   }
 
   Access toPb() {
     Access accessPb = entity.toPb();
     if (role != null) {
       accessPb.setRole(role.name());
     }
+    if (condition != null) {
+      accessPb.setCondition(condition.toPb());
+    }
     return accessPb;
   }
 
   static Acl fromPb(Access access) {
     return Acl.of(
-        Entity.fromPb(access), access.getRole() != null ? Role.valueOf(access.getRole()) : null);
+        Entity.fromPb(access),
+        access.getRole() != null ? Role.valueOf(access.getRole()) : null,
+        access.getCondition() != null ? Expr.fromPb(access.getCondition()) : null);
   }
 }
diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQuery.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQuery.java
@@ -289,6 +289,24 @@ public static DatasetOption fields(DatasetField... fields) {
       return new DatasetOption(
           BigQueryRpc.Option.FIELDS, Helper.selector(DatasetField.REQUIRED_FIELDS, fields));
     }
+
+    /**
+     * Returns an option to specify the dataset's access policy version for conditional access. If
+     * this option is not provided the field remains unset and conditional access cannot be used.
+     * Valid values are 0, 1, and 3. Requests specifying an invalid value will be rejected. Requests
+     * for conditional access policy binding in datasets must specify version 3. Datasets with no
+     * conditional role bindings in access policy may specify any valid value or leave the field
+     * unset. This field will be mapped to <a
+     * href="https://cloud.google.com/iam/docs/policies#versions">IAM Policy version</a> and will be
+     * used to fetch the policy from IAM. If unset or if 0 or 1 the value is used for a dataset with
+     * conditional bindings, access entry with condition will have role string appended by
+     * 'withcond' string followed by a hash value. Please refer to <a
+     * href="https://cloud.google.com/iam/docs/troubleshooting-withcond">Troubleshooting
+     * withcond</a> for more details.
+     */
+    public static DatasetOption accessPolicyVersion(Integer accessPolicyVersion) {
+      return new DatasetOption(BigQueryRpc.Option.ACCESS_POLICY_VERSION, accessPolicyVersion);
+    }
   }
 
   /** Class for specifying dataset delete options. */
diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/BigQueryRpc.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/BigQueryRpc.java
@@ -59,7 +59,8 @@ enum Option {
     REQUESTED_POLICY_VERSION("requestedPolicyVersion"),
     TABLE_METADATA_VIEW("view"),
     RETRY_OPTIONS("retryOptions"),
-    BIGQUERY_RETRY_CONFIG("bigQueryRetryConfig");
+    BIGQUERY_RETRY_CONFIG("bigQueryRetryConfig"),
+    ACCESS_POLICY_VERSION("accessPolicyVersion");
 
     private final String value;
 
diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpBigQueryRpc.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpBigQueryRpc.java
@@ -130,12 +130,19 @@ private void validateRPC() throws BigQueryException, IOException {
   public Dataset getDataset(String projectId, String datasetId, Map<Option, ?> options) {
     try {
       validateRPC();
-      return bigquery
-          .datasets()
-          .get(projectId, datasetId)
-          .setFields(Option.FIELDS.getString(options))
-          .setPrettyPrint(false)
-          .execute();
+
+      Bigquery.Datasets.Get bqGetRequest =
+          bigquery
+              .datasets()
+              .get(projectId, datasetId)
+              .setFields(Option.FIELDS.getString(options))
+              .setPrettyPrint(false);
+      for (Map.Entry<Option, ?> entry : options.entrySet()) {
+        if (entry.getKey() == Option.ACCESS_POLICY_VERSION && entry.getValue() != null) {
+          bqGetRequest.setAccessPolicyVersion((Integer) entry.getValue());
+        }
+      }
+      return bqGetRequest.execute();
     } catch (IOException ex) {
       BigQueryException serviceException = translate(ex);
       if (serviceException.getCode() == HTTP_NOT_FOUND) {
@@ -174,12 +181,18 @@ public Tuple<String, Iterable<Dataset>> listDatasets(String projectId, Map<Optio
   public Dataset create(Dataset dataset, Map<Option, ?> options) {
     try {
       validateRPC();
-      return bigquery
-          .datasets()
-          .insert(dataset.getDatasetReference().getProjectId(), dataset)
-          .setPrettyPrint(false)
-          .setFields(Option.FIELDS.getString(options))
-          .execute();
+      Bigquery.Datasets.Insert bqCreateRequest =
+          bigquery
+              .datasets()
+              .insert(dataset.getDatasetReference().getProjectId(), dataset)
+              .setPrettyPrint(false)
+              .setFields(Option.FIELDS.getString(options));
+      for (Map.Entry<Option, ?> entry : options.entrySet()) {
+        if (entry.getKey() == Option.ACCESS_POLICY_VERSION && entry.getValue() != null) {
+          bqCreateRequest.setAccessPolicyVersion((Integer) entry.getValue());
+        }
+      }
+      return bqCreateRequest.execute();
     } catch (IOException ex) {
       throw translate(ex);
     }
@@ -277,12 +290,18 @@ public Dataset patch(Dataset dataset, Map<Option, ?> options) {
     try {
       validateRPC();
       DatasetReference reference = dataset.getDatasetReference();
-      return bigquery
-          .datasets()
-          .patch(reference.getProjectId(), reference.getDatasetId(), dataset)
-          .setPrettyPrint(false)
-          .setFields(Option.FIELDS.getString(options))
-          .execute();
+      Bigquery.Datasets.Patch bqPatchRequest =
+          bigquery
+              .datasets()
+              .patch(reference.getProjectId(), reference.getDatasetId(), dataset)
+              .setPrettyPrint(false)
+              .setFields(Option.FIELDS.getString(options));
+      for (Map.Entry<Option, ?> entry : options.entrySet()) {
+        if (entry.getKey() == Option.ACCESS_POLICY_VERSION && entry.getValue() != null) {
+          bqPatchRequest.setAccessPolicyVersion((Integer) entry.getValue());
+        }
+      }
+      return bqPatchRequest.execute();
     } catch (IOException ex) {
       throw translate(ex);
     }
diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/AclTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/AclTest.java
@@ -23,6 +23,7 @@
 import com.google.cloud.bigquery.Acl.Domain;
 import com.google.cloud.bigquery.Acl.Entity;
 import com.google.cloud.bigquery.Acl.Entity.Type;
+import com.google.cloud.bigquery.Acl.Expr;
 import com.google.cloud.bigquery.Acl.Group;
 import com.google.cloud.bigquery.Acl.IamMember;
 import com.google.cloud.bigquery.Acl.Role;
@@ -136,4 +137,13 @@ public void testOf() {
     assertEquals(routine, acl.getEntity());
     assertEquals(null, acl.getRole());
   }
+
+  @Test
+  public void testOfWithCondition() {
+    Expr expr = new Expr("expression", "title", "description", "location");
+    Acl acl = Acl.of(Group.ofAllAuthenticatedUsers(), Role.READER, expr);
+    Dataset.Access pb = acl.toPb();
+    assertEquals(acl, Acl.fromPb(pb));
+    assertEquals(acl.getCondition(), expr);
+  }
 }
diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java
@@ -30,6 +30,7 @@
 import com.google.cloud.RetryOption;
 import com.google.cloud.ServiceOptions;
 import com.google.cloud.Tuple;
+import com.google.cloud.bigquery.BigQuery.DatasetOption;
 import com.google.cloud.bigquery.BigQuery.JobOption;
 import com.google.cloud.bigquery.BigQuery.QueryResultsOption;
 import com.google.cloud.bigquery.InsertAllRequest.RowToInsert;
@@ -572,6 +573,20 @@ public void testCreateDatasetWithSelectedFields() {
     verify(bigqueryRpcMock).create(eq(DATASET_INFO_WITH_PROJECT.toPb()), capturedOptions.capture());
   }
 
+  @Test
+  public void testCreateDatasetWithAccessPolicy() {
+    DatasetInfo datasetInfo = DATASET_INFO.setProjectId(OTHER_PROJECT);
+    DatasetOption datasetOption = DatasetOption.accessPolicyVersion(3);
+    when(bigqueryRpcMock.create(datasetInfo.toPb(), optionMap(datasetOption)))
+        .thenReturn(datasetInfo.toPb());
+    BigQueryOptions bigQueryOptions =
+        createBigQueryOptionsForProject(OTHER_PROJECT, rpcFactoryMock);
+    bigquery = bigQueryOptions.getService();
+    Dataset dataset = bigquery.create(datasetInfo, datasetOption);
+    assertEquals(new Dataset(bigquery, new DatasetInfo.BuilderImpl(datasetInfo)), dataset);
+    verify(bigqueryRpcMock).create(datasetInfo.toPb(), optionMap(datasetOption));
+  }
+
   @Test
   public void testGetDataset() {
     when(bigqueryRpcMock.getDataset(PROJECT, DATASET, EMPTY_RPC_OPTIONS))
diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java
