bigquery/connection_test.go at main · viant/bigquery · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
package bigquery

import (
	"testing"

	"github.com/stretchr/testify/assert"
)

func TestJobConfiguration_Reservation(t *testing.T) {
	var testCases = []struct {
		description         string
		cfg                 *Config
		query               string
		expectedReservation string
		expectedPriority    string
	}{
		{
			description: "reservation set from config",
			cfg: &Config{
				ProjectID:   "myproject",
				DatasetID:   "mydataset",
				Priority:    PriorityInteractive,
				Reservation: "projects/myproject/locations/us/reservations/myreservation",
			},
			query:               "SELECT 1",
			expectedReservation: "projects/myproject/locations/us/reservations/myreservation",
			expectedPriority:    PriorityInteractive,
		},
		{
			description: "no reservation when not configured",
			cfg: &Config{
				ProjectID: "myproject",
				DatasetID: "mydataset",
				Priority:  PriorityInteractive,
			},
			query:               "SELECT 1",
			expectedReservation: "",
			expectedPriority:    PriorityInteractive,
		},
		{
			description: "reservation with batch priority",
			cfg: &Config{
				ProjectID:   "myproject",
				DatasetID:   "mydataset",
				Priority:    PriorityBatch,
				Reservation: "projects/myproject/locations/us/reservations/prod",
			},
			query:               "SELECT 1",
			expectedReservation: "projects/myproject/locations/us/reservations/prod",
			expectedPriority:    PriorityBatch,
		},
		{
			description: "priority from hint overrides config priority but reservation still set",
			cfg: &Config{
				ProjectID:   "myproject",
				DatasetID:   "mydataset",
				Priority:    PriorityInteractive,
				Reservation: "projects/myproject/locations/us/reservations/myreservation",
			},
			query:               `SELECT /*+ {"Priority": "BATCH"} +*/ 1`,
			expectedReservation: "projects/myproject/locations/us/reservations/myreservation",
			expectedPriority:    PriorityBatch,
		},
	}

	for _, tc := range testCases {
		t.Run(tc.description, func(t *testing.T) {
			conn := &connection{cfg: tc.cfg, projectID: tc.cfg.ProjectID}
			job, err := conn.jobConfiguration(tc.query)
			if !assert.NoError(t, err) {
				return
			}
			assert.Equal(t, tc.expectedReservation, job.Configuration.Reservation)
			assert.Equal(t, tc.expectedPriority, job.Configuration.Query.Priority)
		})
	}
}