1// Copyright 2015 Google LLC
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//      http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15package bigquery
16
17import (
18	"context"
19	"errors"
20	"fmt"
21	"strconv"
22	"testing"
23	"time"
24
25	"cloud.google.com/go/internal/testutil"
26	"github.com/google/go-cmp/cmp"
27	bq "google.golang.org/api/bigquery/v2"
28	itest "google.golang.org/api/iterator/testing"
29)
30
31// readServiceStub services read requests by returning data from an in-memory list of values.
32type listTablesStub struct {
33	expectedProject, expectedDataset string
34	tables                           []*bq.TableListTables
35}
36
37func (s *listTablesStub) listTables(it *TableIterator, pageSize int, pageToken string) (*bq.TableList, error) {
38	if it.dataset.ProjectID != s.expectedProject {
39		return nil, fmt.Errorf("wrong project id: %q", it.dataset.ProjectID)
40	}
41	if it.dataset.DatasetID != s.expectedDataset {
42		return nil, fmt.Errorf("wrong dataset id: %q", it.dataset.DatasetID)
43	}
44	const maxPageSize = 2
45	if pageSize <= 0 || pageSize > maxPageSize {
46		pageSize = maxPageSize
47	}
48	start := 0
49	if pageToken != "" {
50		var err error
51		start, err = strconv.Atoi(pageToken)
52		if err != nil {
53			return nil, err
54		}
55	}
56	end := start + pageSize
57	if end > len(s.tables) {
58		end = len(s.tables)
59	}
60	nextPageToken := ""
61	if end < len(s.tables) {
62		nextPageToken = strconv.Itoa(end)
63	}
64	return &bq.TableList{
65		Tables:        s.tables[start:end],
66		NextPageToken: nextPageToken,
67	}, nil
68}
69
70func TestTables(t *testing.T) {
71	c := &Client{projectID: "p1"}
72	inTables := []*bq.TableListTables{
73		{TableReference: &bq.TableReference{ProjectId: "p1", DatasetId: "d1", TableId: "t1"}},
74		{TableReference: &bq.TableReference{ProjectId: "p1", DatasetId: "d1", TableId: "t2"}},
75		{TableReference: &bq.TableReference{ProjectId: "p1", DatasetId: "d1", TableId: "t3"}},
76	}
77	outTables := []*Table{
78		{ProjectID: "p1", DatasetID: "d1", TableID: "t1", c: c},
79		{ProjectID: "p1", DatasetID: "d1", TableID: "t2", c: c},
80		{ProjectID: "p1", DatasetID: "d1", TableID: "t3", c: c},
81	}
82
83	lts := &listTablesStub{
84		expectedProject: "p1",
85		expectedDataset: "d1",
86		tables:          inTables,
87	}
88	old := listTables
89	listTables = lts.listTables // cannot use t.Parallel with this test
90	defer func() { listTables = old }()
91
92	msg, ok := itest.TestIterator(outTables,
93		func() interface{} { return c.Dataset("d1").Tables(context.Background()) },
94		func(it interface{}) (interface{}, error) { return it.(*TableIterator).Next() })
95	if !ok {
96		t.Error(msg)
97	}
98}
99
100// listModelsStub services list requests by returning data from an in-memory list of values.
101type listModelsStub struct {
102	expectedProject, expectedDataset string
103	models                           []*bq.Model
104}
105
106func (s *listModelsStub) listModels(it *ModelIterator, pageSize int, pageToken string) (*bq.ListModelsResponse, error) {
107	if it.dataset.ProjectID != s.expectedProject {
108		return nil, errors.New("wrong project id")
109	}
110	if it.dataset.DatasetID != s.expectedDataset {
111		return nil, errors.New("wrong dataset id")
112	}
113	const maxPageSize = 2
114	if pageSize <= 0 || pageSize > maxPageSize {
115		pageSize = maxPageSize
116	}
117	start := 0
118	if pageToken != "" {
119		var err error
120		start, err = strconv.Atoi(pageToken)
121		if err != nil {
122			return nil, err
123		}
124	}
125	end := start + pageSize
126	if end > len(s.models) {
127		end = len(s.models)
128	}
129	nextPageToken := ""
130	if end < len(s.models) {
131		nextPageToken = strconv.Itoa(end)
132	}
133	return &bq.ListModelsResponse{
134		Models:        s.models[start:end],
135		NextPageToken: nextPageToken,
136	}, nil
137}
138
139func TestModels(t *testing.T) {
140	c := &Client{projectID: "p1"}
141	inModels := []*bq.Model{
142		{ModelReference: &bq.ModelReference{ProjectId: "p1", DatasetId: "d1", ModelId: "m1"}},
143		{ModelReference: &bq.ModelReference{ProjectId: "p1", DatasetId: "d1", ModelId: "m2"}},
144		{ModelReference: &bq.ModelReference{ProjectId: "p1", DatasetId: "d1", ModelId: "m3"}},
145	}
146	outModels := []*Model{
147		{ProjectID: "p1", DatasetID: "d1", ModelID: "m1", c: c},
148		{ProjectID: "p1", DatasetID: "d1", ModelID: "m2", c: c},
149		{ProjectID: "p1", DatasetID: "d1", ModelID: "m3", c: c},
150	}
151
152	lms := &listModelsStub{
153		expectedProject: "p1",
154		expectedDataset: "d1",
155		models:          inModels,
156	}
157	old := listModels
158	listModels = lms.listModels // cannot use t.Parallel with this test
159	defer func() { listModels = old }()
160
161	msg, ok := itest.TestIterator(outModels,
162		func() interface{} { return c.Dataset("d1").Models(context.Background()) },
163		func(it interface{}) (interface{}, error) { return it.(*ModelIterator).Next() })
164	if !ok {
165		t.Error(msg)
166	}
167}
168
169// listRoutinesStub services list requests by returning data from an in-memory list of values.
170type listRoutinesStub struct {
171	routines []*bq.Routine
172}
173
174func (s *listRoutinesStub) listRoutines(it *RoutineIterator, pageSize int, pageToken string) (*bq.ListRoutinesResponse, error) {
175	const maxPageSize = 2
176	if pageSize <= 0 || pageSize > maxPageSize {
177		pageSize = maxPageSize
178	}
179	start := 0
180	if pageToken != "" {
181		var err error
182		start, err = strconv.Atoi(pageToken)
183		if err != nil {
184			return nil, err
185		}
186	}
187	end := start + pageSize
188	if end > len(s.routines) {
189		end = len(s.routines)
190	}
191	nextPageToken := ""
192	if end < len(s.routines) {
193		nextPageToken = strconv.Itoa(end)
194	}
195	return &bq.ListRoutinesResponse{
196		Routines:      s.routines[start:end],
197		NextPageToken: nextPageToken,
198	}, nil
199}
200
201func TestRoutines(t *testing.T) {
202	c := &Client{projectID: "p1"}
203	inRoutines := []*bq.Routine{
204		{RoutineReference: &bq.RoutineReference{ProjectId: "p1", DatasetId: "d1", RoutineId: "r1"}},
205		{RoutineReference: &bq.RoutineReference{ProjectId: "p1", DatasetId: "d1", RoutineId: "r2"}},
206		{RoutineReference: &bq.RoutineReference{ProjectId: "p1", DatasetId: "d1", RoutineId: "r3"}},
207	}
208	outRoutines := []*Routine{
209		{ProjectID: "p1", DatasetID: "d1", RoutineID: "r1", c: c},
210		{ProjectID: "p1", DatasetID: "d1", RoutineID: "r2", c: c},
211		{ProjectID: "p1", DatasetID: "d1", RoutineID: "r3", c: c},
212	}
213
214	lms := &listRoutinesStub{
215		routines: inRoutines,
216	}
217	old := listRoutines
218	listRoutines = lms.listRoutines // cannot use t.Parallel with this test
219	defer func() { listRoutines = old }()
220
221	msg, ok := itest.TestIterator(outRoutines,
222		func() interface{} { return c.Dataset("d1").Routines(context.Background()) },
223		func(it interface{}) (interface{}, error) { return it.(*RoutineIterator).Next() })
224	if !ok {
225		t.Error(msg)
226	}
227}
228
229type listDatasetsStub struct {
230	expectedProject string
231	datasets        []*bq.DatasetListDatasets
232	hidden          map[*bq.DatasetListDatasets]bool
233}
234
235func (s *listDatasetsStub) listDatasets(it *DatasetIterator, pageSize int, pageToken string) (*bq.DatasetList, error) {
236	const maxPageSize = 2
237	if pageSize <= 0 || pageSize > maxPageSize {
238		pageSize = maxPageSize
239	}
240	if it.Filter != "" {
241		return nil, errors.New("filter not supported")
242	}
243	if it.ProjectID != s.expectedProject {
244		return nil, errors.New("bad project ID")
245	}
246	start := 0
247	if pageToken != "" {
248		var err error
249		start, err = strconv.Atoi(pageToken)
250		if err != nil {
251			return nil, err
252		}
253	}
254	var (
255		i             int
256		result        []*bq.DatasetListDatasets
257		nextPageToken string
258	)
259	for i = start; len(result) < pageSize && i < len(s.datasets); i++ {
260		if s.hidden[s.datasets[i]] && !it.ListHidden {
261			continue
262		}
263		result = append(result, s.datasets[i])
264	}
265	if i < len(s.datasets) {
266		nextPageToken = strconv.Itoa(i)
267	}
268	return &bq.DatasetList{
269		Datasets:      result,
270		NextPageToken: nextPageToken,
271	}, nil
272}
273
274func TestDatasets(t *testing.T) {
275	client := &Client{projectID: "p"}
276	inDatasets := []*bq.DatasetListDatasets{
277		{DatasetReference: &bq.DatasetReference{ProjectId: "p", DatasetId: "a"}},
278		{DatasetReference: &bq.DatasetReference{ProjectId: "p", DatasetId: "b"}},
279		{DatasetReference: &bq.DatasetReference{ProjectId: "p", DatasetId: "hidden"}},
280		{DatasetReference: &bq.DatasetReference{ProjectId: "p", DatasetId: "c"}},
281	}
282	outDatasets := []*Dataset{
283		{"p", "a", client},
284		{"p", "b", client},
285		{"p", "hidden", client},
286		{"p", "c", client},
287	}
288	lds := &listDatasetsStub{
289		expectedProject: "p",
290		datasets:        inDatasets,
291		hidden:          map[*bq.DatasetListDatasets]bool{inDatasets[2]: true},
292	}
293	old := listDatasets
294	listDatasets = lds.listDatasets // cannot use t.Parallel with this test
295	defer func() { listDatasets = old }()
296
297	msg, ok := itest.TestIterator(outDatasets,
298		func() interface{} { it := client.Datasets(context.Background()); it.ListHidden = true; return it },
299		func(it interface{}) (interface{}, error) { return it.(*DatasetIterator).Next() })
300	if !ok {
301		t.Fatalf("ListHidden=true: %s", msg)
302	}
303
304	msg, ok = itest.TestIterator([]*Dataset{outDatasets[0], outDatasets[1], outDatasets[3]},
305		func() interface{} { it := client.Datasets(context.Background()); it.ListHidden = false; return it },
306		func(it interface{}) (interface{}, error) { return it.(*DatasetIterator).Next() })
307	if !ok {
308		t.Fatalf("ListHidden=false: %s", msg)
309	}
310}
311
312func TestDatasetToBQ(t *testing.T) {
313	for _, test := range []struct {
314		in   *DatasetMetadata
315		want *bq.Dataset
316	}{
317		{nil, &bq.Dataset{}},
318		{&DatasetMetadata{Name: "name"}, &bq.Dataset{FriendlyName: "name"}},
319		{&DatasetMetadata{
320			Name:                   "name",
321			Description:            "desc",
322			DefaultTableExpiration: time.Hour,
323			DefaultEncryptionConfig: &EncryptionConfig{
324				KMSKeyName: "some_key",
325			},
326			Location: "EU",
327			Labels:   map[string]string{"x": "y"},
328			Access:   []*AccessEntry{{Role: OwnerRole, Entity: "example.com", EntityType: DomainEntity}},
329		}, &bq.Dataset{
330			FriendlyName:             "name",
331			Description:              "desc",
332			DefaultTableExpirationMs: 60 * 60 * 1000,
333			DefaultEncryptionConfiguration: &bq.EncryptionConfiguration{
334				KmsKeyName: "some_key",
335			},
336			Location: "EU",
337			Labels:   map[string]string{"x": "y"},
338			Access:   []*bq.DatasetAccess{{Role: "OWNER", Domain: "example.com"}},
339		}},
340	} {
341		got, err := test.in.toBQ()
342		if err != nil {
343			t.Fatal(err)
344		}
345		if !testutil.Equal(got, test.want) {
346			t.Errorf("%v:\ngot  %+v\nwant %+v", test.in, got, test.want)
347		}
348	}
349
350	// Check that non-writeable fields are unset.
351	aTime := time.Date(2017, 1, 26, 0, 0, 0, 0, time.Local)
352	for _, dm := range []*DatasetMetadata{
353		{CreationTime: aTime},
354		{LastModifiedTime: aTime},
355		{FullID: "x"},
356		{ETag: "e"},
357	} {
358		if _, err := dm.toBQ(); err == nil {
359			t.Errorf("%+v: got nil, want error", dm)
360		}
361	}
362}
363
364func TestBQToDatasetMetadata(t *testing.T) {
365	cTime := time.Date(2017, 1, 26, 0, 0, 0, 0, time.Local)
366	cMillis := cTime.UnixNano() / 1e6
367	mTime := time.Date(2017, 10, 31, 0, 0, 0, 0, time.Local)
368	mMillis := mTime.UnixNano() / 1e6
369	q := &bq.Dataset{
370		CreationTime:             cMillis,
371		LastModifiedTime:         mMillis,
372		FriendlyName:             "name",
373		Description:              "desc",
374		DefaultTableExpirationMs: 60 * 60 * 1000,
375		DefaultEncryptionConfiguration: &bq.EncryptionConfiguration{
376			KmsKeyName: "some_key",
377		},
378		Location: "EU",
379		Labels:   map[string]string{"x": "y"},
380		Access: []*bq.DatasetAccess{
381			{Role: "READER", UserByEmail: "joe@example.com"},
382			{Role: "WRITER", GroupByEmail: "users@example.com"},
383		},
384		Etag: "etag",
385	}
386	want := &DatasetMetadata{
387		CreationTime:           cTime,
388		LastModifiedTime:       mTime,
389		Name:                   "name",
390		Description:            "desc",
391		DefaultTableExpiration: time.Hour,
392		DefaultEncryptionConfig: &EncryptionConfig{
393			KMSKeyName: "some_key",
394		},
395		Location: "EU",
396		Labels:   map[string]string{"x": "y"},
397		Access: []*AccessEntry{
398			{Role: ReaderRole, Entity: "joe@example.com", EntityType: UserEmailEntity},
399			{Role: WriterRole, Entity: "users@example.com", EntityType: GroupEmailEntity},
400		},
401		ETag: "etag",
402	}
403	got, err := bqToDatasetMetadata(q)
404	if err != nil {
405		t.Fatal(err)
406	}
407	if diff := testutil.Diff(got, want); diff != "" {
408		t.Errorf("-got, +want:\n%s", diff)
409	}
410}
411
412func TestDatasetMetadataToUpdateToBQ(t *testing.T) {
413	dm := DatasetMetadataToUpdate{
414		Description:            "desc",
415		Name:                   "name",
416		DefaultTableExpiration: time.Hour,
417		DefaultEncryptionConfig: &EncryptionConfig{
418			KMSKeyName: "some_key",
419		},
420	}
421	dm.SetLabel("label", "value")
422	dm.DeleteLabel("del")
423
424	got, err := dm.toBQ()
425	if err != nil {
426		t.Fatal(err)
427	}
428	want := &bq.Dataset{
429		Description:              "desc",
430		FriendlyName:             "name",
431		DefaultTableExpirationMs: 60 * 60 * 1000,
432		DefaultEncryptionConfiguration: &bq.EncryptionConfiguration{
433			KmsKeyName:      "some_key",
434			ForceSendFields: []string{"KmsKeyName"},
435		},
436		Labels:          map[string]string{"label": "value"},
437		ForceSendFields: []string{"Description", "FriendlyName"},
438		NullFields:      []string{"Labels.del"},
439	}
440	if diff := testutil.Diff(got, want); diff != "" {
441		t.Errorf("-got, +want:\n%s", diff)
442	}
443}
444
445func TestConvertAccessEntry(t *testing.T) {
446	c := &Client{projectID: "pid"}
447	for _, e := range []*AccessEntry{
448		{Role: ReaderRole, Entity: "e", EntityType: DomainEntity},
449		{Role: WriterRole, Entity: "e", EntityType: GroupEmailEntity},
450		{Role: OwnerRole, Entity: "e", EntityType: UserEmailEntity},
451		{Role: ReaderRole, Entity: "e", EntityType: SpecialGroupEntity},
452		{Role: ReaderRole, Entity: "e", EntityType: IAMMemberEntity},
453		{Role: ReaderRole, EntityType: ViewEntity,
454			View: &Table{ProjectID: "p", DatasetID: "d", TableID: "t", c: c}},
455	} {
456		q, err := e.toBQ()
457		if err != nil {
458			t.Fatal(err)
459		}
460		got, err := bqToAccessEntry(q, c)
461		if err != nil {
462			t.Fatal(err)
463		}
464		if diff := testutil.Diff(got, e, cmp.AllowUnexported(Table{}, Client{})); diff != "" {
465			t.Errorf("got=-, want=+:\n%s", diff)
466		}
467	}
468
469	e := &AccessEntry{Role: ReaderRole, Entity: "e"}
470	if _, err := e.toBQ(); err == nil {
471		t.Error("got nil, want error")
472	}
473	if _, err := bqToAccessEntry(&bq.DatasetAccess{Role: "WRITER"}, nil); err == nil {
474		t.Error("got nil, want error")
475	}
476}
477