Explorar o código

Add field `location` to `BigQueryConfiguration` (#3795)

Signed-off-by: Yannik Daellenbach <git@daellenbach.org>
Yannik Dällenbach hai 4 horas
pai
achega
7a246b6bc6

+ 23 - 1
pkg/cloud/gcp/bigqueryconfiguration.go

@@ -16,6 +16,7 @@ type BigQueryConfiguration struct {
 	Dataset              string     `json:"dataset"`
 	Table                string     `json:"table"`
 	ExcludePartitionTime bool       `json:"excludePartitionTime"`
+	Location             string     `json:"location"`
 	Authorizer           Authorizer `json:"authorizer"`
 }
 
@@ -76,6 +77,10 @@ func (bqc *BigQueryConfiguration) Equals(config cloud.Config) bool {
 		return false
 	}
 
+	if bqc.Location != thatConfig.Location {
+		return false
+	}
+
 	return true
 }
 
@@ -84,6 +89,7 @@ func (bqc *BigQueryConfiguration) Sanitize() cloud.Config {
 		ProjectID:  bqc.ProjectID,
 		Dataset:    bqc.Dataset,
 		Table:      bqc.Table,
+		Location:   bqc.Location,
 		Authorizer: bqc.Authorizer.Sanitize().(Authorizer),
 	}
 }
@@ -106,7 +112,15 @@ func (bqc *BigQueryConfiguration) GetBigQueryClient(ctx context.Context) (*bigqu
 	if err != nil {
 		return nil, err
 	}
-	return bigquery.NewClient(ctx, bqc.ProjectID, clientOpts...)
+
+	client, err := bigquery.NewClient(ctx, bqc.ProjectID, clientOpts...)
+	if err != nil {
+		return nil, err
+	}
+
+	client.Location = bqc.Location
+
+	return client, nil
 }
 
 // UnmarshalJSON assumes data is save as an BigQueryConfigurationDTO
@@ -137,6 +151,14 @@ func (bqc *BigQueryConfiguration) UnmarshalJSON(b []byte) error {
 	}
 	bqc.Table = table
 
+	if _, ok := fmap["location"]; ok {
+		location, err := cloud.GetInterfaceValue[string](fmap, "location")
+		if err != nil {
+			return fmt.Errorf("BigQueryConfiguration: FromInterface: %s", err.Error())
+		}
+		bqc.Location = location
+	}
+
 	authAny, ok := fmap["authorizer"]
 	if !ok {
 		return fmt.Errorf("StorageConfiguration: UnmarshalJSON: missing authorizer")

+ 29 - 0
pkg/cloud/gcp/bigqueryconfiguration_test.go

@@ -138,6 +138,7 @@ func TestBigQueryConfiguration_Equals(t *testing.T) {
 						"key1": "key2",
 					},
 				},
+				Location: "EU",
 			},
 			right: &BigQueryConfiguration{
 				ProjectID: "projectID",
@@ -149,6 +150,7 @@ func TestBigQueryConfiguration_Equals(t *testing.T) {
 						"key1": "key2",
 					},
 				},
+				Location: "EU",
 			},
 			expected: true,
 		},
@@ -318,6 +320,33 @@ func TestBigQueryConfiguration_Equals(t *testing.T) {
 			},
 			expected: false,
 		},
+		"different location": {
+			left: BigQueryConfiguration{
+				ProjectID: "projectID",
+				Dataset:   "dataset",
+				Table:     "table",
+				Location:  "EU",
+				Authorizer: &ServiceAccountKey{
+					Key: map[string]string{
+						"Key":  "Key",
+						"key1": "key2",
+					},
+				},
+			},
+			right: &BigQueryConfiguration{
+				ProjectID: "projectID",
+				Dataset:   "dataset2",
+				Table:     "table",
+				Location:  "US",
+				Authorizer: &ServiceAccountKey{
+					Key: map[string]string{
+						"Key":  "Key",
+						"key1": "key2",
+					},
+				},
+			},
+			expected: false,
+		},
 	}
 
 	for name, testCase := range testCases {