Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add --data-from-json option for initialization of emulator data #295

Open
wants to merge 4 commits into
base: main
Choose a base branch
from
Open

Add --data-from-json option for initialization of emulator data #295

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
a75d0ee
Add `--data-from-json` option for initialization of emulator data
ohaibbq Apr 10, 2024
3ef4659
add test
ohaibbq Apr 10, 2024
6d46bfa
add mode unmarshaller
ohaibbq Apr 11, 2024
04ec262
Update main.go
ohaibbq Apr 11, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions cmd/bigquery-emulator/main.go
View file
Open in desktop
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 if required. if not specified, it will be on memory" 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 @@ -109,6 +110,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
Open in desktop
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
Open in desktop
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
Open in desktop
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
Open in desktop
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
Loading