feat(bigquery): support IAM conditions in datasets (#11123)

This PR adds support for IAM conditions via the existing dataset access mechanism.  To do so, the following changes are necessary:

* add the `Expr` type for expressing conditions, and wire it into the existing DatasetAccessEntry.
* add an option pattern to the Dataset-related RPC methods
* Add a new WithAccessPolicyVersion option for setting access policies

To expose the new functionality, this PR adds CreateWithOptions, UpdateWithOptions, MetadataWithOptions methods on Dataset that accept the new option.

Author: shollyman

bigquery/dataset.go

@@ -190,12 +190,28 @@ func (d *Dataset) Identifier(f IdentifierFormat) (string, error) {
 	}
 }
 
-// Create creates a dataset in the BigQuery service. An error will be returned if the
-// dataset already exists. Pass in a DatasetMetadata value to configure the dataset.
+// Create creates a dataset in the BigQuery service.
+//
+// An error will be returned if the dataset already exists.
+// Pass in a DatasetMetadata value to configure the dataset.
 func (d *Dataset) Create(ctx context.Context, md *DatasetMetadata) (err error) {
+	return d.CreateWithOptions(ctx, md)
+}
+
+// CreateWithOptions creates a dataset in the BigQuery service, and
+// provides additional options to control the behavior of the call.
+//
+// An error will be returned if the dataset already exists.
+// Pass in a DatasetMetadata value to configure the dataset.
+func (d *Dataset) CreateWithOptions(ctx context.Context, md *DatasetMetadata, opts ...DatasetOption) (err error) {
 	ctx = trace.StartSpan(ctx, "cloud.google.com/go/bigquery.Dataset.Create")
 	defer func() { trace.EndSpan(ctx, err) }()
 
+	cOpt := &dsCallOption{}
+	for _, o := range opts {
+		o(cOpt)
+	}
+
 	ds, err := md.toBQ()
 	if err != nil {
 		return err
@@ -207,6 +223,9 @@ func (d *Dataset) Create(ctx context.Context, md *DatasetMetadata) (err error) {
 	}
 	call := d.c.bqs.Datasets.Insert(d.ProjectID, ds).Context(ctx)
 	setClientHeader(call.Header())
+	if cOpt.accessPolicyVersion != nil {
+		call.AccessPolicyVersion(int64(optional.ToInt(cOpt.accessPolicyVersion)))
+	}
 	_, err = call.Do()
 	return err
 }
@@ -289,11 +308,25 @@ func (d *Dataset) deleteInternal(ctx context.Context, deleteContents bool) (err
 
 // Metadata fetches the metadata for the dataset.
 func (d *Dataset) Metadata(ctx context.Context) (md *DatasetMetadata, err error) {
+	return d.MetadataWithOptions(ctx)
+}
+
+// MetadataWithOptions fetches metadata for the dataset, and provides additional options for
+// controlling the request.
+func (d *Dataset) MetadataWithOptions(ctx context.Context, opts ...DatasetOption) (md *DatasetMetadata, err error) {
 	ctx = trace.StartSpan(ctx, "cloud.google.com/go/bigquery.Dataset.Metadata")
 	defer func() { trace.EndSpan(ctx, err) }()
 
+	cOpt := &dsCallOption{}
+	for _, o := range opts {
+		o(cOpt)
+	}
+
 	call := d.c.bqs.Datasets.Get(d.ProjectID, d.DatasetID).Context(ctx)
 	setClientHeader(call.Header())
+	if cOpt.accessPolicyVersion != nil {
+		call.AccessPolicyVersion(int64(optional.ToInt(cOpt.accessPolicyVersion)))
+	}
 	var ds *bq.Dataset
 	if err := runWithRetry(ctx, func() (err error) {
 		sCtx := trace.StartSpan(ctx, "bigquery.datasets.get")
@@ -306,6 +339,36 @@ func (d *Dataset) Metadata(ctx context.Context) (md *DatasetMetadata, err error)
 	return bqToDatasetMetadata(ds, d.c)
 }
 
+// dsCallOption provides a general option holder for dataset RPCs
+type dsCallOption struct {
+	accessPolicyVersion optional.Int
+}
+
+// DatasetOption provides an option type for customizing requests against the Dataset
+// service.
+type DatasetOption func(*dsCallOption)
+
+// WithAccessPolicyVersion is an option that enabled setting of the Access Policy Version for a request
+// where appropriate.  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.
+//
+// Dataset with no conditional role bindings in access policy may specify any valid value
+// or leave the field unset.
+//
+// This field will be mapped to [IAM Policy version] (https://cloud.google.com/iam/docs/policies#versions)
+// and will be used to fetch policy from IAM. If unset or if 0 or 1 value is used for
+// dataset with conditional bindings, access entry with condition will have role string
+// appended by 'withcond' string followed by a hash value.
+//
+// Please refer https://cloud.google.com/iam/docs/troubleshooting-withcond for more details.
+func WithAccessPolicyVersion(apv int) DatasetOption {
+	return func(o *dsCallOption) {
+		o.accessPolicyVersion = apv
+	}
+}
+
 func bqToDatasetMetadata(d *bq.Dataset, c *Client) (*DatasetMetadata, error) {
 	dm := &DatasetMetadata{
 		CreationTime:               unixMillisToTime(d.CreationTime),
@@ -345,18 +408,36 @@ func bqToDatasetMetadata(d *bq.Dataset, c *Client) (*DatasetMetadata, error) {
 // set the etag argument to the DatasetMetadata.ETag field from the read.
 // Pass the empty string for etag for a "blind write" that will always succeed.
 func (d *Dataset) Update(ctx context.Context, dm DatasetMetadataToUpdate, etag string) (md *DatasetMetadata, err error) {
+	return d.UpdateWithOptions(ctx, dm, etag)
+}
+
+// UpdateWithOptions modifies specific Dataset metadata fields and
+// provides an interface for specifying additional options to the request.
+//
+// To perform a read-modify-write that protects against intervening reads,
+// set the etag argument to the DatasetMetadata.ETag field from the read.
+// Pass the empty string for etag for a "blind write" that will always succeed.
+func (d *Dataset) UpdateWithOptions(ctx context.Context, dm DatasetMetadataToUpdate, etag string, opts ...DatasetOption) (md *DatasetMetadata, err error) {
 	ctx = trace.StartSpan(ctx, "cloud.google.com/go/bigquery.Dataset.Update")
 	defer func() { trace.EndSpan(ctx, err) }()
 
+	cOpt := &dsCallOption{}
+	for _, o := range opts {
+		o(cOpt)
+	}
 	ds, err := dm.toBQ()
 	if err != nil {
 		return nil, err
 	}
+
 	call := d.c.bqs.Datasets.Patch(d.ProjectID, d.DatasetID, ds).Context(ctx)
 	setClientHeader(call.Header())
 	if etag != "" {
 		call.Header().Set("If-Match", etag)
 	}
+	if cOpt.accessPolicyVersion != nil {
+		call.AccessPolicyVersion(int64(optional.ToInt(cOpt.accessPolicyVersion)))
+	}
 	var ds2 *bq.Dataset
 	if err := runWithRetry(ctx, func() (err error) {
 		sCtx := trace.StartSpan(ctx, "bigquery.datasets.patch")
@@ -811,6 +892,50 @@ type AccessEntry struct {
 	View       *Table              // The view granted access (EntityType must be ViewEntity)
 	Routine    *Routine            // The routine granted access (only UDF currently supported)
 	Dataset    *DatasetAccessEntry // The resources within a dataset granted access.
+	Condition  *Expr               // Condition for the access binding.
+}
+
+// Expr represents the conditional information related to dataset access policies.
+type Expr struct {
+	// Textual representation of an expression in Common Expression Language syntax.
+	Expression string
+
+	// 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.
+	Title string
+
+	// Optional. Description of the expression. This is a longer text which
+	// describes the expression, e.g. when hovered over it in a UI.
+	Description string
+
+	// Optional. String indicating the location of the expression for error
+	// reporting, e.g. a file name and a position in the file.
+	Location string
+}
+
+func (ex *Expr) toBQ() *bq.Expr {
+	if ex == nil {
+		return nil
+	}
+	return &bq.Expr{
+		Expression:  ex.Expression,
+		Title:       ex.Title,
+		Description: ex.Description,
+		Location:    ex.Location,
+	}
+}
+
+func bqToExpr(bq *bq.Expr) *Expr {
+	if bq == nil {
+		return nil
+	}
+	return &Expr{
+		Expression:  bq.Expression,
+		Title:       bq.Title,
+		Description: bq.Description,
+		Location:    bq.Location,
+	}
 }
 
 // AccessRole is the level of access to grant to a dataset.
@@ -857,7 +982,10 @@ const (
 )
 
 func (e *AccessEntry) toBQ() (*bq.DatasetAccess, error) {
-	q := &bq.DatasetAccess{Role: string(e.Role)}
+	q := &bq.DatasetAccess{
+		Role:      string(e.Role),
+		Condition: e.Condition.toBQ(),
+	}
 	switch e.EntityType {
 	case DomainEntity:
 		q.Domain = e.Entity
@@ -911,6 +1039,9 @@ func bqToAccessEntry(q *bq.DatasetAccess, c *Client) (*AccessEntry, error) {
 	default:
 		return nil, errors.New("bigquery: invalid access value")
 	}
+	if q.Condition != nil {
+		e.Condition = bqToExpr(q.Condition)
+	}
 	return e, nil
 }
 

bigquery/dataset_integration_test.go

@@ -22,6 +22,7 @@ import (
 	"time"
 
 	"cloud.google.com/go/internal/testutil"
+	"github.com/google/go-cmp/cmp"
 	"github.com/google/go-cmp/cmp/cmpopts"
 )
 
@@ -435,6 +436,102 @@ func TestIntegration_DatasetUpdateAccess(t *testing.T) {
 	}
 }
 
+// This test validates behaviors related to IAM conditions in
+// dataset access control.
+func TestIntegration_DatasetConditions(t *testing.T) {
+	if client == nil {
+		t.Skip("Integration tests skipped")
+	}
+	ctx := context.Background()
+	// Use our test dataset for a base access policy.
+	md, err := dataset.Metadata(ctx)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	wantEntry := &AccessEntry{
+		Role:       ReaderRole,
+		Entity:     "[email protected]",
+		EntityType: UserEmailEntity,
+		Condition: &Expr{
+			Expression:  "request.time < timestamp('2030-01-01T00:00:00Z')",
+			Description: "requests before the year 2030",
+			Title:       "test condition",
+		},
+	}
+	origAccess := append(md.Access, wantEntry)
+
+	ds := client.Dataset(datasetIDs.New())
+	wantMeta := &DatasetMetadata{
+		Access:      origAccess,
+		Description: "test dataset",
+	}
+
+	// First, attempt to create the dataset without specifying a policy access version.
+	err = ds.Create(ctx, wantMeta)
+	if err == nil {
+		t.Fatalf("expected Create failure, but succeeded")
+	}
+
+	err = ds.CreateWithOptions(ctx, wantMeta, WithAccessPolicyVersion(3))
+	if err != nil {
+		t.Fatalf("expected Create to succeed, but failed: %v", err)
+	}
+	defer func() {
+		if err := ds.Delete(ctx); err != nil {
+			t.Logf("defer deletion failed: %v", err)
+		}
+	}()
+
+	// Now, get the dataset without specifying policy version
+	md, err = ds.Metadata(ctx)
+	if err != nil {
+		t.Fatalf("Metadata: %v", err)
+	}
+	for _, entry := range md.Access {
+		if entry.Entity == wantEntry.Entity &&
+			entry.Condition != nil {
+			t.Fatalf("got policy with condition without specifying access policy version")
+		}
+	}
+
+	// Re-fetch metadata with access policy specified.
+	md, err = ds.MetadataWithOptions(ctx, WithAccessPolicyVersion(3))
+	if err != nil {
+		t.Fatalf("Metadata (WithAccessPolicy): %v", err)
+	}
+	var foundEntry bool
+	for _, entry := range md.Access {
+		if entry.Entity == wantEntry.Entity {
+			if cmp.Equal(entry.Condition, wantEntry.Condition) {
+				foundEntry = true
+				break
+			}
+		}
+	}
+	if !foundEntry {
+		t.Fatalf("failed to find wanted entry in access list")
+	}
+
+	newAccess := append(origAccess, &AccessEntry{
+		Role:       ReaderRole,
+		Entity:     "allUsers",
+		EntityType: IAMMemberEntity,
+	})
+
+	// append another entry.  Should fail without sending access policy version since we have conditions present.
+	md, err = ds.Update(ctx, DatasetMetadataToUpdate{Access: newAccess}, "")
+	if err == nil {
+		t.Fatalf("Update succeeded where failure expected: %v", err)
+	}
+
+	md, err = ds.UpdateWithOptions(ctx, DatasetMetadataToUpdate{Access: newAccess}, "", WithAccessPolicyVersion(3))
+	if err != nil {
+		t.Fatalf("Update failed: %v", err)
+	}
+
+}
+
 // Comparison function for AccessEntries to enable order insensitive equality checking.
 func lessAccessEntries(x, y *AccessEntry) bool {
 	if x.Entity < y.Entity {

bigquery/dataset_test.go

@@ -420,6 +420,15 @@ func TestBQToDatasetMetadata(t *testing.T) {
 		Labels:   map[string]string{"x": "y"},
 		Access: []*bq.DatasetAccess{
 			{Role: "READER", UserByEmail: "[email protected]"},
+			{Role: "READER",
+				UserByEmail: "[email protected]",
+				Condition: &bq.Expr{
+					Description: "desc",
+					Expression:  "expr",
+					Location:    "loc",
+					Title:       "title",
+				},
+			},
 			{Role: "WRITER", GroupByEmail: "[email protected]"},
 			{
 				Dataset: &bq.DatasetAccessEntry{
@@ -456,7 +465,20 @@ func TestBQToDatasetMetadata(t *testing.T) {
 		Location:            "EU",
 		Labels:              map[string]string{"x": "y"},
 		Access: []*AccessEntry{
-			{Role: ReaderRole, Entity: "[email protected]", EntityType: UserEmailEntity},
+			{Role: ReaderRole,
+				Entity:     "[email protected]",
+				EntityType: UserEmailEntity,
+			},
+			{Role: ReaderRole,
+				Entity:     "[email protected]",
+				EntityType: UserEmailEntity,
+				Condition: &Expr{
+					Title:       "title",
+					Expression:  "expr",
+					Location:    "loc",
+					Description: "desc",
+				},
+			},
 			{Role: WriterRole, Entity: "[email protected]", EntityType: GroupEmailEntity},
 			{
 				EntityType: DatasetEntity,
@@ -536,6 +558,12 @@ func TestConvertAccessEntry(t *testing.T) {
 		{Role: OwnerRole, Entity: "e", EntityType: UserEmailEntity},
 		{Role: ReaderRole, Entity: "e", EntityType: SpecialGroupEntity},
 		{Role: ReaderRole, Entity: "e", EntityType: IAMMemberEntity},
+		{Role: WriterRole, Entity: "e", EntityType: IAMMemberEntity,
+			Condition: &Expr{Expression: "expr",
+				Title:       "title",
+				Location:    "loc",
+				Description: "desc",
+			}},
 		{Role: ReaderRole, EntityType: ViewEntity,
 			View: &Table{ProjectID: "p", DatasetID: "d", TableID: "t", c: c}},
 		{Role: ReaderRole, EntityType: RoutineEntity,