diff --git a/bigquery/routine.go b/bigquery/routine.go index 2274f33fe360..77fd5d90d57b 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 + + // 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 } // 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_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 5527a751101f..23ba01167178 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 { @@ -53,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) @@ -72,9 +77,22 @@ 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: "empty", + conversion: "FromRoutineMetadata", + in: &RoutineMetadata{}, + want: &bq.Routine{}, + }, + { + name: "basic", + conversion: "ToRoutineMetadata", + in: &bq.Routine{ CreationTime: aTimeMillis, LastModifiedTime: aTimeMillis, DefinitionBody: "body", @@ -89,8 +107,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 +124,106 @@ 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: "basic", + conversion: "FromRoutineMetadata", + in: &RoutineMetadata{ + Description: "desc", + DeterminismLevel: Deterministic, + Body: "body", + Type: "type", + Language: "lang", + ReturnType: &StandardSQLDataType{TypeKind: "INT64"}, + ReturnTableType: &StandardSQLTableType{ + Columns: []*StandardSQLField{ + {Name: "field", Type: &StandardSQLDataType{TypeKind: "FLOAT64"}}, + }, + }, + DataGovernanceType: "DATA_MASKING", + }, + want: &bq.Routine{ + DefinitionBody: "body", + Description: "desc", + 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", + 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 +232,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) }) }