Skip to content

Commit

Permalink
feat(bigquery): RANGE support for basic data movement (#9762)
Browse files Browse the repository at this point in the history 
* feat(bigquery): RANGE type StandardSQLDataType support

This PR augments the StandardSQLDataRepresentation(s) to support
range-specific augmentations, and adds some testing.  Astute observers
will note that this does include mapping changes to param handling, which will
be tested in a subsequent PR that expands RANGE coverage to that area of
the library.

* feat(bigquery): RANGE support (data)

This PR adds a new RangeValue type for conveying parsed range start/end
values, and adds the requisite plumbing for parameter handling (type and
binding) and data result reading.

* more typing work and inference

* round trip

* more testing, cleanup

* more tests

* explicit param type test

* correct supported RANGE element types

* more testing

* add temporary arg to initQueryParameterTestCases to isolate storage work

* reorder error handling
  • Loading branch information
@shollyman
shollyman authored Apr 17, 2024
1 parent a140c7f commit 07f0806
Show file tree
Hide file tree
Showing 9 changed files with 514 additions and 48 deletions.
157 changes: 119 additions & 38 deletions bigquery/integration_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1160,18 +1160,20 @@ type SubTestStruct struct {
}

type TestStruct struct {
Name string
Bytes []byte
Integer int64
Float float64
Boolean bool
Timestamp time.Time
Date civil.Date
Time civil.Time
DateTime civil.DateTime
Numeric *big.Rat
Geography string

Name string
Bytes []byte
Integer int64
Float float64
Boolean bool
Timestamp time.Time
Date civil.Date
Time civil.Time
DateTime civil.DateTime
Numeric *big.Rat
Geography string
RangeDate *RangeValue `bigquery:"rangedate"` //TODO: remove tags when field normalization works
RangeDateTime *RangeValue `bigquery:"rangedatetime"`
RangeTimestamp *RangeValue `bigquery:"rangetimestamp"`
StringArray []string
IntegerArray []int64
FloatArray []float64
Expand Down Expand Up @@ -1200,6 +1202,19 @@ func TestIntegration_InsertAndReadStructs(t *testing.T) {
t.Fatal(err)
}

// Finish declaring the ambigous range element types.
for idx, typ := range map[int]FieldType{
11: DateFieldType,
12: DateTimeFieldType,
13: TimestampFieldType,
} {
if schema[idx].Type != RangeFieldType {
t.Fatalf("mismatch in expected RANGE element in schema field %d", idx)
} else {
schema[idx].RangeElementType = &RangeElementType{Type: typ}
}
}

ctx := context.Background()
table := newTable(t, schema)
defer table.Delete(ctx)
Expand All @@ -1214,6 +1229,15 @@ func TestIntegration_InsertAndReadStructs(t *testing.T) {
dtm2 := civil.DateTime{Date: d2, Time: tm2}
g := "POINT(-122.350220 47.649154)"
g2 := "POINT(-122.0836791 37.421827)"
rangedate := &RangeValue{Start: civil.Date{Year: 2024, Month: 04, Day: 11}}
rangedatetime := &RangeValue{
End: civil.DateTime{
Date: civil.Date{Year: 2024, Month: 04, Day: 11},
Time: civil.Time{Hour: 2, Minute: 4, Second: 6, Nanosecond: 0}},
}
rangetimestamp := &RangeValue{
Start: time.Date(2016, 3, 20, 15, 4, 5, 6000, time.UTC),
}

// Populate the table.
ins := table.Inserter()
Expand All @@ -1230,6 +1254,9 @@ func TestIntegration_InsertAndReadStructs(t *testing.T) {
dtm,
big.NewRat(57, 100),
g,
rangedate,
rangedatetime,
rangetimestamp,
[]string{"a", "b"},
[]int64{1, 2},
[]float64{1, 1.41},
Expand All @@ -1255,16 +1282,19 @@ func TestIntegration_InsertAndReadStructs(t *testing.T) {
},
},
{
Name: "b",
Bytes: []byte("byte2"),
Integer: 24,
Float: 4.13,
Boolean: false,
Timestamp: ts,
Date: d,
Time: tm,
DateTime: dtm,
Numeric: big.NewRat(4499, 10000),
Name: "b",
Bytes: []byte("byte2"),
Integer: 24,
Float: 4.13,
Boolean: false,
Timestamp: ts,
Date: d,
Time: tm,
DateTime: dtm,
Numeric: big.NewRat(4499, 10000),
RangeDate: rangedate,
RangeDateTime: rangedatetime,
RangeTimestamp: rangetimestamp,
},
}
var savers []*StructSaver
Expand Down Expand Up @@ -1866,6 +1896,10 @@ func TestIntegration_StandardQuery(t *testing.T) {
{"ArrayOfStructs", "SELECT [(1, 2, 3), (4, 5, 6)]", []Value{[]Value{ints(1, 2, 3), ints(4, 5, 6)}}},
{"ComplexNested", "SELECT [([1, 2, 3], 4), ([5, 6], 7)]", []Value{[]Value{[]Value{ints(1, 2, 3), int64(4)}, []Value{ints(5, 6), int64(7)}}}},
{"SubSelectArray", "SELECT ARRAY(SELECT STRUCT([1, 2]))", []Value{[]Value{[]Value{ints(1, 2)}}}},
{"RangeOofDateLiteral",
"SELECT RANGE(DATE '2023-03-01', DATE '2024-04-16')",
[]Value{&RangeValue{Start: civil.Date{Year: 2023, Month: 03, Day: 01}, End: civil.Date{Year: 2024, Month: 04, Day: 16}}},
},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
Expand Down Expand Up @@ -2108,17 +2142,19 @@ func TestIntegration_QuerySessionSupport(t *testing.T) {

}

type queryParameterTestCase struct {
name string
query string
parameters []QueryParameter
wantRow []Value
wantConfig interface{}
}

var (
queryParameterTestCases = []struct {
name string
query string
parameters []QueryParameter
wantRow []Value
wantConfig interface{}
}{}
queryParameterTestCases = []queryParameterTestCase{}
)

func initQueryParameterTestCases() {
func initQueryParameterTestCases(includeRangeCases bool) {
d := civil.Date{Year: 2016, Month: 3, Day: 20}
tm := civil.Time{Hour: 15, Minute: 04, Second: 05, Nanosecond: 3008}
rtm := tm
Expand All @@ -2127,6 +2163,12 @@ func initQueryParameterTestCases() {
ts := time.Date(2016, 3, 20, 15, 04, 05, 0, time.UTC)
rat := big.NewRat(13, 10)
bigRat := big.NewRat(12345, 10e10)
rangeTimestamp1 := &RangeValue{
Start: time.Date(2016, 3, 20, 15, 04, 05, 0, time.UTC),
}
rangeTimestamp2 := &RangeValue{
End: time.Date(2016, 3, 20, 15, 04, 05, 0, time.UTC),
}

type ss struct {
String string
Expand All @@ -2139,13 +2181,7 @@ func initQueryParameterTestCases() {
SubStructArray []ss
}

queryParameterTestCases = []struct {
name string
query string
parameters []QueryParameter
wantRow []Value
wantConfig interface{}
}{
queryParameterTestCases = []queryParameterTestCase{
{
"Int64Param",
"SELECT @val",
Expand Down Expand Up @@ -2412,6 +2448,51 @@ func initQueryParameterTestCases() {
},
},
}

if includeRangeCases {
queryParameterTestCases = append(queryParameterTestCases, []queryParameterTestCase{
{
"RangeUnboundedEnd",
"SELECT @val",
[]QueryParameter{
{
Name: "val",
Value: &QueryParameterValue{
Type: StandardSQLDataType{
TypeKind: "RANGE",
RangeElementType: &StandardSQLDataType{
TypeKind: "TIMESTAMP",
},
},
Value: rangeTimestamp1,
},
},
},
[]Value{rangeTimestamp1},
rangeTimestamp1,
},
{
"RangeUnboundedStart",
"SELECT @val",
[]QueryParameter{
{
Name: "val",
Value: &QueryParameterValue{
Type: StandardSQLDataType{
TypeKind: "RANGE",
RangeElementType: &StandardSQLDataType{
TypeKind: "TIMESTAMP",
},
},
Value: rangeTimestamp2,
},
},
},
[]Value{rangeTimestamp2},
rangeTimestamp2,
},
}...)
}
}

func TestIntegration_QueryParameters(t *testing.T) {
Expand All @@ -2420,7 +2501,7 @@ func TestIntegration_QueryParameters(t *testing.T) {
}
ctx := context.Background()

initQueryParameterTestCases()
initQueryParameterTestCases(true)

for _, tc := range queryParameterTestCases {
t.Run(tc.name, func(t *testing.T) {
Expand Down
69 changes: 67 additions & 2 deletions bigquery/params.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -86,6 +86,7 @@ var (
geographyParamType = &bq.QueryParameterType{Type: "GEOGRAPHY"}
intervalParamType = &bq.QueryParameterType{Type: "INTERVAL"}
jsonParamType = &bq.QueryParameterType{Type: "JSON"}
rangeParamType = &bq.QueryParameterType{Type: "RANGE"}
)

var (
Expand All @@ -95,6 +96,7 @@ var (
typeOfGoTime = reflect.TypeOf(time.Time{})
typeOfRat = reflect.TypeOf(&big.Rat{})
typeOfIntervalValue = reflect.TypeOf(&IntervalValue{})
typeOfRangeValue = reflect.TypeOf(&RangeValue{})
typeOfQueryParameterValue = reflect.TypeOf(&QueryParameterValue{})
)

Expand Down Expand Up @@ -315,9 +317,11 @@ func (p QueryParameter) toBQ() (*bq.QueryParameter, error) {
}, nil
}

var errNilParam = fmt.Errorf("bigquery: nil parameter")

func paramType(t reflect.Type, v reflect.Value) (*bq.QueryParameterType, error) {
if t == nil {
return nil, errors.New("bigquery: nil parameter")
return nil, errNilParam
}
switch t {
case typeOfDate, typeOfNullDate:
Expand All @@ -344,6 +348,25 @@ func paramType(t reflect.Type, v reflect.Value) (*bq.QueryParameterType, error)
return geographyParamType, nil
case typeOfNullJSON:
return jsonParamType, nil
case typeOfRangeValue:
iv := v.Interface().(*RangeValue)
// In order to autodetect a Range param correctly, at least one of start,end must be populated.
// Without it, users must declare typing via using QueryParameterValue.
element := iv.Start
if element == nil {
element = iv.End
}
if element == nil {
return nil, fmt.Errorf("unable to determine range element type from RangeValue without a non-nil start or end value")
}
elet, err := paramType(reflect.TypeOf(element), reflect.ValueOf(element))
if err != nil {
return nil, err
}
return &bq.QueryParameterType{
Type: "RANGE",
RangeElementType: elet,
}, nil
case typeOfQueryParameterValue:
return v.Interface().(*QueryParameterValue).toBQParamType(), nil
}
Expand Down Expand Up @@ -410,7 +433,7 @@ func paramType(t reflect.Type, v reflect.Value) (*bq.QueryParameterType, error)
func paramValue(v reflect.Value) (*bq.QueryParameterValue, error) {
res := &bq.QueryParameterValue{}
if !v.IsValid() {
return res, errors.New("bigquery: nil parameter")
return res, errNilParam
}
t := v.Type()
switch t {
Expand Down Expand Up @@ -492,6 +515,28 @@ func paramValue(v reflect.Value) (*bq.QueryParameterValue, error) {
case typeOfIntervalValue:
res.Value = IntervalString(v.Interface().(*IntervalValue))
return res, nil
case typeOfRangeValue:
// RangeValue is a compound type, and we must process the start/end to
// fully populate the value.
res.RangeValue = &bq.RangeValue{}
iv := v.Interface().(*RangeValue)
sVal, err := paramValue(reflect.ValueOf(iv.Start))
if err != nil {
if !errors.Is(err, errNilParam) {
return nil, err
}
} else {
res.RangeValue.Start = sVal
}
eVal, err := paramValue(reflect.ValueOf(iv.End))
if err != nil {
if !errors.Is(err, errNilParam) {
return nil, err
}
} else {
res.RangeValue.End = eVal
}
return res, nil
case typeOfQueryParameterValue:
return v.Interface().(*QueryParameterValue).toBQParamValue()
}
Expand Down Expand Up @@ -592,6 +637,26 @@ func convertParamValue(qval *bq.QueryParameterValue, qtype *bq.QueryParameterTyp
return map[string]interface{}(nil), nil
}
return convertParamStruct(qval.StructValues, qtype.StructTypes)
case "RANGE":
rv := &RangeValue{}
if qval.RangeValue == nil {
return rv, nil
}
if qval.RangeValue.Start != nil {
startVal, err := convertParamValue(qval.RangeValue.Start, qtype.RangeElementType)
if err != nil {
return nil, err
}
rv.Start = startVal
}
if qval.RangeValue.End != nil {
endVal, err := convertParamValue(qval.RangeValue.End, qtype.RangeElementType)
if err != nil {
return nil, err
}
rv.End = endVal
}
return rv, nil
case "TIMESTAMP":
if isNullScalar(qval) {
return NullTimestamp{Valid: false}, nil
Expand Down
Loading

0 comments on commit 07f0806

Please sign in to comment.