Skip to content

Commit

Permalink
feat(bigquery): RANGE type StandardSQLDataType support (#9754)
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
shollyman committed Apr 12, 2024
1 parent 7fc9c54 commit 33666cf
Show file tree
Hide file tree
Showing 3 changed files with 76 additions and 0 deletions.
38 changes: 38 additions & 0 deletions bigquery/routine_integration_test.go
Expand Up @@ -53,6 +53,44 @@ func TestIntegration_RoutineScalarUDF(t *testing.T) {
}
}

func TestIntegration_RoutineRangeType(t *testing.T) {
if client == nil {
t.Skip("Integration tests skipped")
}
ctx := context.Background()

routineID := routineIDs.New()
routine := dataset.Routine(routineID)
err := routine.Create(ctx, &RoutineMetadata{
Type: "SCALAR_FUNCTION",
Language: "SQL",
Body: "RANGE_CONTAINS(r1,r2)",
Arguments: []*RoutineArgument{
{
Name: "r1",
DataType: &StandardSQLDataType{
TypeKind: "RANGE",
RangeElementType: &StandardSQLDataType{
TypeKind: "TIMESTAMP",
},
},
},
{
Name: "r2",
DataType: &StandardSQLDataType{
TypeKind: "RANGE",
RangeElementType: &StandardSQLDataType{
TypeKind: "TIMESTAMP",
},
},
},
},
})
if err != nil {
t.Fatalf("Create: %v", err)
}
}

func TestIntegration_RoutineDataGovernance(t *testing.T) {
if client == nil {
t.Skip("Integration tests skipped")
Expand Down
24 changes: 24 additions & 0 deletions bigquery/standardsql.go
Expand Up @@ -26,6 +26,8 @@ type StandardSQLDataType struct {
// ArrayElementType indicates the type of an array's elements, when the
// TypeKind is ARRAY.
ArrayElementType *StandardSQLDataType
// The type of the range's elements, if TypeKind is RANGE.
RangeElementType *StandardSQLDataType
// StructType indicates the struct definition (fields), when the
// TypeKind is STRUCT.
StructType *StandardSQLStructType
Expand Down Expand Up @@ -60,6 +62,13 @@ func (ssdt *StandardSQLDataType) toBQ() (*bq.StandardSqlDataType, error) {
}
bqdt.StructType = dt
}
if ssdt.RangeElementType != nil {
dt, err := ssdt.RangeElementType.toBQ()
if err != nil {
return nil, err
}
bqdt.RangeElementType = dt
}
return bqdt, nil
}

Expand All @@ -77,6 +86,14 @@ func (ssdt StandardSQLDataType) toBQParamType() *bq.QueryParameterType {
}
return &bq.QueryParameterType{Type: "STRUCT", StructTypes: fts}
}
if ssdt.RangeElementType != nil {
return &bq.QueryParameterType{
Type: string(RangeFieldType),
RangeElementType: &bq.QueryParameterType{
Type: ssdt.RangeElementType.TypeKind,
},
}
}
return &bq.QueryParameterType{Type: ssdt.TypeKind}
}

Expand All @@ -102,6 +119,13 @@ func bqToStandardSQLDataType(bqdt *bq.StandardSqlDataType) (*StandardSQLDataType
}
ssdt.StructType = st
}
if bqdt.RangeElementType != nil {
st, err := bqToStandardSQLDataType(bqdt.RangeElementType)
if err != nil {
return nil, err
}
ssdt.RangeElementType = st
}
return ssdt, nil
}

Expand Down
14 changes: 14 additions & 0 deletions bigquery/standardsql_test.go
Expand Up @@ -49,6 +49,20 @@ func TestBQToStandardSQLDataType(t *testing.T) {
},
},
},
{
&bq.StandardSqlDataType{
TypeKind: "RANGE",
RangeElementType: &bq.StandardSqlDataType{
TypeKind: "DATETIME",
},
},
&StandardSQLDataType{
TypeKind: "RANGE",
RangeElementType: &StandardSQLDataType{
TypeKind: "DATETIME",
},
},
},
} {
got, err := bqToStandardSQLDataType(test.in)
if err != nil {
Expand Down

0 comments on commit 33666cf

Please sign in to comment.