Skip to content

Commit

Permalink
Properly unmarshal JSON values into golang values when inserting tabl…
Browse files Browse the repository at this point in the history
…e data
  • Loading branch information
ohaibbq committed Jul 9, 2024
1 parent 316038b commit 74df6c3
Show file tree
Hide file tree
Showing 2 changed files with 133 additions and 0 deletions.
11 changes: 11 additions & 0 deletions internal/contentdata/repository.go
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ import (
"reflect"
"strings"

"github.com/goccy/go-json"
"github.com/goccy/go-zetasqlite"
"go.uber.org/zap"
bigqueryv2 "google.golang.org/api/bigquery/v2"
Expand Down Expand Up @@ -400,6 +401,7 @@ func (r *Repository) AddTableData(ctx context.Context, tx *connection.Tx, projec
for _, column := range columns {
if value, found := data[column.Name]; found {
isTimestampColumn := column.Type == types.TIMESTAMP
isJsonColumn := column.Type == types.JSON
inputString, isInputString := value.(string)

if isInputString && isTimestampColumn {
Expand All @@ -411,6 +413,15 @@ func (r *Repository) AddTableData(ctx context.Context, tx *connection.Tx, projec
}
}

if isInputString && isJsonColumn {
var jsonValue interface{}
if err := json.Unmarshal([]byte(inputString), &jsonValue); err != nil {
return fmt.Errorf("failed to unmarshal json value [%s]: %w", inputString, err)
}
values = append(values, jsonValue)
continue
}

values = append(values, value)
} else {
values = append(values, nil)
Expand Down
122 changes: 122 additions & 0 deletions server/server_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import (
"compress/gzip"
"context"
"encoding/csv"
"errors"
"fmt"
"io"
"math/big"
Expand Down Expand Up @@ -284,6 +285,127 @@ func TestJob(t *testing.T) {
}
}

type JSONTableSaver struct {
JsonColumn map[string]bigquery.Value
}

// Save implements the bigquery.ValueSaver interface to serialize the "json_column" field as JSON string.
func (s *JSONTableSaver) Save() (row map[string]bigquery.Value, insertID string, err error) {
attr, err := json.Marshal(s.JsonColumn)
if err != nil {
return nil, "", err
}

row = map[string]bigquery.Value{
"json_column": string(attr),
}

return row, "1", nil
}

func TestJson(t *testing.T) {
ctx := context.Background()

const (
projectName = "test"
)

bqServer, err := server.New(server.TempStorage)
if err != nil {
t.Fatal(err)
}
if err := bqServer.SetProject(projectName); err != nil {
t.Fatal(err)
}

project := types.NewProject("test",
types.NewDataset("dataset",
types.NewTable("table_a", []*types.Column{
{Name: "json_column", Type: "JSON", Mode: "REQUIRED"},
}, nil,
),
),
)

if err := bqServer.Load(server.StructSource(project)); err != nil {
t.Fatal(err)
}

testServer := bqServer.TestServer()
defer func() {
testServer.Close()
bqServer.Stop(ctx)
}()

client, err := bigquery.NewClient(
ctx,
projectName,
option.WithEndpoint(testServer.URL),
option.WithoutAuthentication(),
)
if err != nil {
t.Fatal(err)
}
defer client.Close()

saver := &JSONTableSaver{JsonColumn: map[string]bigquery.Value{
"a_date": "2024-01-01",
"b_float": 1.2,
"c_nested": map[string]bigquery.Value{
"test": "struct value",
},
"d_list": []bigquery.Value{
"list item",
"list item",
},
}}

table := client.Dataset("dataset").Table("table_a")
if err := table.Inserter().Put(ctx, saver); err != nil {
t.Fatal(err)
}

query := client.Query(`SELECT
JSON_VALUE(json_column, "$.a_date") AS a_date,
JSON_VALUE(json_column, "$.b_float") AS b_float,
JSON_VALUE(json_column, "$.c_nested.test") AS c_nested,
JSON_VALUE(json_column, "$.d_list[0]") AS d_list,
FROM dataset.table_a`)
iter, err := query.Read(ctx)
if err != nil {
t.Fatal(err)
}

var rows [][]bigquery.Value
for {
var row []bigquery.Value
if err := iter.Next(&row); err != nil {
if errors.Is(err, iterator.Done) {
break
}
t.Fatal(err)
}
rows = append(rows, row)
}
if len(rows) != 1 {
t.Fatalf("expected 1 rows, got [%d]", len(rows))
}
for _, row := range rows {
if row[0] != "2024-01-01" {
t.Fatalf(`expected ["2024-01-01"] but got [%s]`, row[0])
}
if row[1] != "1.2" {
t.Fatalf(`expected ["1.2"] but got [%s]`, row[1])
}
if row[2] != "struct value" {
t.Fatalf(`expected ["struct value"] but got [%s]`, row[2])
}
if row[3] != "list item" {
t.Fatalf(`expected ["list item"] but got [%s]`, row[2])
}
}
}

func TestFetchData(t *testing.T) {
ctx := context.Background()

Expand Down

0 comments on commit 74df6c3

Please sign in to comment.