From c1ec0295127fd8e486a0b18010daa1c638bafe35 Mon Sep 17 00:00:00 2001 From: Seth Hollyman Date: Wed, 8 Nov 2023 19:40:07 +0000 Subject: [PATCH 1/4] feat(bigquery): add DataGovernanceType to routines --- bigquery/routine.go | 48 +++++++++++++-------- bigquery/routine_test.go | 93 +++++++++++++++++++++++++--------------- 2 files changed, 89 insertions(+), 52 deletions(-) diff --git a/bigquery/routine.go b/bigquery/routine.go index 2274f33fe360..9746daeaeb74 100644 --- a/bigquery/routine.go +++ b/bigquery/routine.go @@ -206,6 +206,11 @@ type RoutineMetadata struct { // For JAVASCRIPT function, it is the evaluated string in the AS clause of // a CREATE FUNCTION statement. Body string + + // Controls features like data masking. If set to `DATA_MASKING`, the function + // is validated and made available as a masking function. For more information, + // see: https://cloud.google.com/bigquery/docs/user-defined-functions#custom-mask + DataGovernanceType string } // RemoteFunctionOptions contains information for a remote user-defined function. @@ -278,6 +283,7 @@ func (rm *RoutineMetadata) toBQ() (*bq.Routine, error) { r.Language = rm.Language r.RoutineType = rm.Type r.DefinitionBody = rm.Body + r.DataGovernanceType = rm.DataGovernanceType rt, err := rm.ReturnType.toBQ() if err != nil { return nil, err @@ -405,15 +411,16 @@ func routineArgumentsToBQ(in []*RoutineArgument) ([]*bq.Argument, error) { // RoutineMetadataToUpdate governs updating a routine. type RoutineMetadataToUpdate struct { - Arguments []*RoutineArgument - Description optional.String - DeterminismLevel optional.String - Type optional.String - Language optional.String - Body optional.String - ImportedLibraries []string - ReturnType *StandardSQLDataType - ReturnTableType *StandardSQLTableType + Arguments []*RoutineArgument + Description optional.String + DeterminismLevel optional.String + Type optional.String + Language optional.String + Body optional.String + ImportedLibraries []string + ReturnType *StandardSQLDataType + ReturnTableType *StandardSQLTableType + DataGovernanceType optional.String } func (rm *RoutineMetadataToUpdate) toBQ() (*bq.Routine, error) { @@ -491,20 +498,25 @@ func (rm *RoutineMetadataToUpdate) toBQ() (*bq.Routine, error) { r.ReturnTableType = tt forceSend("ReturnTableType") } + if rm.DataGovernanceType != nil { + r.DataGovernanceType = optional.ToString(rm.DataGovernanceType) + forceSend("DataGovernanceType") + } return r, nil } func bqToRoutineMetadata(r *bq.Routine) (*RoutineMetadata, error) { meta := &RoutineMetadata{ - ETag: r.Etag, - Type: r.RoutineType, - CreationTime: unixMillisToTime(r.CreationTime), - Description: r.Description, - DeterminismLevel: RoutineDeterminism(r.DeterminismLevel), - LastModifiedTime: unixMillisToTime(r.LastModifiedTime), - Language: r.Language, - ImportedLibraries: r.ImportedLibraries, - Body: r.DefinitionBody, + ETag: r.Etag, + Type: r.RoutineType, + CreationTime: unixMillisToTime(r.CreationTime), + Description: r.Description, + DeterminismLevel: RoutineDeterminism(r.DeterminismLevel), + LastModifiedTime: unixMillisToTime(r.LastModifiedTime), + Language: r.Language, + ImportedLibraries: r.ImportedLibraries, + Body: r.DefinitionBody, + DataGovernanceType: r.DataGovernanceType, } args, err := bqToArgs(r.Arguments) if err != nil { diff --git a/bigquery/routine_test.go b/bigquery/routine_test.go index 5527a751101f..4b0e1ed02a34 100644 --- a/bigquery/routine_test.go +++ b/bigquery/routine_test.go @@ -72,9 +72,16 @@ func TestRoutineTypeConversions(t *testing.T) { in interface{} want interface{} }{ - {"empty", "ToRoutineMetadata", &bq.Routine{}, &RoutineMetadata{}}, - {"basic", "ToRoutineMetadata", - &bq.Routine{ + { + name: "empty", + conversion: "ToRoutineMetadata", + in: &bq.Routine{}, + want: &RoutineMetadata{}, + }, + { + name: "basic", + conversion: "ToRoutineMetadata", + in: &bq.Routine{ CreationTime: aTimeMillis, LastModifiedTime: aTimeMillis, DefinitionBody: "body", @@ -89,8 +96,9 @@ func TestRoutineTypeConversions(t *testing.T) { {Name: "field", Type: &bq.StandardSqlDataType{TypeKind: "FLOAT64"}}, }, }, + DataGovernanceType: "DATA_MASKING", }, - &RoutineMetadata{ + want: &RoutineMetadata{ CreationTime: aTime, LastModifiedTime: aTime, Description: "desc", @@ -105,55 +113,73 @@ func TestRoutineTypeConversions(t *testing.T) { {Name: "field", Type: &StandardSQLDataType{TypeKind: "FLOAT64"}}, }, }, - }}, - {"body_and_libs", "FromRoutineMetadataToUpdate", - &RoutineMetadataToUpdate{ - Body: "body", - ImportedLibraries: []string{"foo", "bar"}, - ReturnType: &StandardSQLDataType{TypeKind: "FOO"}, + DataGovernanceType: "DATA_MASKING", }, - &bq.Routine{ - DefinitionBody: "body", - ImportedLibraries: []string{"foo", "bar"}, - ReturnType: &bq.StandardSqlDataType{TypeKind: "FOO"}, - ForceSendFields: []string{"DefinitionBody", "ImportedLibraries", "ReturnType"}, - }}, - {"null_fields", "FromRoutineMetadataToUpdate", - &RoutineMetadataToUpdate{ + }, + { + name: "body_and_libs", + conversion: "FromRoutineMetadataToUpdate", + in: &RoutineMetadataToUpdate{ + Body: "body", + ImportedLibraries: []string{"foo", "bar"}, + ReturnType: &StandardSQLDataType{TypeKind: "FOO"}, + DataGovernanceType: "DATA_MASKING", + }, + want: &bq.Routine{ + DefinitionBody: "body", + ImportedLibraries: []string{"foo", "bar"}, + ReturnType: &bq.StandardSqlDataType{TypeKind: "FOO"}, + DataGovernanceType: "DATA_MASKING", + ForceSendFields: []string{"DefinitionBody", "ImportedLibraries", "ReturnType", "DataGovernanceType"}, + }, + }, + { + name: "null_fields", + conversion: "FromRoutineMetadataToUpdate", + in: &RoutineMetadataToUpdate{ Type: "type", Arguments: []*RoutineArgument{}, ImportedLibraries: []string{}, }, - &bq.Routine{ + want: &bq.Routine{ RoutineType: "type", ForceSendFields: []string{"RoutineType"}, NullFields: []string{"Arguments", "ImportedLibraries"}, - }}, - {"empty", "ToRoutineArgument", - &bq.Argument{}, - &RoutineArgument{}}, - {"basic", "ToRoutineArgument", - &bq.Argument{ + }, + }, + {name: "empty", + conversion: "ToRoutineArgument", + in: &bq.Argument{}, + want: &RoutineArgument{}}, + { + name: "basic", + conversion: "ToRoutineArgument", + in: &bq.Argument{ Name: "foo", ArgumentKind: "bar", Mode: "baz", }, - &RoutineArgument{ + want: &RoutineArgument{ Name: "foo", Kind: "bar", Mode: "baz", - }}, - {"empty", "FromRoutineArgument", - &RoutineArgument{}, - &bq.Argument{}, + }, }, - {"basic", "FromRoutineArgument", - &RoutineArgument{ + { + name: "empty", + conversion: "FromRoutineArgument", + in: &RoutineArgument{}, + want: &bq.Argument{}, + }, + { + name: "basic", + conversion: "FromRoutineArgument", + in: &RoutineArgument{ Name: "foo", Kind: "bar", Mode: "baz", }, - &bq.Argument{ + want: &bq.Argument{ Name: "foo", ArgumentKind: "bar", Mode: "baz", @@ -162,7 +188,6 @@ func TestRoutineTypeConversions(t *testing.T) { for _, test := range tests { t.Run(fmt.Sprintf("%s/%s", test.conversion, test.name), func(t *testing.T) { - t.Parallel() testRoutineConversion(t, test.conversion, test.in, test.want) }) } From 9003f868ec88deb7dac64898d0a0a70423ef6e8d Mon Sep 17 00:00:00 2001 From: Seth Hollyman Date: Wed, 8 Nov 2023 19:43:37 +0000 Subject: [PATCH 2/4] comment cleanup --- bigquery/routine.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bigquery/routine.go b/bigquery/routine.go index 9746daeaeb74..77fd5d90d57b 100644 --- a/bigquery/routine.go +++ b/bigquery/routine.go @@ -207,7 +207,7 @@ type RoutineMetadata struct { // a CREATE FUNCTION statement. Body string - // Controls features like data masking. If set to `DATA_MASKING`, the function + // For data governance use cases. If set to "DATA_MASKING", the function // is validated and made available as a masking function. For more information, // see: https://cloud.google.com/bigquery/docs/user-defined-functions#custom-mask DataGovernanceType string From 276805b0efd87aec3256809785cd5459c440792a Mon Sep 17 00:00:00 2001 From: Seth Hollyman Date: Thu, 9 Nov 2023 22:02:05 +0000 Subject: [PATCH 3/4] address reviewer feedback (more tests, linebreak) --- bigquery/routine_test.go | 53 +++++++++++++++++++++++++++++++++++++++- 1 file changed, 52 insertions(+), 1 deletion(-) diff --git a/bigquery/routine_test.go b/bigquery/routine_test.go index 4b0e1ed02a34..e2abf47ea9e7 100644 --- a/bigquery/routine_test.go +++ b/bigquery/routine_test.go @@ -32,6 +32,12 @@ func testRoutineConversion(t *testing.T, conversion string, in interface{}, want t.Fatalf("failed input type conversion (bq.Routine): %v", in) } got, err = bqToRoutineMetadata(input) + case "FromRoutineMetadata": + input, ok := in.(*RoutineMetadata) + if !ok { + t.Fatalf("failed input type conversion (bq.RoutineMetadata): %v", in) + } + got, err = input.toBQ() case "FromRoutineMetadataToUpdate": input, ok := in.(*RoutineMetadataToUpdate) if !ok { @@ -78,6 +84,12 @@ func TestRoutineTypeConversions(t *testing.T) { in: &bq.Routine{}, want: &RoutineMetadata{}, }, + { + name: "empty", + conversion: "FromRoutineMetadata", + in: &RoutineMetadata{}, + want: &bq.Routine{}, + }, { name: "basic", conversion: "ToRoutineMetadata", @@ -116,6 +128,44 @@ func TestRoutineTypeConversions(t *testing.T) { DataGovernanceType: "DATA_MASKING", }, }, + { + name: "basic", + conversion: "FromRoutineMetadata", + in: &RoutineMetadata{ + CreationTime: aTime, + LastModifiedTime: aTime, + Description: "desc", + DeterminismLevel: Deterministic, + Body: "body", + ETag: "etag", + Type: "type", + Language: "lang", + ReturnType: &StandardSQLDataType{TypeKind: "INT64"}, + ReturnTableType: &StandardSQLTableType{ + Columns: []*StandardSQLField{ + {Name: "field", Type: &StandardSQLDataType{TypeKind: "FLOAT64"}}, + }, + }, + DataGovernanceType: "DATA_MASKING", + }, + want: &bq.Routine{ + CreationTime: aTimeMillis, + LastModifiedTime: aTimeMillis, + DefinitionBody: "body", + Description: "desc", + Etag: "etag", + DeterminismLevel: "DETERMINISTIC", + RoutineType: "type", + Language: "lang", + ReturnType: &bq.StandardSqlDataType{TypeKind: "INT64"}, + ReturnTableType: &bq.StandardSqlTableType{ + Columns: []*bq.StandardSqlField{ + {Name: "field", Type: &bq.StandardSqlDataType{TypeKind: "FLOAT64"}}, + }, + }, + DataGovernanceType: "DATA_MASKING", + }, + }, { name: "body_and_libs", conversion: "FromRoutineMetadataToUpdate", @@ -147,7 +197,8 @@ func TestRoutineTypeConversions(t *testing.T) { NullFields: []string{"Arguments", "ImportedLibraries"}, }, }, - {name: "empty", + { + name: "empty", conversion: "ToRoutineArgument", in: &bq.Argument{}, want: &RoutineArgument{}}, From 52d9394e19e393de9449cbc456e9c82b6da4d094 Mon Sep 17 00:00:00 2001 From: Seth Hollyman Date: Thu, 9 Nov 2023 22:28:39 +0000 Subject: [PATCH 4/4] add it tests, cleanup unit --- bigquery/routine_integration_test.go | 43 +++++++++++++++++++++++----- bigquery/routine_test.go | 9 +----- 2 files changed, 37 insertions(+), 15 deletions(-) diff --git a/bigquery/routine_integration_test.go b/bigquery/routine_integration_test.go index d2a7cd66dc82..ffd8228417c3 100644 --- a/bigquery/routine_integration_test.go +++ b/bigquery/routine_integration_test.go @@ -19,11 +19,11 @@ import ( "fmt" "testing" + "cloud.google.com/go/bigquery/connection/apiv1/connectionpb" "cloud.google.com/go/internal" "cloud.google.com/go/internal/testutil" gax "github.com/googleapis/gax-go/v2" "google.golang.org/api/iterator" - "google.golang.org/genproto/googleapis/cloud/bigquery/connection/v1" ) func TestIntegration_RoutineScalarUDF(t *testing.T) { @@ -53,6 +53,35 @@ func TestIntegration_RoutineScalarUDF(t *testing.T) { } } +func TestIntegration_RoutineDataGovernance(t *testing.T) { + if client == nil { + t.Skip("Integration tests skipped") + } + ctx := context.Background() + + // Create a scalar UDF routine via API. + routineID := routineIDs.New() + routine := dataset.Routine(routineID) + err := routine.Create(ctx, &RoutineMetadata{ + Type: "SCALAR_FUNCTION", + Language: "SQL", + Body: "x", + Arguments: []*RoutineArgument{ + { + Name: "x", + DataType: &StandardSQLDataType{ + TypeKind: "INT64", + }, + }, + }, + ReturnType: &StandardSQLDataType{TypeKind: "INT64"}, + DataGovernanceType: "DATA_MASKING", + }) + if err != nil { + t.Fatalf("Create: %v", err) + } +} + func TestIntegration_RoutineJSUDF(t *testing.T) { if client == nil { t.Skip("Integration tests skipped") @@ -146,27 +175,27 @@ func TestIntegration_RoutineRemoteUDF(t *testing.T) { func createConnection(ctx context.Context, t *testing.T, parent, name string) (cleanup func(), connectionID string, err error) { fullname := fmt.Sprintf("%s/connections/%s", parent, name) - conn, err := connectionsClient.CreateConnection(ctx, &connection.CreateConnectionRequest{ + conn, err := connectionsClient.CreateConnection(ctx, &connectionpb.CreateConnectionRequest{ Parent: parent, ConnectionId: name, - Connection: &connection.Connection{ + Connection: &connectionpb.Connection{ FriendlyName: name, - Properties: &connection.Connection_CloudResource{ - CloudResource: &connection.CloudResourceProperties{}, + Properties: &connectionpb.Connection_CloudResource{ + CloudResource: &connectionpb.CloudResourceProperties{}, }, }, }) if err != nil { return } - conn, err = connectionsClient.GetConnection(ctx, &connection.GetConnectionRequest{ + conn, err = connectionsClient.GetConnection(ctx, &connectionpb.GetConnectionRequest{ Name: fullname, }) if err != nil { return } cleanup = func() { - err := connectionsClient.DeleteConnection(ctx, &connection.DeleteConnectionRequest{ + err := connectionsClient.DeleteConnection(ctx, &connectionpb.DeleteConnectionRequest{ Name: fullname, }) if err != nil { diff --git a/bigquery/routine_test.go b/bigquery/routine_test.go index e2abf47ea9e7..23ba01167178 100644 --- a/bigquery/routine_test.go +++ b/bigquery/routine_test.go @@ -59,9 +59,8 @@ func testRoutineConversion(t *testing.T, conversion string, in interface{}, want default: t.Fatalf("invalid comparison: %s", conversion) } - if err != nil { - t.Fatalf("failed conversion function for %q", conversion) + t.Fatalf("failed conversion function for %q: %v", conversion, err) } if diff := testutil.Diff(got, want); diff != "" { t.Fatalf("%+v: -got, +want:\n%s", in, diff) @@ -132,12 +131,9 @@ func TestRoutineTypeConversions(t *testing.T) { name: "basic", conversion: "FromRoutineMetadata", in: &RoutineMetadata{ - CreationTime: aTime, - LastModifiedTime: aTime, Description: "desc", DeterminismLevel: Deterministic, Body: "body", - ETag: "etag", Type: "type", Language: "lang", ReturnType: &StandardSQLDataType{TypeKind: "INT64"}, @@ -149,11 +145,8 @@ func TestRoutineTypeConversions(t *testing.T) { DataGovernanceType: "DATA_MASKING", }, want: &bq.Routine{ - CreationTime: aTimeMillis, - LastModifiedTime: aTimeMillis, DefinitionBody: "body", Description: "desc", - Etag: "etag", DeterminismLevel: "DETERMINISTIC", RoutineType: "type", Language: "lang",