Skip to content

Commit

Permalink
Merge pull request #9 from Recidiviz/dan/candidate-json-source
Browse files Browse the repository at this point in the history 
Add `--data-from-json` option for initialization of emulator data
  • Loading branch information
@ohaibbq
ohaibbq committed Apr 11, 2024
2 parents f593ce6 + dd7aa97 commit 70d2848
Show file tree
Hide file tree
Showing 5 changed files with 237 additions and 21 deletions.
7 changes: 7 additions & 0 deletions cmd/bigquery-emulator/main.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ type option struct {
LogFormat server.LogFormat `description:"specify the log format (console/json)" long:"log-format" default:"console"`
Database string `description:"specify the database file, use :memory: for in-memory storage. if not specified, it will be a temp file" long:"database"`
DataFromYAML string `description:"specify the path to the YAML file that contains the initial data" long:"data-from-yaml"`
DataFromJSON string `description:"specify the path to the JSON file that contains the initial data (faster for large, multi-megabyte files)" long:"data-from-json"`
Version bool `description:"print version" long:"version" short:"v"`
}

Expand Down Expand Up @@ -111,6 +112,12 @@ func runServer(args []string, opt option) error {
}
}

if opt.DataFromJSON != "" {
if err := bqServer.Load(server.JSONSource(opt.DataFromJSON)); err != nil {
return err
}
}

ctx := context.Background()
interrupt := make(chan os.Signal, 1)
signal.Notify(interrupt, os.Interrupt, syscall.SIGTERM)
Expand Down
60 changes: 60 additions & 0 deletions server/server_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -211,6 +211,66 @@ func findDatasets(t *testing.T, ctx context.Context, client *bigquery.Client) []
return datasets
}

func TestDataFromJSON(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)
}
if err := bqServer.Load(server.JSONSource(filepath.Join("testdata", "data.json"))); 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()

query := client.Query("SELECT * FROM dataset1.table_a")
job, err := query.Run(ctx)
if err != nil {
t.Fatal(err)
}
if _, err := job.Config(); err != nil {
t.Fatal(err)
}
if _, err := job.Wait(ctx); err != nil {
t.Fatal(err)
}

gotJob, err := client.JobFromID(ctx, job.ID())
if err != nil {
t.Fatal(err)
}
if gotJob.ID() != job.ID() {
t.Fatalf("failed to get job expected ID %s. but got %s", job.ID(), gotJob.ID())
}

if jobs := findJobs(t, ctx, client); len(jobs) != 1 {
t.Fatalf("failed to find jobs. expected 1 jobs but found %d jobs", len(jobs))
}
}

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

Expand Down
29 changes: 29 additions & 0 deletions server/source.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@ import (
"bytes"
"context"
"errors"
"github.com/goccy/go-json"
"io"
"os"

"github.com/go-playground/validator/v10"
Expand Down Expand Up @@ -36,6 +38,33 @@ func YAMLSource(path string) Source {
}
}

func JSONSource(path string) Source {
return func(s *Server) error {
jsonFile, err := os.Open(path)
if err != nil {
return err
}

content, err := io.ReadAll(jsonFile)
if err != nil {
return err
}

err = jsonFile.Close()
if err != nil {
return err
}

var v struct {
Projects []*types.Project `json:"projects"`
}
if err := json.Unmarshal([]byte(content), &v); err != nil {
return err
}
return s.addProjects(context.Background(), v.Projects)
}
}

func StructSource(projects ...*types.Project) Source {
return func(s *Server) error {
return s.addProjects(context.Background(), projects)
Expand Down
108 changes: 108 additions & 0 deletions server/testdata/data.json
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,108 @@
{
"projects": [
{
"id": "test",
"datasets": [
{
"id": "dataset1",
"tables": [
{
"id": "table_a",
"columns": [
{
"name": "id",
"type": "INTEGER",
"mode": "REQUIRED"
},
{
"name": "name",
"type": "STRING",
"mode": "required"
},
{
"name": "structarr",
"type": "STRUCT",
"mode": "repeated",
"fields": [
{
"name": "key",
"type": "STRING"
},
{
"name": "value",
"type": "JSON"
}
]
},
{
"name": "birthday",
"type": "DATE"
},
{
"name": "skillNum",
"type": "NUMERIC"
},
{
"name": "created_at",
"type": "TIMESTAMP"
}
],
"data": [
{
"id": 1,
"name": "alice",
"structarr": [
{
"key": "profile",
"value": "{\"age\": 10}"
}
],
"birthday": "2012-01-01",
"skillNum": 3,
"created_at": "2022-01-01T12:00:00"
},
{
"id": 2,
"name": "bob",
"structarr": [
{
"key": "profile",
"value": "{\"age\": 15}"
}
],
"birthday": "2007-02-01",
"skillNum": 5,
"created_at": "2022-01-02T18:00:00"
}
]
},
{
"id": "table_b",
"columns": [
{
"name": "num",
"type": "NUMERIC"
},
{
"name": "bignum",
"type": "BIGNUMERIC"
},
{
"name": "interval",
"type": "INTERVAL"
}
],
"data": [
{
"num": 1.2345,
"bignum": 1.234567891234,
"interval": "1-6 15 0:0:0"
}
]
}
]
}
]
}
]
}
54 changes: 33 additions & 21 deletions types/types.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,23 +13,23 @@ import (
)

type Project struct {
ID string `yaml:"id" validate:"required"`
Datasets []*Dataset `yaml:"datasets" validate:"required"`
Jobs []*Job `yaml:"jobs"`
ID string `json:"id" yaml:"id" validate:"required"`
Datasets []*Dataset `json:"datasets" yaml:"datasets" validate:"required"`
Jobs []*Job `json:"jobs" yaml:"jobs"`
}

type Dataset struct {
ID string `yaml:"id" validate:"required"`
Tables []*Table `yaml:"tables"`
Models []*Model `yaml:"models"`
Routines []*Routine `yaml:"routines"`
ID string `json:"id" yaml:"id" validate:"required"`
Tables []*Table `json:"tables" yaml:"tables"`
Models []*Model `json:"models" yaml:"models"`
Routines []*Routine `json:"routines" yaml:"routines"`
}

type Table struct {
ID string `yaml:"id" validate:"required"`
Columns []*Column `yaml:"columns" validate:"required"`
Data Data `yaml:"data"`
Metadata map[string]interface{} `yaml:"metadata"`
ID string `json:"id" yaml:"id" validate:"required"`
Columns []*Column `json:"columns" yaml:"columns" validate:"required"`
Data Data `json:"data" yaml:"data"`
Metadata map[string]interface{} `json:"metadata" yaml:"metadata"`
}

func (t *Table) ToBigqueryV2(projectID, datasetID string) *bigqueryv2.Table {
Expand Down Expand Up @@ -79,17 +79,29 @@ func (m *Mode) UnmarshalYAML(b []byte) error {
return nil
}

func (m *Mode) UnmarshalJSON(b []byte) error {
switch strings.ToLower(strings.Trim(string(b), `"`)) {
case strings.ToLower(string(NullableMode)):
*m = NullableMode
case strings.ToLower(string(RequiredMode)):
*m = RequiredMode
case strings.ToLower(string(RepeatedMode)):
*m = RepeatedMode
}
return nil
}

const (
NullableMode Mode = "NULLABLE"
RequiredMode Mode = "REQUIRED"
RepeatedMode Mode = "REPEATED"
)

type Column struct {
Name string `yaml:"name" validate:"required"`
Type Type `yaml:"type" validate:"type"`
Mode Mode `yaml:"mode" validate:"mode"`
Fields []*Column `yaml:"fields"`
Name string `json:"name" yaml:"name" validate:"required"`
Type Type `json:"type" yaml:"type" validate:"type"`
Mode Mode `json:"mode" yaml:"mode" validate:"mode"`
Fields []*Column `json:"fields" yaml:"fields"`
}

func (c *Column) FormatType() string {
Expand Down Expand Up @@ -136,18 +148,18 @@ func tableFieldSchemaFromColumn(c *Column) *bigqueryv2.TableFieldSchema {
}

type Job struct {
ID string `yaml:"id" validate:"required"`
Metadata map[string]interface{} `yaml:"metadata"`
ID string `json:"id" yaml:"id" validate:"required"`
Metadata map[string]interface{} `json:"metadata" yaml:"metadata"`
}

type Model struct {
ID string `yaml:"id" validate:"required"`
Metadata map[string]interface{} `yaml:"metadata"`
ID string `json:"id" yaml:"id" validate:"required"`
Metadata map[string]interface{} `json:"metadata" yaml:"metadata"`
}

type Routine struct {
ID string `yaml:"id" validate:"required"`
Metadata map[string]interface{} `yaml:"metadata"`
ID string `json:"id" yaml:"id" validate:"required"`
Metadata map[string]interface{} `json:"metadata" yaml:"metadata"`
}

type Type string
Expand Down

0 comments on commit 70d2848

Please sign in to comment.