test: select * from dataset1.user

myuon · Mar 10, 2024 · 1d805a5 · 1d805a5
1 parent 45db356
commit 1d805a5
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 63 deletions.
diff --git a/gallon/output_bigquery.go b/gallon/output_bigquery.go
@@ -25,6 +25,7 @@ import (
 type OutputPluginBigQuery struct {
 	logger               logr.Logger
 	client               *bigquery.Client
+	endpoint             *string
 	datasetId            string
 	tableId              string
 	schema               bigquery.Schema
@@ -34,6 +35,7 @@ type OutputPluginBigQuery struct {
 
 func NewOutputPluginBigQuery(
 	client *bigquery.Client,
+	endpoint *string,
 	datasetId string,
 	tableId string,
 	schema bigquery.Schema,
@@ -42,6 +44,7 @@ func NewOutputPluginBigQuery(
 ) *OutputPluginBigQuery {
 	return &OutputPluginBigQuery{
 		client:               client,
+		endpoint:             endpoint,
 		datasetId:            datasetId,
 		tableId:              tableId,
 		schema:               schema,
@@ -320,6 +323,7 @@ func NewOutputPluginBigQueryFromConfig(configYml []byte) (*OutputPluginBigQuery,
 
 	return NewOutputPluginBigQuery(
 		client,
+		config.Endpoint,
 		config.DatasetId,
 		config.TableId,
 		schema,

diff --git a/test/bigquery/bigquery_test.go b/test/bigquery/bigquery_test.go
@@ -3,15 +3,13 @@ package bigquery
 import (
 	"cloud.google.com/go/bigquery"
 	"context"
+	"errors"
 	"fmt"
-	"github.com/aws/aws-sdk-go-v2/aws"
-	"github.com/aws/aws-sdk-go-v2/service/dynamodb"
-	"github.com/aws/aws-sdk-go-v2/service/dynamodb/types"
-	"github.com/brianvoe/gofakeit/v6"
 	"github.com/myuon/gallon/cmd"
 	"github.com/ory/dockertest/v3"
 	"github.com/ory/dockertest/v3/docker"
 	"go.uber.org/zap"
+	"google.golang.org/api/iterator"
 	"google.golang.org/api/option"
 	"log"
 	"os"
@@ -35,62 +33,6 @@ type UserTable struct {
 var client *bigquery.Client
 var endpoint string
 
-func NewFakeUserTable() (UserTable, error) {
-	v := UserTable{}
-	if err := gofakeit.Struct(&v); err != nil {
-		return v, err
-	}
-
-	return v, nil
-}
-
-func Migrate(client *dynamodb.Client) error {
-	ctx := context.Background()
-	if _, err := client.CreateTable(ctx, &dynamodb.CreateTableInput{
-		TableName: aws.String("users"),
-		AttributeDefinitions: []types.AttributeDefinition{
-			{
-				AttributeName: aws.String("id"),
-				AttributeType: types.ScalarAttributeTypeS,
-			},
-		},
-		KeySchema: []types.KeySchemaElement{
-			{
-				AttributeName: aws.String("id"),
-				KeyType:       types.KeyTypeHash,
-			},
-		},
-		BillingMode: types.BillingModePayPerRequest,
-	}); err != nil {
-		return err
-	}
-
-	for i := 0; i < 1000; i++ {
-		v, err := NewFakeUserTable()
-		if err != nil {
-			return err
-		}
-
-		record := map[string]types.AttributeValue{
-			"id":         &types.AttributeValueMemberS{Value: v.ID},
-			"name":       &types.AttributeValueMemberS{Value: v.Name},
-			"age":        &types.AttributeValueMemberN{Value: fmt.Sprintf("%v", v.Age)},
-			"created_at": &types.AttributeValueMemberN{Value: fmt.Sprintf("%v", v.CreatedAt)},
-		}
-
-		if _, err := client.PutItem(ctx, &dynamodb.PutItemInput{
-			TableName: aws.String("users"),
-			Item:      record,
-		}); err != nil {
-			return err
-		}
-	}
-
-	log.Printf("Migrated %v rows", 1000)
-
-	return nil
-}
-
 func TestMain(m *testing.M) {
 	var exitCode int
 	defer func() {
@@ -138,7 +80,10 @@ func TestMain(m *testing.M) {
 	port := resource.GetPort("9050/tcp")
 	endpoint = fmt.Sprintf("http://localhost:%v", port)
 
-	client, _ := bigquery.NewClient(context.Background(), "test", option.WithEndpoint(endpoint))
+	client, err = bigquery.NewClient(context.Background(), "test", option.WithEndpoint(endpoint))
+	if err != nil {
+		log.Panicf("Could not create client: %v", err)
+	}
 
 	if err := pool.Retry(func() error {
 		log.Println("Trying to connect to database...")
@@ -188,8 +133,7 @@ out:
 		t.Errorf("Could not run command: %s", err)
 	}
 
-	table := client.Dataset("dataset1").Table("user")
-	job, err := client.Query(fmt.Sprintf("SELECT * FROM %v", table.FullyQualifiedName())).Run(context.Background())
+	job, err := client.Query("SELECT * FROM `dataset1.user`").Run(context.Background())
 	if err != nil {
 		t.Errorf("Could not run query: %s", err)
 	}
@@ -206,6 +150,9 @@ out:
 	for {
 		var v UserTable
 		err := it.Next(&v)
+		if errors.Is(err, iterator.Done) {
+			break
+		}
 		if err != nil {
 			t.Errorf("Could not iterate: %s", err)
 		}

diff --git a/test/bigquery/testdata/data.yaml b/test/bigquery/testdata/data.yaml
@@ -2,3 +2,14 @@ projects:
   - id: test
     datasets:
       - id: dataset1
+        tables:
+          - id: user
+            columns:
+              - name: id
+                type: STRING
+              - name: name
+                type: STRING
+              - name: age
+                type: INTEGER
+              - name: created_at
+                type: INTEGER