Просмотр исходного кода

Add field `queryProjectID` to `BigQueryConfiguration` (#3821)

Signed-off-by: Varshith <kvarshithgowda@gmail.com>
Co-authored-by: Alex Meijer <ameijer@users.noreply.github.com>
Varshith 1 неделя назад
Родитель
Сommit
ca98c342b5
2 измененных файлов с 91 добавлено и 8 удалено
  1. 33 6
      pkg/cloud/gcp/bigqueryconfiguration.go
  2. 58 2
      pkg/cloud/gcp/bigqueryconfiguration_test.go

+ 33 - 6
pkg/cloud/gcp/bigqueryconfiguration.go

@@ -17,6 +17,7 @@ type BigQueryConfiguration struct {
 	Table                string     `json:"table"`
 	ExcludePartitionTime bool       `json:"excludePartitionTime"`
 	Location             string     `json:"location"`
+	QueryProjectID       string     `json:"queryProjectID"`
 	Authorizer           Authorizer `json:"authorizer"`
 }
 
@@ -81,16 +82,29 @@ func (bqc *BigQueryConfiguration) Equals(config cloud.Config) bool {
 		return false
 	}
 
+	bqcEffective := bqc.QueryProjectID
+	if bqcEffective == "" {
+		bqcEffective = bqc.ProjectID
+	}
+	thatEffective := thatConfig.QueryProjectID
+	if thatEffective == "" {
+		thatEffective = thatConfig.ProjectID
+	}
+	if bqcEffective != thatEffective {
+		return false
+	}
+
 	return true
 }
 
 func (bqc *BigQueryConfiguration) Sanitize() cloud.Config {
 	return &BigQueryConfiguration{
-		ProjectID:  bqc.ProjectID,
-		Dataset:    bqc.Dataset,
-		Table:      bqc.Table,
-		Location:   bqc.Location,
-		Authorizer: bqc.Authorizer.Sanitize().(Authorizer),
+		ProjectID:      bqc.ProjectID,
+		Dataset:        bqc.Dataset,
+		Table:          bqc.Table,
+		Location:       bqc.Location,
+		QueryProjectID: bqc.QueryProjectID,
+		Authorizer:     bqc.Authorizer.Sanitize().(Authorizer),
 	}
 }
 
@@ -113,7 +127,12 @@ func (bqc *BigQueryConfiguration) GetBigQueryClient(ctx context.Context) (*bigqu
 		return nil, err
 	}
 
-	client, err := bigquery.NewClient(ctx, bqc.ProjectID, clientOpts...)
+	queryProjectID := bqc.QueryProjectID
+	if queryProjectID == "" {
+		queryProjectID = bqc.ProjectID
+	}
+
+	client, err := bigquery.NewClient(ctx, queryProjectID, clientOpts...)
 	if err != nil {
 		return nil, err
 	}
@@ -159,6 +178,14 @@ func (bqc *BigQueryConfiguration) UnmarshalJSON(b []byte) error {
 		bqc.Location = location
 	}
 
+	if _, ok := fmap["queryProjectID"]; ok {
+		queryProjectID, err := cloud.GetInterfaceValue[string](fmap, "queryProjectID")
+		if err != nil {
+			return fmt.Errorf("BigQueryConfiguration: UnmarshalJSON: %w", err)
+		}
+		bqc.QueryProjectID = queryProjectID
+	}
+
 	authAny, ok := fmap["authorizer"]
 	if !ok {
 		return fmt.Errorf("StorageConfiguration: UnmarshalJSON: missing authorizer")

+ 58 - 2
pkg/cloud/gcp/bigqueryconfiguration_test.go

@@ -138,7 +138,8 @@ func TestBigQueryConfiguration_Equals(t *testing.T) {
 						"key1": "key2",
 					},
 				},
-				Location: "EU",
+				Location:       "EU",
+				QueryProjectID: "queryProjectID",
 			},
 			right: &BigQueryConfiguration{
 				ProjectID: "projectID",
@@ -150,7 +151,8 @@ func TestBigQueryConfiguration_Equals(t *testing.T) {
 						"key1": "key2",
 					},
 				},
-				Location: "EU",
+				Location:       "EU",
+				QueryProjectID: "queryProjectID",
 			},
 			expected: true,
 		},
@@ -347,6 +349,60 @@ func TestBigQueryConfiguration_Equals(t *testing.T) {
 			},
 			expected: false,
 		},
+		"different queryProjectID": {
+			left: BigQueryConfiguration{
+				ProjectID:      "projectID",
+				Dataset:        "dataset",
+				Table:          "table",
+				QueryProjectID: "queryProjectID",
+				Authorizer: &ServiceAccountKey{
+					Key: map[string]string{
+						"Key":  "Key",
+						"key1": "key2",
+					},
+				},
+			},
+			right: &BigQueryConfiguration{
+				ProjectID:      "projectID",
+				Dataset:        "dataset",
+				Table:          "table",
+				QueryProjectID: "queryProjectID2",
+				Authorizer: &ServiceAccountKey{
+					Key: map[string]string{
+						"Key":  "Key",
+						"key1": "key2",
+					},
+				},
+			},
+			expected: false,
+		},
+		"empty queryProjectID equals projectID": {
+			left: BigQueryConfiguration{
+				ProjectID:      "projectID",
+				Dataset:        "dataset",
+				Table:          "table",
+				QueryProjectID: "",
+				Authorizer: &ServiceAccountKey{
+					Key: map[string]string{
+						"Key":  "Key",
+						"key1": "key2",
+					},
+				},
+			},
+			right: &BigQueryConfiguration{
+				ProjectID:      "projectID",
+				Dataset:        "dataset",
+				Table:          "table",
+				QueryProjectID: "projectID",
+				Authorizer: &ServiceAccountKey{
+					Key: map[string]string{
+						"Key":  "Key",
+						"key1": "key2",
+					},
+				},
+			},
+			expected: true,
+		},
 	}
 
 	for name, testCase := range testCases {