From 2f417a39d93402fbb1e5e3001645019782d7d656 Mon Sep 17 00:00:00 2001
From: shollyman <shollyman@google.com>
Date: Thu, 14 Jan 2021 09:24:03 -0800
Subject: [PATCH] feat(bigquery): add support for allowing Javascript UDFs to
 indicate determinism (#3534)

BigQuery supports javascript UDFs, which allow users to define function
bodies to be executed via javascrtipt.  Allowing users to indicate
determinism level of the UDF helps the query engine make more efficient
choices around aspects like query caching.

Fixes: #3533
---
 bigquery/integration_test.go | 23 ++++++++++++++++++----
 bigquery/routine.go          | 38 +++++++++++++++++++++++++++++++++---
 bigquery/routine_test.go     |  2 ++
 3 files changed, 56 insertions(+), 7 deletions(-)

diff --git a/bigquery/integration_test.go b/bigquery/integration_test.go
index 94d98f1401d..093188187e0 100644
--- a/bigquery/integration_test.go
+++ b/bigquery/integration_test.go
@@ -2916,18 +2916,33 @@ func TestIntegration_RoutineJSUDF(t *testing.T) {
 	// Create a scalar UDF routine via API.
 	routineID := routineIDs.New()
 	routine := dataset.Routine(routineID)
-	err := routine.Create(ctx, &RoutineMetadata{
+	meta := &RoutineMetadata{
 		Language: "JAVASCRIPT", Type: "SCALAR_FUNCTION",
-		Description: "capitalizes using javascript",
+		Description:      "capitalizes using javascript",
+		DeterminismLevel: Deterministic,
 		Arguments: []*RoutineArgument{
 			{Name: "instr", Kind: "FIXED_TYPE", DataType: &StandardSQLDataType{TypeKind: "STRING"}},
 		},
 		ReturnType: &StandardSQLDataType{TypeKind: "STRING"},
 		Body:       "return instr.toUpperCase();",
-	})
-	if err != nil {
+	}
+	if err := routine.Create(ctx, meta); err != nil {
 		t.Fatalf("Create: %v", err)
 	}
+
+	newMeta := &RoutineMetadataToUpdate{
+		Language:    meta.Language,
+		Body:        meta.Body,
+		Arguments:   meta.Arguments,
+		Description: meta.Description,
+		ReturnType:  meta.ReturnType,
+		Type:        meta.Type,
+
+		DeterminismLevel: NotDeterministic,
+	}
+	if _, err := routine.Update(ctx, newMeta, ""); err != nil {
+		t.Fatalf("Update: %v", err)
+	}
 }
 
 func TestIntegration_RoutineComplexTypes(t *testing.T) {
diff --git a/bigquery/routine.go b/bigquery/routine.go
index 58f1a78b798..6f6c3a177f4 100644
--- a/bigquery/routine.go
+++ b/bigquery/routine.go
@@ -129,13 +129,27 @@ func (r *Routine) Delete(ctx context.Context) (err error) {
 	return req.Do()
 }
 
+// RoutineDeterminism specifies the level of determinism that javascript User Defined Functions
+// exhibit.
+type RoutineDeterminism string
+
+const (
+	// Deterministic indicates that two calls with the same input to a UDF yield the same output.
+	Deterministic RoutineDeterminism = "DETERMINISTIC"
+	// NotDeterministic indicates that the output of the UDF is not guaranteed to yield the same
+	// output each time for a given set of inputs.
+	NotDeterministic RoutineDeterminism = "NOT_DETERMINISTIC"
+)
+
 // RoutineMetadata represents details of a given BigQuery Routine.
 type RoutineMetadata struct {
 	ETag string
 	// Type indicates the type of routine, such as SCALAR_FUNCTION or PROCEDURE.
-	Type             string
-	CreationTime     time.Time
-	Description      string
+	Type         string
+	CreationTime time.Time
+	Description  string
+	// DeterminismLevel is only applicable to Javascript UDFs.
+	DeterminismLevel RoutineDeterminism
 	LastModifiedTime time.Time
 	// Language of the routine, such as SQL or JAVASCRIPT.
 	Language string
@@ -161,6 +175,7 @@ func (rm *RoutineMetadata) toBQ() (*bq.Routine, error) {
 		return r, nil
 	}
 	r.Description = rm.Description
+	r.DeterminismLevel = string(rm.DeterminismLevel)
 	r.Language = rm.Language
 	r.RoutineType = rm.Type
 	r.DefinitionBody = rm.Body
@@ -280,6 +295,7 @@ func routineArgumentsToBQ(in []*RoutineArgument) ([]*bq.Argument, error) {
 type RoutineMetadataToUpdate struct {
 	Arguments         []*RoutineArgument
 	Description       optional.String
+	DeterminismLevel  optional.String
 	Type              optional.String
 	Language          optional.String
 	Body              optional.String
@@ -299,6 +315,21 @@ func (rm *RoutineMetadataToUpdate) toBQ() (*bq.Routine, error) {
 		r.Description = optional.ToString(rm.Description)
 		forceSend("Description")
 	}
+	if rm.DeterminismLevel != nil {
+		processed := false
+		// Allow either string or RoutineDeterminism, a type based on string.
+		if x, ok := rm.DeterminismLevel.(RoutineDeterminism); ok {
+			r.DeterminismLevel = string(x)
+			processed = true
+		}
+		if x, ok := rm.DeterminismLevel.(string); ok {
+			r.DeterminismLevel = x
+			processed = true
+		}
+		if !processed {
+			panic(fmt.Sprintf("DeterminismLevel should be either type string or RoutineDetermism in update, got %T", rm.DeterminismLevel))
+		}
+	}
 	if rm.Arguments != nil {
 		if len(rm.Arguments) == 0 {
 			nullField("Arguments")
@@ -348,6 +379,7 @@ func bqToRoutineMetadata(r *bq.Routine) (*RoutineMetadata, error) {
 		Type:              r.RoutineType,
 		CreationTime:      unixMillisToTime(r.CreationTime),
 		Description:       r.Description,
+		DeterminismLevel:  RoutineDeterminism(r.DeterminismLevel),
 		LastModifiedTime:  unixMillisToTime(r.LastModifiedTime),
 		Language:          r.Language,
 		ImportedLibraries: r.ImportedLibraries,
diff --git a/bigquery/routine_test.go b/bigquery/routine_test.go
index 00122a44e0c..f0c122ea113 100644
--- a/bigquery/routine_test.go
+++ b/bigquery/routine_test.go
@@ -80,6 +80,7 @@ func TestRoutineTypeConversions(t *testing.T) {
 				DefinitionBody:   "body",
 				Description:      "desc",
 				Etag:             "etag",
+				DeterminismLevel: "DETERMINISTIC",
 				RoutineType:      "type",
 				Language:         "lang",
 				ReturnType:       &bq.StandardSqlDataType{TypeKind: "INT64"},
@@ -88,6 +89,7 @@ func TestRoutineTypeConversions(t *testing.T) {
 				CreationTime:     aTime,
 				LastModifiedTime: aTime,
 				Description:      "desc",
+				DeterminismLevel: Deterministic,
 				Body:             "body",
 				ETag:             "etag",
 				Type:             "type",