From 20054504408a6fba4bba18ee6321db8d0003eff2 Mon Sep 17 00:00:00 2001
From: shollyman <shollyman@google.com>
Date: Thu, 29 Oct 2020 13:55:21 -0700
Subject: [PATCH] feat(bigquery): Add support for authorized UDFs (#3084)

* feat(bigquery): Add support for authorized UDFs

Much like authorized views, BigQuery now supports authorized User
Defined Functions (UDFs). This PR augments the `Access` field of a
dataset entity to support routine references.
---
 bigquery/dataset.go          |  9 +++++++++
 bigquery/dataset_test.go     |  4 +++-
 bigquery/integration_test.go | 19 ++++++++++++++++++-
 bigquery/routine.go          |  8 ++++++++
 4 files changed, 38 insertions(+), 2 deletions(-)

diff --git a/bigquery/dataset.go b/bigquery/dataset.go
index 48cb8e36063..91b649eed8c 100644
--- a/bigquery/dataset.go
+++ b/bigquery/dataset.go
@@ -642,6 +642,7 @@ type AccessEntry struct {
 	EntityType EntityType // The type of entity
 	Entity     string     // The entity (individual or group) granted access
 	View       *Table     // The view granted access (EntityType must be ViewEntity)
+	Routine    *Routine   // The routine granted access (only UDF currently supported)
 }
 
 // AccessRole is the level of access to grant to a dataset.
@@ -679,6 +680,9 @@ const (
 	// IAMMemberEntity represents entities present in IAM but not represented using
 	// the other entity types.
 	IAMMemberEntity
+
+	// RoutineEntity is a BigQuery routine, referencing a User Defined Function (UDF).
+	RoutineEntity
 )
 
 func (e *AccessEntry) toBQ() (*bq.DatasetAccess, error) {
@@ -696,6 +700,8 @@ func (e *AccessEntry) toBQ() (*bq.DatasetAccess, error) {
 		q.View = e.View.toBQ()
 	case IAMMemberEntity:
 		q.IamMember = e.Entity
+	case RoutineEntity:
+		q.Routine = e.Routine.toBQ()
 	default:
 		return nil, fmt.Errorf("bigquery: unknown entity type %d", e.EntityType)
 	}
@@ -723,6 +729,9 @@ func bqToAccessEntry(q *bq.DatasetAccess, c *Client) (*AccessEntry, error) {
 	case q.IamMember != "":
 		e.Entity = q.IamMember
 		e.EntityType = IAMMemberEntity
+	case q.Routine != nil:
+		e.Routine = c.DatasetInProject(q.Routine.ProjectId, q.Routine.DatasetId).Routine(q.Routine.RoutineId)
+		e.EntityType = RoutineEntity
 	default:
 		return nil, errors.New("bigquery: invalid access value")
 	}
diff --git a/bigquery/dataset_test.go b/bigquery/dataset_test.go
index 782272cf3b9..351947b73e9 100644
--- a/bigquery/dataset_test.go
+++ b/bigquery/dataset_test.go
@@ -452,6 +452,8 @@ func TestConvertAccessEntry(t *testing.T) {
 		{Role: ReaderRole, Entity: "e", EntityType: IAMMemberEntity},
 		{Role: ReaderRole, EntityType: ViewEntity,
 			View: &Table{ProjectID: "p", DatasetID: "d", TableID: "t", c: c}},
+		{Role: ReaderRole, EntityType: RoutineEntity,
+			Routine: &Routine{ProjectID: "p", DatasetID: "d", RoutineID: "r", c: c}},
 	} {
 		q, err := e.toBQ()
 		if err != nil {
@@ -461,7 +463,7 @@ func TestConvertAccessEntry(t *testing.T) {
 		if err != nil {
 			t.Fatal(err)
 		}
-		if diff := testutil.Diff(got, e, cmp.AllowUnexported(Table{}, Client{})); diff != "" {
+		if diff := testutil.Diff(got, e, cmp.AllowUnexported(Table{}, Client{}, Routine{})); diff != "" {
 			t.Errorf("got=-, want=+:\n%s", diff)
 		}
 	}
diff --git a/bigquery/integration_test.go b/bigquery/integration_test.go
index 58742dc5a7d..14573e2bb3a 100644
--- a/bigquery/integration_test.go
+++ b/bigquery/integration_test.go
@@ -716,6 +716,19 @@ func TestIntegration_DatasetUpdateAccess(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
+
+	// Create a sample UDF so we can verify adding authorized UDFs
+	routineID := routineIDs.New()
+	routine := dataset.Routine(routineID)
+
+	sql := fmt.Sprintf(`
+			CREATE FUNCTION `+"`%s`"+`(x INT64) AS (x * 3);`,
+		routine.FullyQualifiedName())
+	if err := runQueryJob(ctx, sql); err != nil {
+		t.Fatal(err)
+	}
+	defer routine.Delete(ctx)
+
 	origAccess := append([]*AccessEntry(nil), md.Access...)
 	newEntries := []*AccessEntry{
 		{
@@ -728,6 +741,10 @@ func TestIntegration_DatasetUpdateAccess(t *testing.T) {
 			Entity:     "allUsers",
 			EntityType: IAMMemberEntity,
 		},
+		{
+			EntityType: RoutineEntity,
+			Routine:    routine,
+		},
 	}
 
 	newAccess := append(md.Access, newEntries...)
@@ -743,7 +760,7 @@ func TestIntegration_DatasetUpdateAccess(t *testing.T) {
 		}
 	}()
 
-	if diff := testutil.Diff(md.Access, newAccess, cmpopts.SortSlices(lessAccessEntries)); diff != "" {
+	if diff := testutil.Diff(md.Access, newAccess, cmpopts.SortSlices(lessAccessEntries), cmpopts.IgnoreUnexported(Routine{})); diff != "" {
 		t.Fatalf("got=-, want=+:\n%s", diff)
 	}
 }
diff --git a/bigquery/routine.go b/bigquery/routine.go
index a7026b6fc18..58f1a78b798 100644
--- a/bigquery/routine.go
+++ b/bigquery/routine.go
@@ -36,6 +36,14 @@ type Routine struct {
 	c *Client
 }
 
+func (r *Routine) toBQ() *bq.RoutineReference {
+	return &bq.RoutineReference{
+		ProjectId: r.ProjectID,
+		DatasetId: r.DatasetID,
+		RoutineId: r.RoutineID,
+	}
+}
+
 // FullyQualifiedName returns an identifer for the routine in project.dataset.routine format.
 func (r *Routine) FullyQualifiedName() string {
 	return fmt.Sprintf("%s.%s.%s", r.ProjectID, r.DatasetID, r.RoutineID)