Skip to content

Commit 57491ae

Browse files
shollymanshollyman
authored
feat(bigquery): add DataGovernanceType to routines (#8990)
* feat(bigquery): add DataGovernanceType to routines * comment cleanup * address reviewer feedback (more tests, linebreak) * add it tests, cleanup unit
1 parent a0b64f8 commit 57491ae

File tree

3 files changed

+171
-61
lines changed

3 files changed

+171
-61
lines changed

β€Žbigquery/routine.go

Lines changed: 30 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -206,6 +206,11 @@ type RoutineMetadata struct {
206206
// For JAVASCRIPT function, it is the evaluated string in the AS clause of
207207
// a CREATE FUNCTION statement.
208208
Body string
209+
210+
// For data governance use cases. If set to "DATA_MASKING", the function
211+
// is validated and made available as a masking function. For more information,
212+
// see: https://cloud.google.com/bigquery/docs/user-defined-functions#custom-mask
213+
DataGovernanceType string
209214
}
210215

211216
// RemoteFunctionOptions contains information for a remote user-defined function.
@@ -278,6 +283,7 @@ func (rm *RoutineMetadata) toBQ() (*bq.Routine, error) {
278283
r.Language = rm.Language
279284
r.RoutineType = rm.Type
280285
r.DefinitionBody = rm.Body
286+
r.DataGovernanceType = rm.DataGovernanceType
281287
rt, err := rm.ReturnType.toBQ()
282288
if err != nil {
283289
return nil, err
@@ -405,15 +411,16 @@ func routineArgumentsToBQ(in []*RoutineArgument) ([]*bq.Argument, error) {
405411

406412
// RoutineMetadataToUpdate governs updating a routine.
407413
type RoutineMetadataToUpdate struct {
408-
Arguments []*RoutineArgument
409-
Description optional.String
410-
DeterminismLevel optional.String
411-
Type optional.String
412-
Language optional.String
413-
Body optional.String
414-
ImportedLibraries []string
415-
ReturnType *StandardSQLDataType
416-
ReturnTableType *StandardSQLTableType
414+
Arguments []*RoutineArgument
415+
Description optional.String
416+
DeterminismLevel optional.String
417+
Type optional.String
418+
Language optional.String
419+
Body optional.String
420+
ImportedLibraries []string
421+
ReturnType *StandardSQLDataType
422+
ReturnTableType *StandardSQLTableType
423+
DataGovernanceType optional.String
417424
}
418425

419426
func (rm *RoutineMetadataToUpdate) toBQ() (*bq.Routine, error) {
@@ -491,20 +498,25 @@ func (rm *RoutineMetadataToUpdate) toBQ() (*bq.Routine, error) {
491498
r.ReturnTableType = tt
492499
forceSend("ReturnTableType")
493500
}
501+
if rm.DataGovernanceType != nil {
502+
r.DataGovernanceType = optional.ToString(rm.DataGovernanceType)
503+
forceSend("DataGovernanceType")
504+
}
494505
return r, nil
495506
}
496507

497508
func bqToRoutineMetadata(r *bq.Routine) (*RoutineMetadata, error) {
498509
meta := &RoutineMetadata{
499-
ETag: r.Etag,
500-
Type: r.RoutineType,
501-
CreationTime: unixMillisToTime(r.CreationTime),
502-
Description: r.Description,
503-
DeterminismLevel: RoutineDeterminism(r.DeterminismLevel),
504-
LastModifiedTime: unixMillisToTime(r.LastModifiedTime),
505-
Language: r.Language,
506-
ImportedLibraries: r.ImportedLibraries,
507-
Body: r.DefinitionBody,
510+
ETag: r.Etag,
511+
Type: r.RoutineType,
512+
CreationTime: unixMillisToTime(r.CreationTime),
513+
Description: r.Description,
514+
DeterminismLevel: RoutineDeterminism(r.DeterminismLevel),
515+
LastModifiedTime: unixMillisToTime(r.LastModifiedTime),
516+
Language: r.Language,
517+
ImportedLibraries: r.ImportedLibraries,
518+
Body: r.DefinitionBody,
519+
DataGovernanceType: r.DataGovernanceType,
508520
}
509521
args, err := bqToArgs(r.Arguments)
510522
if err != nil {

β€Žbigquery/routine_integration_test.go

Lines changed: 36 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -19,11 +19,11 @@ import (
1919
"fmt"
2020
"testing"
2121

22+
"cloud.google.com/go/bigquery/connection/apiv1/connectionpb"
2223
"cloud.google.com/go/internal"
2324
"cloud.google.com/go/internal/testutil"
2425
gax "github.com/googleapis/gax-go/v2"
2526
"google.golang.org/api/iterator"
26-
"google.golang.org/genproto/googleapis/cloud/bigquery/connection/v1"
2727
)
2828

2929
func TestIntegration_RoutineScalarUDF(t *testing.T) {
@@ -53,6 +53,35 @@ func TestIntegration_RoutineScalarUDF(t *testing.T) {
5353
}
5454
}
5555

56+
func TestIntegration_RoutineDataGovernance(t *testing.T) {
57+
if client == nil {
58+
t.Skip("Integration tests skipped")
59+
}
60+
ctx := context.Background()
61+
62+
// Create a scalar UDF routine via API.
63+
routineID := routineIDs.New()
64+
routine := dataset.Routine(routineID)
65+
err := routine.Create(ctx, &RoutineMetadata{
66+
Type: "SCALAR_FUNCTION",
67+
Language: "SQL",
68+
Body: "x",
69+
Arguments: []*RoutineArgument{
70+
{
71+
Name: "x",
72+
DataType: &StandardSQLDataType{
73+
TypeKind: "INT64",
74+
},
75+
},
76+
},
77+
ReturnType: &StandardSQLDataType{TypeKind: "INT64"},
78+
DataGovernanceType: "DATA_MASKING",
79+
})
80+
if err != nil {
81+
t.Fatalf("Create: %v", err)
82+
}
83+
}
84+
5685
func TestIntegration_RoutineJSUDF(t *testing.T) {
5786
if client == nil {
5887
t.Skip("Integration tests skipped")
@@ -146,27 +175,27 @@ func TestIntegration_RoutineRemoteUDF(t *testing.T) {
146175

147176
func createConnection(ctx context.Context, t *testing.T, parent, name string) (cleanup func(), connectionID string, err error) {
148177
fullname := fmt.Sprintf("%s/connections/%s", parent, name)
149-
conn, err := connectionsClient.CreateConnection(ctx, &connection.CreateConnectionRequest{
178+
conn, err := connectionsClient.CreateConnection(ctx, &connectionpb.CreateConnectionRequest{
150179
Parent: parent,
151180
ConnectionId: name,
152-
Connection: &connection.Connection{
181+
Connection: &connectionpb.Connection{
153182
FriendlyName: name,
154-
Properties: &connection.Connection_CloudResource{
155-
CloudResource: &connection.CloudResourceProperties{},
183+
Properties: &connectionpb.Connection_CloudResource{
184+
CloudResource: &connectionpb.CloudResourceProperties{},
156185
},
157186
},
158187
})
159188
if err != nil {
160189
return
161190
}
162-
conn, err = connectionsClient.GetConnection(ctx, &connection.GetConnectionRequest{
191+
conn, err = connectionsClient.GetConnection(ctx, &connectionpb.GetConnectionRequest{
163192
Name: fullname,
164193
})
165194
if err != nil {
166195
return
167196
}
168197
cleanup = func() {
169-
err := connectionsClient.DeleteConnection(ctx, &connection.DeleteConnectionRequest{
198+
err := connectionsClient.DeleteConnection(ctx, &connectionpb.DeleteConnectionRequest{
170199
Name: fullname,
171200
})
172201
if err != nil {

β€Žbigquery/routine_test.go

Lines changed: 105 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,12 @@ func testRoutineConversion(t *testing.T, conversion string, in interface{}, want
3232
t.Fatalf("failed input type conversion (bq.Routine): %v", in)
3333
}
3434
got, err = bqToRoutineMetadata(input)
35+
case "FromRoutineMetadata":
36+
input, ok := in.(*RoutineMetadata)
37+
if !ok {
38+
t.Fatalf("failed input type conversion (bq.RoutineMetadata): %v", in)
39+
}
40+
got, err = input.toBQ()
3541
case "FromRoutineMetadataToUpdate":
3642
input, ok := in.(*RoutineMetadataToUpdate)
3743
if !ok {
@@ -53,9 +59,8 @@ func testRoutineConversion(t *testing.T, conversion string, in interface{}, want
5359
default:
5460
t.Fatalf("invalid comparison: %s", conversion)
5561
}
56-
5762
if err != nil {
58-
t.Fatalf("failed conversion function for %q", conversion)
63+
t.Fatalf("failed conversion function for %q: %v", conversion, err)
5964
}
6065
if diff := testutil.Diff(got, want); diff != "" {
6166
t.Fatalf("%+v: -got, +want:\n%s", in, diff)
@@ -72,9 +77,22 @@ func TestRoutineTypeConversions(t *testing.T) {
7277
in interface{}
7378
want interface{}
7479
}{
75-
{"empty", "ToRoutineMetadata", &bq.Routine{}, &RoutineMetadata{}},
76-
{"basic", "ToRoutineMetadata",
77-
&bq.Routine{
80+
{
81+
name: "empty",
82+
conversion: "ToRoutineMetadata",
83+
in: &bq.Routine{},
84+
want: &RoutineMetadata{},
85+
},
86+
{
87+
name: "empty",
88+
conversion: "FromRoutineMetadata",
89+
in: &RoutineMetadata{},
90+
want: &bq.Routine{},
91+
},
92+
{
93+
name: "basic",
94+
conversion: "ToRoutineMetadata",
95+
in: &bq.Routine{
7896
CreationTime: aTimeMillis,
7997
LastModifiedTime: aTimeMillis,
8098
DefinitionBody: "body",
@@ -89,8 +107,9 @@ func TestRoutineTypeConversions(t *testing.T) {
89107
{Name: "field", Type: &bq.StandardSqlDataType{TypeKind: "FLOAT64"}},
90108
},
91109
},
110+
DataGovernanceType: "DATA_MASKING",
92111
},
93-
&RoutineMetadata{
112+
want: &RoutineMetadata{
94113
CreationTime: aTime,
95114
LastModifiedTime: aTime,
96115
Description: "desc",
@@ -105,55 +124,106 @@ func TestRoutineTypeConversions(t *testing.T) {
105124
{Name: "field", Type: &StandardSQLDataType{TypeKind: "FLOAT64"}},
106125
},
107126
},
108-
}},
109-
{"body_and_libs", "FromRoutineMetadataToUpdate",
110-
&RoutineMetadataToUpdate{
111-
Body: "body",
112-
ImportedLibraries: []string{"foo", "bar"},
113-
ReturnType: &StandardSQLDataType{TypeKind: "FOO"},
127+
DataGovernanceType: "DATA_MASKING",
114128
},
115-
&bq.Routine{
116-
DefinitionBody: "body",
117-
ImportedLibraries: []string{"foo", "bar"},
118-
ReturnType: &bq.StandardSqlDataType{TypeKind: "FOO"},
119-
ForceSendFields: []string{"DefinitionBody", "ImportedLibraries", "ReturnType"},
120-
}},
121-
{"null_fields", "FromRoutineMetadataToUpdate",
122-
&RoutineMetadataToUpdate{
129+
},
130+
{
131+
name: "basic",
132+
conversion: "FromRoutineMetadata",
133+
in: &RoutineMetadata{
134+
Description: "desc",
135+
DeterminismLevel: Deterministic,
136+
Body: "body",
137+
Type: "type",
138+
Language: "lang",
139+
ReturnType: &StandardSQLDataType{TypeKind: "INT64"},
140+
ReturnTableType: &StandardSQLTableType{
141+
Columns: []*StandardSQLField{
142+
{Name: "field", Type: &StandardSQLDataType{TypeKind: "FLOAT64"}},
143+
},
144+
},
145+
DataGovernanceType: "DATA_MASKING",
146+
},
147+
want: &bq.Routine{
148+
DefinitionBody: "body",
149+
Description: "desc",
150+
DeterminismLevel: "DETERMINISTIC",
151+
RoutineType: "type",
152+
Language: "lang",
153+
ReturnType: &bq.StandardSqlDataType{TypeKind: "INT64"},
154+
ReturnTableType: &bq.StandardSqlTableType{
155+
Columns: []*bq.StandardSqlField{
156+
{Name: "field", Type: &bq.StandardSqlDataType{TypeKind: "FLOAT64"}},
157+
},
158+
},
159+
DataGovernanceType: "DATA_MASKING",
160+
},
161+
},
162+
{
163+
name: "body_and_libs",
164+
conversion: "FromRoutineMetadataToUpdate",
165+
in: &RoutineMetadataToUpdate{
166+
Body: "body",
167+
ImportedLibraries: []string{"foo", "bar"},
168+
ReturnType: &StandardSQLDataType{TypeKind: "FOO"},
169+
DataGovernanceType: "DATA_MASKING",
170+
},
171+
want: &bq.Routine{
172+
DefinitionBody: "body",
173+
ImportedLibraries: []string{"foo", "bar"},
174+
ReturnType: &bq.StandardSqlDataType{TypeKind: "FOO"},
175+
DataGovernanceType: "DATA_MASKING",
176+
ForceSendFields: []string{"DefinitionBody", "ImportedLibraries", "ReturnType", "DataGovernanceType"},
177+
},
178+
},
179+
{
180+
name: "null_fields",
181+
conversion: "FromRoutineMetadataToUpdate",
182+
in: &RoutineMetadataToUpdate{
123183
Type: "type",
124184
Arguments: []*RoutineArgument{},
125185
ImportedLibraries: []string{},
126186
},
127-
&bq.Routine{
187+
want: &bq.Routine{
128188
RoutineType: "type",
129189
ForceSendFields: []string{"RoutineType"},
130190
NullFields: []string{"Arguments", "ImportedLibraries"},
131-
}},
132-
{"empty", "ToRoutineArgument",
133-
&bq.Argument{},
134-
&RoutineArgument{}},
135-
{"basic", "ToRoutineArgument",
136-
&bq.Argument{
191+
},
192+
},
193+
{
194+
name: "empty",
195+
conversion: "ToRoutineArgument",
196+
in: &bq.Argument{},
197+
want: &RoutineArgument{}},
198+
{
199+
name: "basic",
200+
conversion: "ToRoutineArgument",
201+
in: &bq.Argument{
137202
Name: "foo",
138203
ArgumentKind: "bar",
139204
Mode: "baz",
140205
},
141-
&RoutineArgument{
206+
want: &RoutineArgument{
142207
Name: "foo",
143208
Kind: "bar",
144209
Mode: "baz",
145-
}},
146-
{"empty", "FromRoutineArgument",
147-
&RoutineArgument{},
148-
&bq.Argument{},
210+
},
149211
},
150-
{"basic", "FromRoutineArgument",
151-
&RoutineArgument{
212+
{
213+
name: "empty",
214+
conversion: "FromRoutineArgument",
215+
in: &RoutineArgument{},
216+
want: &bq.Argument{},
217+
},
218+
{
219+
name: "basic",
220+
conversion: "FromRoutineArgument",
221+
in: &RoutineArgument{
152222
Name: "foo",
153223
Kind: "bar",
154224
Mode: "baz",
155225
},
156-
&bq.Argument{
226+
want: &bq.Argument{
157227
Name: "foo",
158228
ArgumentKind: "bar",
159229
Mode: "baz",
@@ -162,7 +232,6 @@ func TestRoutineTypeConversions(t *testing.T) {
162232

163233
for _, test := range tests {
164234
t.Run(fmt.Sprintf("%s/%s", test.conversion, test.name), func(t *testing.T) {
165-
t.Parallel()
166235
testRoutineConversion(t, test.conversion, test.in, test.want)
167236
})
168237
}

0 commit comments

Comments
 (0)