Commit 420350df authored by Zhou Yaochen's avatar Zhou Yaochen
Browse files

update ignore

parent fc1f6363

Too many changes to show.

To preserve performance only 723 of 723+ files are displayed.
// Copyright 2018 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package bigquery
import (
"context"
"testing"
"cloud.google.com/go/internal/testutil"
)
func TestOCTracing(t *testing.T) {
ctx := context.Background()
client := getClient(t)
defer client.Close()
te := testutil.NewTestExporter()
defer te.Unregister()
q := client.Query("select *")
q.Run(ctx) // Doesn't matter if we get an error; span should be created either way
if len(te.Spans) == 0 {
t.Fatalf("Expected some spans to be created, but got %d", 0)
}
}
// Copyright 2016 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package bigquery
import (
"encoding/base64"
"errors"
"fmt"
"math/big"
"reflect"
"regexp"
"time"
"cloud.google.com/go/civil"
"cloud.google.com/go/internal/fields"
bq "google.golang.org/api/bigquery/v2"
)
var (
// See https://cloud.google.com/bigquery/docs/reference/standard-sql/data-types#timestamp-type.
timestampFormat = "2006-01-02 15:04:05.999999-07:00"
// See https://cloud.google.com/bigquery/docs/reference/rest/v2/tables#schema.fields.name
validFieldName = regexp.MustCompile("^[a-zA-Z_][a-zA-Z0-9_]{0,127}$")
)
const nullableTagOption = "nullable"
func bqTagParser(t reflect.StructTag) (name string, keep bool, other interface{}, err error) {
name, keep, opts, err := fields.ParseStandardTag("bigquery", t)
if err != nil {
return "", false, nil, err
}
if name != "" && !validFieldName.MatchString(name) {
return "", false, nil, errInvalidFieldName
}
for _, opt := range opts {
if opt != nullableTagOption {
return "", false, nil, fmt.Errorf(
"bigquery: invalid tag option %q. The only valid option is %q",
opt, nullableTagOption)
}
}
return name, keep, opts, nil
}
var fieldCache = fields.NewCache(bqTagParser, nil, nil)
var (
int64ParamType = &bq.QueryParameterType{Type: "INT64"}
float64ParamType = &bq.QueryParameterType{Type: "FLOAT64"}
boolParamType = &bq.QueryParameterType{Type: "BOOL"}
stringParamType = &bq.QueryParameterType{Type: "STRING"}
bytesParamType = &bq.QueryParameterType{Type: "BYTES"}
dateParamType = &bq.QueryParameterType{Type: "DATE"}
timeParamType = &bq.QueryParameterType{Type: "TIME"}
dateTimeParamType = &bq.QueryParameterType{Type: "DATETIME"}
timestampParamType = &bq.QueryParameterType{Type: "TIMESTAMP"}
numericParamType = &bq.QueryParameterType{Type: "NUMERIC"}
)
var (
typeOfDate = reflect.TypeOf(civil.Date{})
typeOfTime = reflect.TypeOf(civil.Time{})
typeOfDateTime = reflect.TypeOf(civil.DateTime{})
typeOfGoTime = reflect.TypeOf(time.Time{})
typeOfRat = reflect.TypeOf(&big.Rat{})
)
// A QueryParameter is a parameter to a query.
type QueryParameter struct {
// Name is used for named parameter mode.
// It must match the name in the query case-insensitively.
Name string
// Value is the value of the parameter.
//
// When you create a QueryParameter to send to BigQuery, the following Go types
// are supported, with their corresponding Bigquery types:
// int, int8, int16, int32, int64, uint8, uint16, uint32: INT64
// Note that uint, uint64 and uintptr are not supported, because
// they may contain values that cannot fit into a 64-bit signed integer.
// float32, float64: FLOAT64
// bool: BOOL
// string: STRING
// []byte: BYTES
// time.Time: TIMESTAMP
// *big.Rat: NUMERIC
// Arrays and slices of the above.
// Structs of the above. Only the exported fields are used.
//
// When a QueryParameter is returned inside a QueryConfig from a call to
// Job.Config:
// Integers are of type int64.
// Floating-point values are of type float64.
// Arrays are of type []interface{}, regardless of the array element type.
// Structs are of type map[string]interface{}.
Value interface{}
}
func (p QueryParameter) toBQ() (*bq.QueryParameter, error) {
pv, err := paramValue(reflect.ValueOf(p.Value))
if err != nil {
return nil, err
}
pt, err := paramType(reflect.TypeOf(p.Value))
if err != nil {
return nil, err
}
return &bq.QueryParameter{
Name: p.Name,
ParameterValue: &pv,
ParameterType: pt,
}, nil
}
func paramType(t reflect.Type) (*bq.QueryParameterType, error) {
if t == nil {
return nil, errors.New("bigquery: nil parameter")
}
switch t {
case typeOfDate:
return dateParamType, nil
case typeOfTime:
return timeParamType, nil
case typeOfDateTime:
return dateTimeParamType, nil
case typeOfGoTime:
return timestampParamType, nil
case typeOfRat:
return numericParamType, nil
}
switch t.Kind() {
case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64, reflect.Uint8, reflect.Uint16, reflect.Uint32:
return int64ParamType, nil
case reflect.Float32, reflect.Float64:
return float64ParamType, nil
case reflect.Bool:
return boolParamType, nil
case reflect.String:
return stringParamType, nil
case reflect.Slice:
if t.Elem().Kind() == reflect.Uint8 {
return bytesParamType, nil
}
fallthrough
case reflect.Array:
et, err := paramType(t.Elem())
if err != nil {
return nil, err
}
return &bq.QueryParameterType{Type: "ARRAY", ArrayType: et}, nil
case reflect.Ptr:
if t.Elem().Kind() != reflect.Struct {
break
}
t = t.Elem()
fallthrough
case reflect.Struct:
var fts []*bq.QueryParameterTypeStructTypes
fields, err := fieldCache.Fields(t)
if err != nil {
return nil, err
}
for _, f := range fields {
pt, err := paramType(f.Type)
if err != nil {
return nil, err
}
fts = append(fts, &bq.QueryParameterTypeStructTypes{
Name: f.Name,
Type: pt,
})
}
return &bq.QueryParameterType{Type: "STRUCT", StructTypes: fts}, nil
}
return nil, fmt.Errorf("bigquery: Go type %s cannot be represented as a parameter type", t)
}
func paramValue(v reflect.Value) (bq.QueryParameterValue, error) {
var res bq.QueryParameterValue
if !v.IsValid() {
return res, errors.New("bigquery: nil parameter")
}
t := v.Type()
switch t {
case typeOfDate:
res.Value = v.Interface().(civil.Date).String()
return res, nil
case typeOfTime:
// civil.Time has nanosecond resolution, but BigQuery TIME only microsecond.
// (If we send nanoseconds, then when we try to read the result we get "query job
// missing destination table").
res.Value = CivilTimeString(v.Interface().(civil.Time))
return res, nil
case typeOfDateTime:
res.Value = CivilDateTimeString(v.Interface().(civil.DateTime))
return res, nil
case typeOfGoTime:
res.Value = v.Interface().(time.Time).Format(timestampFormat)
return res, nil
case typeOfRat:
res.Value = NumericString(v.Interface().(*big.Rat))
return res, nil
}
switch t.Kind() {
case reflect.Slice:
if t.Elem().Kind() == reflect.Uint8 {
res.Value = base64.StdEncoding.EncodeToString(v.Interface().([]byte))
return res, nil
}
fallthrough
case reflect.Array:
var vals []*bq.QueryParameterValue
for i := 0; i < v.Len(); i++ {
val, err := paramValue(v.Index(i))
if err != nil {
return bq.QueryParameterValue{}, err
}
vals = append(vals, &val)
}
return bq.QueryParameterValue{ArrayValues: vals}, nil
case reflect.Ptr:
if t.Elem().Kind() != reflect.Struct {
return res, fmt.Errorf("bigquery: Go type %s cannot be represented as a parameter value", t)
}
t = t.Elem()
v = v.Elem()
if !v.IsValid() {
// nil pointer becomes empty value
return res, nil
}
fallthrough
case reflect.Struct:
fields, err := fieldCache.Fields(t)
if err != nil {
return bq.QueryParameterValue{}, err
}
res.StructValues = map[string]bq.QueryParameterValue{}
for _, f := range fields {
fv := v.FieldByIndex(f.Index)
fp, err := paramValue(fv)
if err != nil {
return bq.QueryParameterValue{}, err
}
res.StructValues[f.Name] = fp
}
return res, nil
}
// None of the above: assume a scalar type. (If it's not a valid type,
// paramType will catch the error.)
res.Value = fmt.Sprint(v.Interface())
return res, nil
}
func bqToQueryParameter(q *bq.QueryParameter) (QueryParameter, error) {
p := QueryParameter{Name: q.Name}
val, err := convertParamValue(q.ParameterValue, q.ParameterType)
if err != nil {
return QueryParameter{}, err
}
p.Value = val
return p, nil
}
var paramTypeToFieldType = map[string]FieldType{
int64ParamType.Type: IntegerFieldType,
float64ParamType.Type: FloatFieldType,
boolParamType.Type: BooleanFieldType,
stringParamType.Type: StringFieldType,
bytesParamType.Type: BytesFieldType,
dateParamType.Type: DateFieldType,
timeParamType.Type: TimeFieldType,
numericParamType.Type: NumericFieldType,
}
// Convert a parameter value from the service to a Go value. This is similar to, but
// not quite the same as, converting data values.
func convertParamValue(qval *bq.QueryParameterValue, qtype *bq.QueryParameterType) (interface{}, error) {
switch qtype.Type {
case "ARRAY":
if qval == nil {
return []interface{}(nil), nil
}
return convertParamArray(qval.ArrayValues, qtype.ArrayType)
case "STRUCT":
if qval == nil {
return map[string]interface{}(nil), nil
}
return convertParamStruct(qval.StructValues, qtype.StructTypes)
case "TIMESTAMP":
return time.Parse(timestampFormat, qval.Value)
case "DATETIME":
return parseCivilDateTime(qval.Value)
default:
return convertBasicType(qval.Value, paramTypeToFieldType[qtype.Type])
}
}
// convertParamArray converts a query parameter array value to a Go value. It
// always returns a []interface{}.
func convertParamArray(elVals []*bq.QueryParameterValue, elType *bq.QueryParameterType) ([]interface{}, error) {
var vals []interface{}
for _, el := range elVals {
val, err := convertParamValue(el, elType)
if err != nil {
return nil, err
}
vals = append(vals, val)
}
return vals, nil
}
// convertParamStruct converts a query parameter struct value into a Go value. It
// always returns a map[string]interface{}.
func convertParamStruct(sVals map[string]bq.QueryParameterValue, sTypes []*bq.QueryParameterTypeStructTypes) (map[string]interface{}, error) {
vals := map[string]interface{}{}
for _, st := range sTypes {
if sv, ok := sVals[st.Name]; ok {
val, err := convertParamValue(&sv, st.Type)
if err != nil {
return nil, err
}
vals[st.Name] = val
} else {
vals[st.Name] = nil
}
}
return vals, nil
}
// Copyright 2016 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package bigquery
import (
"context"
"errors"
"math"
"math/big"
"reflect"
"testing"
"time"
"cloud.google.com/go/civil"
"cloud.google.com/go/internal/testutil"
"github.com/google/go-cmp/cmp"
bq "google.golang.org/api/bigquery/v2"
)
var scalarTests = []struct {
val interface{} // The Go value
wantVal string // paramValue's desired output
wantType *bq.QueryParameterType // paramType's desired output
}{
{int64(0), "0", int64ParamType},
{3.14, "3.14", float64ParamType},
{3.14159e-87, "3.14159e-87", float64ParamType},
{true, "true", boolParamType},
{"string", "string", stringParamType},
{"\u65e5\u672c\u8a9e\n", "\u65e5\u672c\u8a9e\n", stringParamType},
{math.NaN(), "NaN", float64ParamType},
{[]byte("foo"), "Zm9v", bytesParamType}, // base64 encoding of "foo"
{time.Date(2016, 3, 20, 4, 22, 9, 5000, time.FixedZone("neg1-2", -3720)),
"2016-03-20 04:22:09.000005-01:02",
timestampParamType},
{civil.Date{Year: 2016, Month: 3, Day: 20}, "2016-03-20", dateParamType},
{civil.Time{Hour: 4, Minute: 5, Second: 6, Nanosecond: 789000000}, "04:05:06.789000", timeParamType},
{civil.DateTime{Date: civil.Date{Year: 2016, Month: 3, Day: 20}, Time: civil.Time{Hour: 4, Minute: 5, Second: 6, Nanosecond: 789000000}},
"2016-03-20 04:05:06.789000",
dateTimeParamType},
{big.NewRat(12345, 1000), "12.345000000", numericParamType},
}
type (
S1 struct {
A int
B *S2
C bool
}
S2 struct {
D string
e int
}
)
var (
s1 = S1{
A: 1,
B: &S2{D: "s"},
C: true,
}
s1ParamType = &bq.QueryParameterType{
Type: "STRUCT",
StructTypes: []*bq.QueryParameterTypeStructTypes{
{Name: "A", Type: int64ParamType},
{Name: "B", Type: &bq.QueryParameterType{
Type: "STRUCT",
StructTypes: []*bq.QueryParameterTypeStructTypes{
{Name: "D", Type: stringParamType},
},
}},
{Name: "C", Type: boolParamType},
},
}
s1ParamValue = bq.QueryParameterValue{
StructValues: map[string]bq.QueryParameterValue{
"A": sval("1"),
"B": {
StructValues: map[string]bq.QueryParameterValue{
"D": sval("s"),
},
},
"C": sval("true"),
},
}
s1ParamReturnValue = map[string]interface{}{
"A": int64(1),
"B": map[string]interface{}{"D": "s"},
"C": true,
}
)
func sval(s string) bq.QueryParameterValue {
return bq.QueryParameterValue{Value: s}
}
func TestParamValueScalar(t *testing.T) {
for _, test := range scalarTests {
got, err := paramValue(reflect.ValueOf(test.val))
if err != nil {
t.Errorf("%v: got %v, want nil", test.val, err)
continue
}
want := sval(test.wantVal)
if !testutil.Equal(got, want) {
t.Errorf("%v:\ngot %+v\nwant %+v", test.val, got, want)
}
}
}
func TestParamValueArray(t *testing.T) {
qpv := bq.QueryParameterValue{ArrayValues: []*bq.QueryParameterValue{
{Value: "1"},
{Value: "2"},
},
}
for _, test := range []struct {
val interface{}
want bq.QueryParameterValue
}{
{[]int(nil), bq.QueryParameterValue{}},
{[]int{}, bq.QueryParameterValue{}},
{[]int{1, 2}, qpv},
{[2]int{1, 2}, qpv},
} {
got, err := paramValue(reflect.ValueOf(test.val))
if err != nil {
t.Fatal(err)
}
if !testutil.Equal(got, test.want) {
t.Errorf("%#v:\ngot %+v\nwant %+v", test.val, got, test.want)
}
}
}
func TestParamValueStruct(t *testing.T) {
got, err := paramValue(reflect.ValueOf(s1))
if err != nil {
t.Fatal(err)
}
if !testutil.Equal(got, s1ParamValue) {
t.Errorf("got %+v\nwant %+v", got, s1ParamValue)
}
}
func TestParamValueErrors(t *testing.T) {
// paramValue lets a few invalid types through, but paramType catches them.
// Since we never call one without the other that's fine.
for _, val := range []interface{}{nil, new([]int)} {
_, err := paramValue(reflect.ValueOf(val))
if err == nil {
t.Errorf("%v (%T): got nil, want error", val, val)
}
}
}
func TestParamType(t *testing.T) {
for _, test := range scalarTests {
got, err := paramType(reflect.TypeOf(test.val))
if err != nil {
t.Fatal(err)
}
if !testutil.Equal(got, test.wantType) {
t.Errorf("%v (%T): got %v, want %v", test.val, test.val, got, test.wantType)
}
}
for _, test := range []struct {
val interface{}
want *bq.QueryParameterType
}{
{uint32(32767), int64ParamType},
{[]byte("foo"), bytesParamType},
{[]int{}, &bq.QueryParameterType{Type: "ARRAY", ArrayType: int64ParamType}},
{[3]bool{}, &bq.QueryParameterType{Type: "ARRAY", ArrayType: boolParamType}},
{S1{}, s1ParamType},
} {
got, err := paramType(reflect.TypeOf(test.val))
if err != nil {
t.Fatal(err)
}
if !testutil.Equal(got, test.want) {
t.Errorf("%v (%T): got %v, want %v", test.val, test.val, got, test.want)
}
}
}
func TestParamTypeErrors(t *testing.T) {
for _, val := range []interface{}{
nil, uint(0), new([]int), make(chan int),
} {
_, err := paramType(reflect.TypeOf(val))
if err == nil {
t.Errorf("%v (%T): got nil, want error", val, val)
}
}
}
func TestConvertParamValue(t *testing.T) {
// Scalars.
for _, test := range scalarTests {
pval, err := paramValue(reflect.ValueOf(test.val))
if err != nil {
t.Fatal(err)
}
ptype, err := paramType(reflect.TypeOf(test.val))
if err != nil {
t.Fatal(err)
}
got, err := convertParamValue(&pval, ptype)
if err != nil {
t.Fatalf("convertParamValue(%+v, %+v): %v", pval, ptype, err)
}
if !testutil.Equal(got, test.val) {
t.Errorf("%#v: got %#v", test.val, got)
}
}
// Arrays.
for _, test := range []struct {
pval *bq.QueryParameterValue
want []interface{}
}{
{
&bq.QueryParameterValue{},
nil,
},
{
&bq.QueryParameterValue{
ArrayValues: []*bq.QueryParameterValue{{Value: "1"}, {Value: "2"}},
},
[]interface{}{int64(1), int64(2)},
},
} {
ptype := &bq.QueryParameterType{Type: "ARRAY", ArrayType: int64ParamType}
got, err := convertParamValue(test.pval, ptype)
if err != nil {
t.Fatalf("%+v: %v", test.pval, err)
}
if !testutil.Equal(got, test.want) {
t.Errorf("%+v: got %+v, want %+v", test.pval, got, test.want)
}
}
// Structs.
got, err := convertParamValue(&s1ParamValue, s1ParamType)
if err != nil {
t.Fatal(err)
}
if !testutil.Equal(got, s1ParamReturnValue) {
t.Errorf("got %+v, want %+v", got, s1ParamReturnValue)
}
}
func TestIntegration_ScalarParam(t *testing.T) {
roundToMicros := cmp.Transformer("RoundToMicros",
func(t time.Time) time.Time { return t.Round(time.Microsecond) })
c := getClient(t)
for _, test := range scalarTests {
gotData, gotParam, err := paramRoundTrip(c, test.val)
if err != nil {
t.Fatal(err)
}
if !testutil.Equal(gotData, test.val, roundToMicros) {
t.Errorf("\ngot %#v (%T)\nwant %#v (%T)", gotData, gotData, test.val, test.val)
}
if !testutil.Equal(gotParam, test.val, roundToMicros) {
t.Errorf("\ngot %#v (%T)\nwant %#v (%T)", gotParam, gotParam, test.val, test.val)
}
}
}
func TestIntegration_OtherParam(t *testing.T) {
c := getClient(t)
for _, test := range []struct {
val interface{}
wantData interface{}
wantParam interface{}
}{
{[]int(nil), []Value(nil), []interface{}(nil)},
{[]int{}, []Value(nil), []interface{}(nil)},
{
[]int{1, 2},
[]Value{int64(1), int64(2)},
[]interface{}{int64(1), int64(2)},
},
{
[3]int{1, 2, 3},
[]Value{int64(1), int64(2), int64(3)},
[]interface{}{int64(1), int64(2), int64(3)},
},
{
S1{},
[]Value{int64(0), nil, false},
map[string]interface{}{
"A": int64(0),
"B": nil,
"C": false,
},
},
{
s1,
[]Value{int64(1), []Value{"s"}, true},
s1ParamReturnValue,
},
} {
gotData, gotParam, err := paramRoundTrip(c, test.val)
if err != nil {
t.Fatal(err)
}
if !testutil.Equal(gotData, test.wantData) {
t.Errorf("%#v:\ngot %#v (%T)\nwant %#v (%T)",
test.val, gotData, gotData, test.wantData, test.wantData)
}
if !testutil.Equal(gotParam, test.wantParam) {
t.Errorf("%#v:\ngot %#v (%T)\nwant %#v (%T)",
test.val, gotParam, gotParam, test.wantParam, test.wantParam)
}
}
}
// paramRoundTrip passes x as a query parameter to BigQuery. It returns
// the resulting data value from running the query and the parameter value from
// the returned job configuration.
func paramRoundTrip(c *Client, x interface{}) (data Value, param interface{}, err error) {
ctx := context.Background()
q := c.Query("select ?")
q.Parameters = []QueryParameter{{Value: x}}
job, err := q.Run(ctx)
if err != nil {
return nil, nil, err
}
it, err := job.Read(ctx)
if err != nil {
return nil, nil, err
}
var val []Value
err = it.Next(&val)
if err != nil {
return nil, nil, err
}
if len(val) != 1 {
return nil, nil, errors.New("wrong number of values")
}
conf, err := job.Config()
if err != nil {
return nil, nil, err
}
return val[0], conf.(*QueryConfig).Parameters[0].Value, nil
}
// Copyright 2015 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package bigquery
import (
"context"
"errors"
"cloud.google.com/go/internal/trace"
bq "google.golang.org/api/bigquery/v2"
)
// QueryConfig holds the configuration for a query job.
type QueryConfig struct {
// Dst is the table into which the results of the query will be written.
// If this field is nil, a temporary table will be created.
Dst *Table
// The query to execute. See https://cloud.google.com/bigquery/query-reference for details.
Q string
// DefaultProjectID and DefaultDatasetID specify the dataset to use for unqualified table names in the query.
// If DefaultProjectID is set, DefaultDatasetID must also be set.
DefaultProjectID string
DefaultDatasetID string
// TableDefinitions describes data sources outside of BigQuery.
// The map keys may be used as table names in the query string.
//
// When a QueryConfig is returned from Job.Config, the map values
// are always of type *ExternalDataConfig.
TableDefinitions map[string]ExternalData
// CreateDisposition specifies the circumstances under which the destination table will be created.
// The default is CreateIfNeeded.
CreateDisposition TableCreateDisposition
// WriteDisposition specifies how existing data in the destination table is treated.
// The default is WriteEmpty.
WriteDisposition TableWriteDisposition
// DisableQueryCache prevents results being fetched from the query cache.
// If this field is false, results are fetched from the cache if they are available.
// The query cache is a best-effort cache that is flushed whenever tables in the query are modified.
// Cached results are only available when TableID is unspecified in the query's destination Table.
// For more information, see https://cloud.google.com/bigquery/querying-data#querycaching
DisableQueryCache bool
// DisableFlattenedResults prevents results being flattened.
// If this field is false, results from nested and repeated fields are flattened.
// DisableFlattenedResults implies AllowLargeResults
// For more information, see https://cloud.google.com/bigquery/docs/data#nested
DisableFlattenedResults bool
// AllowLargeResults allows the query to produce arbitrarily large result tables.
// The destination must be a table.
// When using this option, queries will take longer to execute, even if the result set is small.
// For additional limitations, see https://cloud.google.com/bigquery/querying-data#largequeryresults
AllowLargeResults bool
// Priority specifies the priority with which to schedule the query.
// The default priority is InteractivePriority.
// For more information, see https://cloud.google.com/bigquery/querying-data#batchqueries
Priority QueryPriority
// MaxBillingTier sets the maximum billing tier for a Query.
// Queries that have resource usage beyond this tier will fail (without
// incurring a charge). If this field is zero, the project default will be used.
MaxBillingTier int
// MaxBytesBilled limits the number of bytes billed for
// this job. Queries that would exceed this limit will fail (without incurring
// a charge).
// If this field is less than 1, the project default will be
// used.
MaxBytesBilled int64
// UseStandardSQL causes the query to use standard SQL. The default.
// Deprecated: use UseLegacySQL.
UseStandardSQL bool
// UseLegacySQL causes the query to use legacy SQL.
UseLegacySQL bool
// Parameters is a list of query parameters. The presence of parameters
// implies the use of standard SQL.
// If the query uses positional syntax ("?"), then no parameter may have a name.
// If the query uses named syntax ("@p"), then all parameters must have names.
// It is illegal to mix positional and named syntax.
Parameters []QueryParameter
// TimePartitioning specifies time-based partitioning
// for the destination table.
TimePartitioning *TimePartitioning
// Clustering specifies the data clustering configuration for the destination table.
Clustering *Clustering
// The labels associated with this job.
Labels map[string]string
// If true, don't actually run this job. A valid query will return a mostly
// empty response with some processing statistics, while an invalid query will
// return the same error it would if it wasn't a dry run.
//
// Query.Read will fail with dry-run queries. Call Query.Run instead, and then
// call LastStatus on the returned job to get statistics. Calling Status on a
// dry-run job will fail.
DryRun bool
// Custom encryption configuration (e.g., Cloud KMS keys).
DestinationEncryptionConfig *EncryptionConfig
// Allows the schema of the destination table to be updated as a side effect of
// the query job.
SchemaUpdateOptions []string
}
func (qc *QueryConfig) toBQ() (*bq.JobConfiguration, error) {
qconf := &bq.JobConfigurationQuery{
Query: qc.Q,
CreateDisposition: string(qc.CreateDisposition),
WriteDisposition: string(qc.WriteDisposition),
AllowLargeResults: qc.AllowLargeResults,
Priority: string(qc.Priority),
MaximumBytesBilled: qc.MaxBytesBilled,
TimePartitioning: qc.TimePartitioning.toBQ(),
Clustering: qc.Clustering.toBQ(),
DestinationEncryptionConfiguration: qc.DestinationEncryptionConfig.toBQ(),
SchemaUpdateOptions: qc.SchemaUpdateOptions,
}
if len(qc.TableDefinitions) > 0 {
qconf.TableDefinitions = make(map[string]bq.ExternalDataConfiguration)
}
for name, data := range qc.TableDefinitions {
qconf.TableDefinitions[name] = data.toBQ()
}
if qc.DefaultProjectID != "" || qc.DefaultDatasetID != "" {
qconf.DefaultDataset = &bq.DatasetReference{
DatasetId: qc.DefaultDatasetID,
ProjectId: qc.DefaultProjectID,
}
}
if tier := int64(qc.MaxBillingTier); tier > 0 {
qconf.MaximumBillingTier = &tier
}
f := false
if qc.DisableQueryCache {
qconf.UseQueryCache = &f
}
if qc.DisableFlattenedResults {
qconf.FlattenResults = &f
// DisableFlattenResults implies AllowLargeResults.
qconf.AllowLargeResults = true
}
if qc.UseStandardSQL && qc.UseLegacySQL {
return nil, errors.New("bigquery: cannot provide both UseStandardSQL and UseLegacySQL")
}
if len(qc.Parameters) > 0 && qc.UseLegacySQL {
return nil, errors.New("bigquery: cannot provide both Parameters (implying standard SQL) and UseLegacySQL")
}
ptrue := true
pfalse := false
if qc.UseLegacySQL {
qconf.UseLegacySql = &ptrue
} else {
qconf.UseLegacySql = &pfalse
}
if qc.Dst != nil && !qc.Dst.implicitTable() {
qconf.DestinationTable = qc.Dst.toBQ()
}
for _, p := range qc.Parameters {
qp, err := p.toBQ()
if err != nil {
return nil, err
}
qconf.QueryParameters = append(qconf.QueryParameters, qp)
}
return &bq.JobConfiguration{
Labels: qc.Labels,
DryRun: qc.DryRun,
Query: qconf,
}, nil
}
func bqToQueryConfig(q *bq.JobConfiguration, c *Client) (*QueryConfig, error) {
qq := q.Query
qc := &QueryConfig{
Labels: q.Labels,
DryRun: q.DryRun,
Q: qq.Query,
CreateDisposition: TableCreateDisposition(qq.CreateDisposition),
WriteDisposition: TableWriteDisposition(qq.WriteDisposition),
AllowLargeResults: qq.AllowLargeResults,
Priority: QueryPriority(qq.Priority),
MaxBytesBilled: qq.MaximumBytesBilled,
UseLegacySQL: qq.UseLegacySql == nil || *qq.UseLegacySql,
TimePartitioning: bqToTimePartitioning(qq.TimePartitioning),
Clustering: bqToClustering(qq.Clustering),
DestinationEncryptionConfig: bqToEncryptionConfig(qq.DestinationEncryptionConfiguration),
SchemaUpdateOptions: qq.SchemaUpdateOptions,
}
qc.UseStandardSQL = !qc.UseLegacySQL
if len(qq.TableDefinitions) > 0 {
qc.TableDefinitions = make(map[string]ExternalData)
}
for name, qedc := range qq.TableDefinitions {
edc, err := bqToExternalDataConfig(&qedc)
if err != nil {
return nil, err
}
qc.TableDefinitions[name] = edc
}
if qq.DefaultDataset != nil {
qc.DefaultProjectID = qq.DefaultDataset.ProjectId
qc.DefaultDatasetID = qq.DefaultDataset.DatasetId
}
if qq.MaximumBillingTier != nil {
qc.MaxBillingTier = int(*qq.MaximumBillingTier)
}
if qq.UseQueryCache != nil && !*qq.UseQueryCache {
qc.DisableQueryCache = true
}
if qq.FlattenResults != nil && !*qq.FlattenResults {
qc.DisableFlattenedResults = true
}
if qq.DestinationTable != nil {
qc.Dst = bqToTable(qq.DestinationTable, c)
}
for _, qp := range qq.QueryParameters {
p, err := bqToQueryParameter(qp)
if err != nil {
return nil, err
}
qc.Parameters = append(qc.Parameters, p)
}
return qc, nil
}
// QueryPriority specifies a priority with which a query is to be executed.
type QueryPriority string
const (
// BatchPriority specifies that the query should be scheduled with the
// batch priority. BigQuery queues each batch query on your behalf, and
// starts the query as soon as idle resources are available, usually within
// a few minutes. If BigQuery hasn't started the query within 24 hours,
// BigQuery changes the job priority to interactive. Batch queries don't
// count towards your concurrent rate limit, which can make it easier to
// start many queries at once.
//
// More information can be found at https://cloud.google.com/bigquery/docs/running-queries#batchqueries.
BatchPriority QueryPriority = "BATCH"
// InteractivePriority specifies that the query should be scheduled with
// interactive priority, which means that the query is executed as soon as
// possible. Interactive queries count towards your concurrent rate limit
// and your daily limit. It is the default priority with which queries get
// executed.
//
// More information can be found at https://cloud.google.com/bigquery/docs/running-queries#queries.
InteractivePriority QueryPriority = "INTERACTIVE"
)
// A Query queries data from a BigQuery table. Use Client.Query to create a Query.
type Query struct {
JobIDConfig
QueryConfig
client *Client
}
// Query creates a query with string q.
// The returned Query may optionally be further configured before its Run method is called.
func (c *Client) Query(q string) *Query {
return &Query{
client: c,
QueryConfig: QueryConfig{Q: q},
}
}
// Run initiates a query job.
func (q *Query) Run(ctx context.Context) (j *Job, err error) {
ctx = trace.StartSpan(ctx, "cloud.google.com/go/bigquery.Query.Run")
defer func() { trace.EndSpan(ctx, err) }()
job, err := q.newJob()
if err != nil {
return nil, err
}
j, err = q.client.insertJob(ctx, job, nil)
if err != nil {
return nil, err
}
return j, nil
}
func (q *Query) newJob() (*bq.Job, error) {
config, err := q.QueryConfig.toBQ()
if err != nil {
return nil, err
}
return &bq.Job{
JobReference: q.JobIDConfig.createJobRef(q.client),
Configuration: config,
}, nil
}
// Read submits a query for execution and returns the results via a RowIterator.
// It is a shorthand for Query.Run followed by Job.Read.
func (q *Query) Read(ctx context.Context) (*RowIterator, error) {
job, err := q.Run(ctx)
if err != nil {
return nil, err
}
return job.Read(ctx)
}
// Copyright 2015 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package bigquery
import (
"testing"
"time"
"cloud.google.com/go/internal/testutil"
"github.com/google/go-cmp/cmp"
bq "google.golang.org/api/bigquery/v2"
)
func defaultQueryJob() *bq.Job {
pfalse := false
return &bq.Job{
JobReference: &bq.JobReference{JobId: "RANDOM", ProjectId: "client-project-id"},
Configuration: &bq.JobConfiguration{
Query: &bq.JobConfigurationQuery{
DestinationTable: &bq.TableReference{
ProjectId: "client-project-id",
DatasetId: "dataset-id",
TableId: "table-id",
},
Query: "query string",
DefaultDataset: &bq.DatasetReference{
ProjectId: "def-project-id",
DatasetId: "def-dataset-id",
},
UseLegacySql: &pfalse,
},
},
}
}
var defaultQuery = &QueryConfig{
Q: "query string",
DefaultProjectID: "def-project-id",
DefaultDatasetID: "def-dataset-id",
}
func TestQuery(t *testing.T) {
defer fixRandomID("RANDOM")()
c := &Client{
projectID: "client-project-id",
}
testCases := []struct {
dst *Table
src *QueryConfig
jobIDConfig JobIDConfig
want *bq.Job
}{
{
dst: c.Dataset("dataset-id").Table("table-id"),
src: defaultQuery,
want: defaultQueryJob(),
},
{
dst: c.Dataset("dataset-id").Table("table-id"),
src: &QueryConfig{
Q: "query string",
Labels: map[string]string{"a": "b"},
DryRun: true,
},
want: func() *bq.Job {
j := defaultQueryJob()
j.Configuration.Labels = map[string]string{"a": "b"}
j.Configuration.DryRun = true
j.Configuration.Query.DefaultDataset = nil
return j
}(),
},
{
dst: c.Dataset("dataset-id").Table("table-id"),
jobIDConfig: JobIDConfig{JobID: "jobID", AddJobIDSuffix: true},
src: &QueryConfig{Q: "query string"},
want: func() *bq.Job {
j := defaultQueryJob()
j.Configuration.Query.DefaultDataset = nil
j.JobReference.JobId = "jobID-RANDOM"
return j
}(),
},
{
dst: &Table{},
src: defaultQuery,
want: func() *bq.Job {
j := defaultQueryJob()
j.Configuration.Query.DestinationTable = nil
return j
}(),
},
{
dst: c.Dataset("dataset-id").Table("table-id"),
src: &QueryConfig{
Q: "query string",
TableDefinitions: map[string]ExternalData{
"atable": func() *GCSReference {
g := NewGCSReference("uri")
g.AllowJaggedRows = true
g.AllowQuotedNewlines = true
g.Compression = Gzip
g.Encoding = UTF_8
g.FieldDelimiter = ";"
g.IgnoreUnknownValues = true
g.MaxBadRecords = 1
g.Quote = "'"
g.SkipLeadingRows = 2
g.Schema = Schema{{Name: "name", Type: StringFieldType}}
return g
}(),
},
},
want: func() *bq.Job {
j := defaultQueryJob()
j.Configuration.Query.DefaultDataset = nil
td := make(map[string]bq.ExternalDataConfiguration)
quote := "'"
td["atable"] = bq.ExternalDataConfiguration{
Compression: "GZIP",
IgnoreUnknownValues: true,
MaxBadRecords: 1,
SourceFormat: "CSV", // must be explicitly set.
SourceUris: []string{"uri"},
CsvOptions: &bq.CsvOptions{
AllowJaggedRows: true,
AllowQuotedNewlines: true,
Encoding: "UTF-8",
FieldDelimiter: ";",
SkipLeadingRows: 2,
Quote: &quote,
},
Schema: &bq.TableSchema{
Fields: []*bq.TableFieldSchema{
{Name: "name", Type: "STRING"},
},
},
}
j.Configuration.Query.TableDefinitions = td
return j
}(),
},
{
dst: &Table{
ProjectID: "project-id",
DatasetID: "dataset-id",
TableID: "table-id",
},
src: &QueryConfig{
Q: "query string",
DefaultProjectID: "def-project-id",
DefaultDatasetID: "def-dataset-id",
CreateDisposition: CreateNever,
WriteDisposition: WriteTruncate,
},
want: func() *bq.Job {
j := defaultQueryJob()
j.Configuration.Query.DestinationTable.ProjectId = "project-id"
j.Configuration.Query.WriteDisposition = "WRITE_TRUNCATE"
j.Configuration.Query.CreateDisposition = "CREATE_NEVER"
return j
}(),
},
{
dst: c.Dataset("dataset-id").Table("table-id"),
src: &QueryConfig{
Q: "query string",
DefaultProjectID: "def-project-id",
DefaultDatasetID: "def-dataset-id",
DisableQueryCache: true,
},
want: func() *bq.Job {
j := defaultQueryJob()
f := false
j.Configuration.Query.UseQueryCache = &f
return j
}(),
},
{
dst: c.Dataset("dataset-id").Table("table-id"),
src: &QueryConfig{
Q: "query string",
DefaultProjectID: "def-project-id",
DefaultDatasetID: "def-dataset-id",
AllowLargeResults: true,
},
want: func() *bq.Job {
j := defaultQueryJob()
j.Configuration.Query.AllowLargeResults = true
return j
}(),
},
{
dst: c.Dataset("dataset-id").Table("table-id"),
src: &QueryConfig{
Q: "query string",
DefaultProjectID: "def-project-id",
DefaultDatasetID: "def-dataset-id",
DisableFlattenedResults: true,
},
want: func() *bq.Job {
j := defaultQueryJob()
f := false
j.Configuration.Query.FlattenResults = &f
j.Configuration.Query.AllowLargeResults = true
return j
}(),
},
{
dst: c.Dataset("dataset-id").Table("table-id"),
src: &QueryConfig{
Q: "query string",
DefaultProjectID: "def-project-id",
DefaultDatasetID: "def-dataset-id",
Priority: QueryPriority("low"),
},
want: func() *bq.Job {
j := defaultQueryJob()
j.Configuration.Query.Priority = "low"
return j
}(),
},
{
dst: c.Dataset("dataset-id").Table("table-id"),
src: &QueryConfig{
Q: "query string",
DefaultProjectID: "def-project-id",
DefaultDatasetID: "def-dataset-id",
MaxBillingTier: 3,
MaxBytesBilled: 5,
},
want: func() *bq.Job {
j := defaultQueryJob()
tier := int64(3)
j.Configuration.Query.MaximumBillingTier = &tier
j.Configuration.Query.MaximumBytesBilled = 5
return j
}(),
},
{
dst: c.Dataset("dataset-id").Table("table-id"),
src: &QueryConfig{
Q: "query string",
DefaultProjectID: "def-project-id",
DefaultDatasetID: "def-dataset-id",
UseStandardSQL: true,
},
want: defaultQueryJob(),
},
{
dst: c.Dataset("dataset-id").Table("table-id"),
src: &QueryConfig{
Q: "query string",
DefaultProjectID: "def-project-id",
DefaultDatasetID: "def-dataset-id",
UseLegacySQL: true,
},
want: func() *bq.Job {
j := defaultQueryJob()
ptrue := true
j.Configuration.Query.UseLegacySql = &ptrue
j.Configuration.Query.ForceSendFields = nil
return j
}(),
},
}
for i, tc := range testCases {
query := c.Query("")
query.JobIDConfig = tc.jobIDConfig
query.QueryConfig = *tc.src
query.Dst = tc.dst
got, err := query.newJob()
if err != nil {
t.Errorf("#%d: err calling query: %v", i, err)
continue
}
checkJob(t, i, got, tc.want)
// Round-trip.
jc, err := bqToJobConfig(got.Configuration, c)
if err != nil {
t.Fatalf("#%d: %v", i, err)
}
wantConfig := query.QueryConfig
// We set AllowLargeResults to true when DisableFlattenedResults is true.
if wantConfig.DisableFlattenedResults {
wantConfig.AllowLargeResults = true
}
// A QueryConfig with neither UseXXXSQL field set is equivalent
// to one where UseStandardSQL = true.
if !wantConfig.UseLegacySQL && !wantConfig.UseStandardSQL {
wantConfig.UseStandardSQL = true
}
// Treat nil and empty tables the same, and ignore the client.
tableEqual := func(t1, t2 *Table) bool {
if t1 == nil {
t1 = &Table{}
}
if t2 == nil {
t2 = &Table{}
}
return t1.ProjectID == t2.ProjectID && t1.DatasetID == t2.DatasetID && t1.TableID == t2.TableID
}
// A table definition that is a GCSReference round-trips as an ExternalDataConfig.
// TODO(jba): see if there is a way to express this with a transformer.
gcsRefToEDC := func(g *GCSReference) *ExternalDataConfig {
q := g.toBQ()
e, _ := bqToExternalDataConfig(&q)
return e
}
externalDataEqual := func(e1, e2 ExternalData) bool {
if r, ok := e1.(*GCSReference); ok {
e1 = gcsRefToEDC(r)
}
if r, ok := e2.(*GCSReference); ok {
e2 = gcsRefToEDC(r)
}
return cmp.Equal(e1, e2)
}
diff := testutil.Diff(jc.(*QueryConfig), &wantConfig,
cmp.Comparer(tableEqual),
cmp.Comparer(externalDataEqual),
)
if diff != "" {
t.Errorf("#%d: (got=-, want=+:\n%s", i, diff)
}
}
}
func TestConfiguringQuery(t *testing.T) {
c := &Client{
projectID: "project-id",
}
query := c.Query("q")
query.JobID = "ajob"
query.DefaultProjectID = "def-project-id"
query.DefaultDatasetID = "def-dataset-id"
query.TimePartitioning = &TimePartitioning{Expiration: 1234 * time.Second, Field: "f"}
query.Clustering = &Clustering{
Fields: []string{"cfield1"},
}
query.DestinationEncryptionConfig = &EncryptionConfig{KMSKeyName: "keyName"}
query.SchemaUpdateOptions = []string{"ALLOW_FIELD_ADDITION"}
// Note: Other configuration fields are tested in other tests above.
// A lot of that can be consolidated once Client.Copy is gone.
pfalse := false
want := &bq.Job{
Configuration: &bq.JobConfiguration{
Query: &bq.JobConfigurationQuery{
Query: "q",
DefaultDataset: &bq.DatasetReference{
ProjectId: "def-project-id",
DatasetId: "def-dataset-id",
},
UseLegacySql: &pfalse,
TimePartitioning: &bq.TimePartitioning{ExpirationMs: 1234000, Field: "f", Type: "DAY"},
Clustering: &bq.Clustering{Fields: []string{"cfield1"}},
DestinationEncryptionConfiguration: &bq.EncryptionConfiguration{KmsKeyName: "keyName"},
SchemaUpdateOptions: []string{"ALLOW_FIELD_ADDITION"},
},
},
JobReference: &bq.JobReference{
JobId: "ajob",
ProjectId: "project-id",
},
}
got, err := query.newJob()
if err != nil {
t.Fatalf("err calling Query.newJob: %v", err)
}
if diff := testutil.Diff(got, want); diff != "" {
t.Errorf("querying: -got +want:\n%s", diff)
}
}
func TestQueryLegacySQL(t *testing.T) {
c := &Client{projectID: "project-id"}
q := c.Query("q")
q.UseStandardSQL = true
q.UseLegacySQL = true
_, err := q.newJob()
if err == nil {
t.Error("UseStandardSQL and UseLegacySQL: got nil, want error")
}
q = c.Query("q")
q.Parameters = []QueryParameter{{Name: "p", Value: 3}}
q.UseLegacySQL = true
_, err = q.newJob()
if err == nil {
t.Error("Parameters and UseLegacySQL: got nil, want error")
}
}
// Copyright 2018 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package bigquery
import (
"math/rand"
"os"
"sync"
"time"
)
// Support for random values (typically job IDs and insert IDs).
const alphanum = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789"
var (
rngMu sync.Mutex
rng = rand.New(rand.NewSource(time.Now().UnixNano() ^ int64(os.Getpid())))
)
// For testing.
var randomIDFn = randomID
// As of August 2017, the BigQuery service uses 27 alphanumeric characters for
// suffixes.
const randomIDLen = 27
func randomID() string {
// This is used for both job IDs and insert IDs.
var b [randomIDLen]byte
rngMu.Lock()
for i := 0; i < len(b); i++ {
b[i] = alphanum[rng.Intn(len(alphanum))]
}
rngMu.Unlock()
return string(b[:])
}
// Seed seeds this package's random number generator, used for generating job and
// insert IDs. Use Seed to obtain repeatable, deterministic behavior from bigquery
// clients. Seed should be called before any clients are created.
func Seed(s int64) {
rng = rand.New(rand.NewSource(s))
}
// Copyright 2015 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package bigquery
import (
"context"
"errors"
"testing"
"cloud.google.com/go/internal/testutil"
"github.com/google/go-cmp/cmp"
bq "google.golang.org/api/bigquery/v2"
"google.golang.org/api/iterator"
)
type pageFetcherArgs struct {
table *Table
schema Schema
startIndex uint64
pageSize int64
pageToken string
}
// pageFetcherReadStub services read requests by returning data from an in-memory list of values.
type pageFetcherReadStub struct {
// values and pageTokens are used as sources of data to return in response to calls to readTabledata or readQuery.
values [][][]Value // contains pages / rows / columns.
pageTokens map[string]string // maps incoming page token to returned page token.
// arguments are recorded for later inspection.
calls []pageFetcherArgs
}
func (s *pageFetcherReadStub) fetchPage(ctx context.Context, t *Table, schema Schema, startIndex uint64, pageSize int64, pageToken string) (*fetchPageResult, error) {
s.calls = append(s.calls,
pageFetcherArgs{t, schema, startIndex, pageSize, pageToken})
result := &fetchPageResult{
pageToken: s.pageTokens[pageToken],
rows: s.values[0],
}
s.values = s.values[1:]
return result, nil
}
func waitForQueryStub(context.Context, string) (Schema, uint64, error) {
return nil, 1, nil
}
func TestRead(t *testing.T) {
// The data for the service stub to return is populated for each test case in the testCases for loop.
ctx := context.Background()
c := &Client{projectID: "project-id"}
pf := &pageFetcherReadStub{}
queryJob := &Job{
projectID: "project-id",
jobID: "job-id",
c: c,
config: &bq.JobConfiguration{
Query: &bq.JobConfigurationQuery{
DestinationTable: &bq.TableReference{
ProjectId: "project-id",
DatasetId: "dataset-id",
TableId: "table-id",
},
},
},
}
for _, readFunc := range []func() *RowIterator{
func() *RowIterator {
return c.Dataset("dataset-id").Table("table-id").read(ctx, pf.fetchPage)
},
func() *RowIterator {
it, err := queryJob.read(ctx, waitForQueryStub, pf.fetchPage)
if err != nil {
t.Fatal(err)
}
return it
},
} {
testCases := []struct {
data [][][]Value
pageTokens map[string]string
want [][]Value
}{
{
data: [][][]Value{{{1, 2}, {11, 12}}, {{30, 40}, {31, 41}}},
pageTokens: map[string]string{"": "a", "a": ""},
want: [][]Value{{1, 2}, {11, 12}, {30, 40}, {31, 41}},
},
{
data: [][][]Value{{{1, 2}, {11, 12}}, {{30, 40}, {31, 41}}},
pageTokens: map[string]string{"": ""}, // no more pages after first one.
want: [][]Value{{1, 2}, {11, 12}},
},
}
for _, tc := range testCases {
pf.values = tc.data
pf.pageTokens = tc.pageTokens
if got, ok := collectValues(t, readFunc()); ok {
if !testutil.Equal(got, tc.want) {
t.Errorf("reading: got:\n%v\nwant:\n%v", got, tc.want)
}
}
}
}
}
func collectValues(t *testing.T, it *RowIterator) ([][]Value, bool) {
var got [][]Value
for {
var vals []Value
err := it.Next(&vals)
if err == iterator.Done {
break
}
if err != nil {
t.Errorf("err calling Next: %v", err)
return nil, false
}
got = append(got, vals)
}
return got, true
}
func TestNoMoreValues(t *testing.T) {
c := &Client{projectID: "project-id"}
pf := &pageFetcherReadStub{
values: [][][]Value{{{1, 2}, {11, 12}}},
}
it := c.Dataset("dataset-id").Table("table-id").read(context.Background(), pf.fetchPage)
var vals []Value
// We expect to retrieve two values and then fail on the next attempt.
if err := it.Next(&vals); err != nil {
t.Fatalf("Next: got: %v: want: nil", err)
}
if err := it.Next(&vals); err != nil {
t.Fatalf("Next: got: %v: want: nil", err)
}
if err := it.Next(&vals); err != iterator.Done {
t.Fatalf("Next: got: %v: want: iterator.Done", err)
}
}
var errBang = errors.New("bang")
func errorFetchPage(context.Context, *Table, Schema, uint64, int64, string) (*fetchPageResult, error) {
return nil, errBang
}
func TestReadError(t *testing.T) {
// test that service read errors are propagated back to the caller.
c := &Client{projectID: "project-id"}
it := c.Dataset("dataset-id").Table("table-id").read(context.Background(), errorFetchPage)
var vals []Value
if err := it.Next(&vals); err != errBang {
t.Fatalf("Get: got: %v: want: %v", err, errBang)
}
}
func TestReadTabledataOptions(t *testing.T) {
// test that read options are propagated.
s := &pageFetcherReadStub{
values: [][][]Value{{{1, 2}}},
}
c := &Client{projectID: "project-id"}
tr := c.Dataset("dataset-id").Table("table-id")
it := tr.read(context.Background(), s.fetchPage)
it.PageInfo().MaxSize = 5
var vals []Value
if err := it.Next(&vals); err != nil {
t.Fatal(err)
}
want := []pageFetcherArgs{{
table: tr,
pageSize: 5,
pageToken: "",
}}
if diff := testutil.Diff(s.calls, want, cmp.AllowUnexported(pageFetcherArgs{}, pageFetcherReadStub{}, Table{}, Client{})); diff != "" {
t.Errorf("reading (got=-, want=+):\n%s", diff)
}
}
func TestReadQueryOptions(t *testing.T) {
// test that read options are propagated.
c := &Client{projectID: "project-id"}
pf := &pageFetcherReadStub{
values: [][][]Value{{{1, 2}}},
}
tr := &bq.TableReference{
ProjectId: "project-id",
DatasetId: "dataset-id",
TableId: "table-id",
}
queryJob := &Job{
projectID: "project-id",
jobID: "job-id",
c: c,
config: &bq.JobConfiguration{
Query: &bq.JobConfigurationQuery{DestinationTable: tr},
},
}
it, err := queryJob.read(context.Background(), waitForQueryStub, pf.fetchPage)
if err != nil {
t.Fatalf("err calling Read: %v", err)
}
it.PageInfo().MaxSize = 5
var vals []Value
if err := it.Next(&vals); err != nil {
t.Fatalf("Next: got: %v: want: nil", err)
}
want := []pageFetcherArgs{{
table: bqToTable(tr, c),
pageSize: 5,
pageToken: "",
}}
if !testutil.Equal(pf.calls, want, cmp.AllowUnexported(pageFetcherArgs{}, Table{}, Client{})) {
t.Errorf("reading: got:\n%v\nwant:\n%v", pf.calls, want)
}
}
// Copyright 2015 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package bigquery
import (
"encoding/json"
"errors"
"fmt"
"reflect"
"sync"
bq "google.golang.org/api/bigquery/v2"
)
// Schema describes the fields in a table or query result.
type Schema []*FieldSchema
// FieldSchema describes a single field.
type FieldSchema struct {
// The field name.
// Must contain only letters (a-z, A-Z), numbers (0-9), or underscores (_),
// and must start with a letter or underscore.
// The maximum length is 128 characters.
Name string
// A description of the field. The maximum length is 16,384 characters.
Description string
// Whether the field may contain multiple values.
Repeated bool
// Whether the field is required. Ignored if Repeated is true.
Required bool
// The field data type. If Type is Record, then this field contains a nested schema,
// which is described by Schema.
Type FieldType
// Describes the nested schema if Type is set to Record.
Schema Schema
}
func (fs *FieldSchema) toBQ() *bq.TableFieldSchema {
tfs := &bq.TableFieldSchema{
Description: fs.Description,
Name: fs.Name,
Type: string(fs.Type),
}
if fs.Repeated {
tfs.Mode = "REPEATED"
} else if fs.Required {
tfs.Mode = "REQUIRED"
} // else leave as default, which is interpreted as NULLABLE.
for _, f := range fs.Schema {
tfs.Fields = append(tfs.Fields, f.toBQ())
}
return tfs
}
func (s Schema) toBQ() *bq.TableSchema {
var fields []*bq.TableFieldSchema
for _, f := range s {
fields = append(fields, f.toBQ())
}
return &bq.TableSchema{Fields: fields}
}
func bqToFieldSchema(tfs *bq.TableFieldSchema) *FieldSchema {
fs := &FieldSchema{
Description: tfs.Description,
Name: tfs.Name,
Repeated: tfs.Mode == "REPEATED",
Required: tfs.Mode == "REQUIRED",
Type: FieldType(tfs.Type),
}
for _, f := range tfs.Fields {
fs.Schema = append(fs.Schema, bqToFieldSchema(f))
}
return fs
}
func bqToSchema(ts *bq.TableSchema) Schema {
if ts == nil {
return nil
}
var s Schema
for _, f := range ts.Fields {
s = append(s, bqToFieldSchema(f))
}
return s
}
// FieldType is the type of field.
type FieldType string
const (
// StringFieldType is a string field type.
StringFieldType FieldType = "STRING"
// BytesFieldType is a bytes field type.
BytesFieldType FieldType = "BYTES"
// IntegerFieldType is a integer field type.
IntegerFieldType FieldType = "INTEGER"
// FloatFieldType is a float field type.
FloatFieldType FieldType = "FLOAT"
// BooleanFieldType is a boolean field type.
BooleanFieldType FieldType = "BOOLEAN"
// TimestampFieldType is a timestamp field type.
TimestampFieldType FieldType = "TIMESTAMP"
// RecordFieldType is a record field type. It is typically used to create columns with repeated or nested data.
RecordFieldType FieldType = "RECORD"
// DateFieldType is a date field type.
DateFieldType FieldType = "DATE"
// TimeFieldType is a time field type.
TimeFieldType FieldType = "TIME"
// DateTimeFieldType is a datetime field type.
DateTimeFieldType FieldType = "DATETIME"
// NumericFieldType is a numeric field type. Numeric types include integer types, floating point types and the
// NUMERIC data type.
NumericFieldType FieldType = "NUMERIC"
)
var (
errNoStruct = errors.New("bigquery: can only infer schema from struct or pointer to struct")
errUnsupportedFieldType = errors.New("bigquery: unsupported type of field in struct")
errInvalidFieldName = errors.New("bigquery: invalid name of field in struct")
errBadNullable = errors.New(`bigquery: use "nullable" only for []byte and struct pointers; for all other types, use a NullXXX type`)
errEmptyJSONSchema = errors.New("bigquery: empty JSON schema")
fieldTypes = map[FieldType]bool{
StringFieldType: true,
BytesFieldType: true,
IntegerFieldType: true,
FloatFieldType: true,
BooleanFieldType: true,
TimestampFieldType: true,
RecordFieldType: true,
DateFieldType: true,
TimeFieldType: true,
DateTimeFieldType: true,
NumericFieldType: true,
}
)
var typeOfByteSlice = reflect.TypeOf([]byte{})
// InferSchema tries to derive a BigQuery schema from the supplied struct value.
// Each exported struct field is mapped to a field in the schema.
//
// The following BigQuery types are inferred from the corresponding Go types.
// (This is the same mapping as that used for RowIterator.Next.) Fields inferred
// from these types are marked required (non-nullable).
//
// STRING string
// BOOL bool
// INTEGER int, int8, int16, int32, int64, uint8, uint16, uint32
// FLOAT float32, float64
// BYTES []byte
// TIMESTAMP time.Time
// DATE civil.Date
// TIME civil.Time
// DATETIME civil.DateTime
// NUMERIC *big.Rat
//
// The big.Rat type supports numbers of arbitrary size and precision. Values
// will be rounded to 9 digits after the decimal point before being transmitted
// to BigQuery. See https://cloud.google.com/bigquery/docs/reference/standard-sql/data-types#numeric-type
// for more on NUMERIC.
//
// A Go slice or array type is inferred to be a BigQuery repeated field of the
// element type. The element type must be one of the above listed types.
//
// Nullable fields are inferred from the NullXXX types, declared in this package:
//
// STRING NullString
// BOOL NullBool
// INTEGER NullInt64
// FLOAT NullFloat64
// TIMESTAMP NullTimestamp
// DATE NullDate
// TIME NullTime
// DATETIME NullDateTime
//
// For a nullable BYTES field, use the type []byte and tag the field "nullable" (see below).
// For a nullable NUMERIC field, use the type *big.Rat and tag the field "nullable".
//
// A struct field that is of struct type is inferred to be a required field of type
// RECORD with a schema inferred recursively. For backwards compatibility, a field of
// type pointer to struct is also inferred to be required. To get a nullable RECORD
// field, use the "nullable" tag (see below).
//
// InferSchema returns an error if any of the examined fields is of type uint,
// uint64, uintptr, map, interface, complex64, complex128, func, or chan. Future
// versions may handle these cases without error.
//
// Recursively defined structs are also disallowed.
//
// Struct fields may be tagged in a way similar to the encoding/json package.
// A tag of the form
// bigquery:"name"
// uses "name" instead of the struct field name as the BigQuery field name.
// A tag of the form
// bigquery:"-"
// omits the field from the inferred schema.
// The "nullable" option marks the field as nullable (not required). It is only
// needed for []byte, *big.Rat and pointer-to-struct fields, and cannot appear on other
// fields. In this example, the Go name of the field is retained:
// bigquery:",nullable"
func InferSchema(st interface{}) (Schema, error) {
return inferSchemaReflectCached(reflect.TypeOf(st))
}
var schemaCache sync.Map
type cacheVal struct {
schema Schema
err error
}
func inferSchemaReflectCached(t reflect.Type) (Schema, error) {
var cv cacheVal
v, ok := schemaCache.Load(t)
if ok {
cv = v.(cacheVal)
} else {
s, err := inferSchemaReflect(t)
cv = cacheVal{s, err}
schemaCache.Store(t, cv)
}
return cv.schema, cv.err
}
func inferSchemaReflect(t reflect.Type) (Schema, error) {
rec, err := hasRecursiveType(t, nil)
if err != nil {
return nil, err
}
if rec {
return nil, fmt.Errorf("bigquery: schema inference for recursive type %s", t)
}
return inferStruct(t)
}
func inferStruct(t reflect.Type) (Schema, error) {
switch t.Kind() {
case reflect.Ptr:
if t.Elem().Kind() != reflect.Struct {
return nil, errNoStruct
}
t = t.Elem()
fallthrough
case reflect.Struct:
return inferFields(t)
default:
return nil, errNoStruct
}
}
// inferFieldSchema infers the FieldSchema for a Go type
func inferFieldSchema(rt reflect.Type, nullable bool) (*FieldSchema, error) {
// Only []byte and struct pointers can be tagged nullable.
if nullable && !(rt == typeOfByteSlice || rt.Kind() == reflect.Ptr && rt.Elem().Kind() == reflect.Struct) {
return nil, errBadNullable
}
switch rt {
case typeOfByteSlice:
return &FieldSchema{Required: !nullable, Type: BytesFieldType}, nil
case typeOfGoTime:
return &FieldSchema{Required: true, Type: TimestampFieldType}, nil
case typeOfDate:
return &FieldSchema{Required: true, Type: DateFieldType}, nil
case typeOfTime:
return &FieldSchema{Required: true, Type: TimeFieldType}, nil
case typeOfDateTime:
return &FieldSchema{Required: true, Type: DateTimeFieldType}, nil
case typeOfRat:
return &FieldSchema{Required: !nullable, Type: NumericFieldType}, nil
}
if ft := nullableFieldType(rt); ft != "" {
return &FieldSchema{Required: false, Type: ft}, nil
}
if isSupportedIntType(rt) || isSupportedUintType(rt) {
return &FieldSchema{Required: true, Type: IntegerFieldType}, nil
}
switch rt.Kind() {
case reflect.Slice, reflect.Array:
et := rt.Elem()
if et != typeOfByteSlice && (et.Kind() == reflect.Slice || et.Kind() == reflect.Array) {
// Multi dimensional slices/arrays are not supported by BigQuery
return nil, errUnsupportedFieldType
}
if nullableFieldType(et) != "" {
// Repeated nullable types are not supported by BigQuery.
return nil, errUnsupportedFieldType
}
f, err := inferFieldSchema(et, false)
if err != nil {
return nil, err
}
f.Repeated = true
f.Required = false
return f, nil
case reflect.Ptr:
if rt.Elem().Kind() != reflect.Struct {
return nil, errUnsupportedFieldType
}
fallthrough
case reflect.Struct:
nested, err := inferStruct(rt)
if err != nil {
return nil, err
}
return &FieldSchema{Required: !nullable, Type: RecordFieldType, Schema: nested}, nil
case reflect.String:
return &FieldSchema{Required: !nullable, Type: StringFieldType}, nil
case reflect.Bool:
return &FieldSchema{Required: !nullable, Type: BooleanFieldType}, nil
case reflect.Float32, reflect.Float64:
return &FieldSchema{Required: !nullable, Type: FloatFieldType}, nil
default:
return nil, errUnsupportedFieldType
}
}
// inferFields extracts all exported field types from struct type.
func inferFields(rt reflect.Type) (Schema, error) {
var s Schema
fields, err := fieldCache.Fields(rt)
if err != nil {
return nil, err
}
for _, field := range fields {
var nullable bool
for _, opt := range field.ParsedTag.([]string) {
if opt == nullableTagOption {
nullable = true
break
}
}
f, err := inferFieldSchema(field.Type, nullable)
if err != nil {
return nil, err
}
f.Name = field.Name
s = append(s, f)
}
return s, nil
}
// isSupportedIntType reports whether t is an int type that can be properly
// represented by the BigQuery INTEGER/INT64 type.
func isSupportedIntType(t reflect.Type) bool {
switch t.Kind() {
case reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64, reflect.Int:
return true
default:
return false
}
}
// isSupportedIntType reports whether t is a uint type that can be properly
// represented by the BigQuery INTEGER/INT64 type.
func isSupportedUintType(t reflect.Type) bool {
switch t.Kind() {
case reflect.Uint8, reflect.Uint16, reflect.Uint32:
return true
default:
return false
}
}
// typeList is a linked list of reflect.Types.
type typeList struct {
t reflect.Type
next *typeList
}
func (l *typeList) has(t reflect.Type) bool {
for l != nil {
if l.t == t {
return true
}
l = l.next
}
return false
}
// hasRecursiveType reports whether t or any type inside t refers to itself, directly or indirectly,
// via exported fields. (Schema inference ignores unexported fields.)
func hasRecursiveType(t reflect.Type, seen *typeList) (bool, error) {
for t.Kind() == reflect.Ptr || t.Kind() == reflect.Slice || t.Kind() == reflect.Array {
t = t.Elem()
}
if t.Kind() != reflect.Struct {
return false, nil
}
if seen.has(t) {
return true, nil
}
fields, err := fieldCache.Fields(t)
if err != nil {
return false, err
}
seen = &typeList{t, seen}
// Because seen is a linked list, additions to it from one field's
// recursive call will not affect the value for subsequent fields' calls.
for _, field := range fields {
ok, err := hasRecursiveType(field.Type, seen)
if err != nil {
return false, err
}
if ok {
return true, nil
}
}
return false, nil
}
// bigQuerySchemaJSONField is an individual field in a JSON BigQuery table schema definition
// (as generated by https://github.com/GoogleCloudPlatform/protoc-gen-bq-schema).
type bigQueryJSONField struct {
Description string `json:"description"`
Fields []bigQueryJSONField `json:"fields"`
Mode string `json:"mode"`
Name string `json:"name"`
Type string `json:"type"`
}
// convertSchemaFromJSON generates a Schema:
func convertSchemaFromJSON(fs []bigQueryJSONField) (Schema, error) {
convertedSchema := Schema{}
for _, f := range fs {
convertedFieldSchema := &FieldSchema{
Description: f.Description,
Name: f.Name,
Required: f.Mode == "REQUIRED",
Repeated: f.Mode == "REPEATED",
}
if len(f.Fields) > 0 {
convertedNestedFieldSchema, err := convertSchemaFromJSON(f.Fields)
if err != nil {
return nil, err
}
convertedFieldSchema.Schema = convertedNestedFieldSchema
}
// Check that the field-type (string) maps to a known FieldType:
if _, ok := fieldTypes[FieldType(f.Type)]; !ok {
return nil, fmt.Errorf("unknown field type (%v)", f.Type)
}
convertedFieldSchema.Type = FieldType(f.Type)
convertedSchema = append(convertedSchema, convertedFieldSchema)
}
return convertedSchema, nil
}
// SchemaFromJSON takes a JSON BigQuery table schema definition
// (as generated by https://github.com/GoogleCloudPlatform/protoc-gen-bq-schema)
// and returns a fully-populated Schema.
func SchemaFromJSON(schemaJSON []byte) (Schema, error) {
var bigQuerySchema []bigQueryJSONField
// Make sure we actually have some content:
if len(schemaJSON) == 0 {
return nil, errEmptyJSONSchema
}
if err := json.Unmarshal(schemaJSON, &bigQuerySchema); err != nil {
return nil, err
}
return convertSchemaFromJSON(bigQuerySchema)
}
// Copyright 2015 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package bigquery
import (
"fmt"
"math/big"
"reflect"
"testing"
"time"
"cloud.google.com/go/civil"
"cloud.google.com/go/internal/pretty"
"cloud.google.com/go/internal/testutil"
bq "google.golang.org/api/bigquery/v2"
)
func (fs *FieldSchema) GoString() string {
if fs == nil {
return "<nil>"
}
return fmt.Sprintf("{Name:%s Description:%s Repeated:%t Required:%t Type:%s Schema:%s}",
fs.Name,
fs.Description,
fs.Repeated,
fs.Required,
fs.Type,
fmt.Sprintf("%#v", fs.Schema),
)
}
func bqTableFieldSchema(desc, name, typ, mode string) *bq.TableFieldSchema {
return &bq.TableFieldSchema{
Description: desc,
Name: name,
Mode: mode,
Type: typ,
}
}
func fieldSchema(desc, name, typ string, repeated, required bool) *FieldSchema {
return &FieldSchema{
Description: desc,
Name: name,
Repeated: repeated,
Required: required,
Type: FieldType(typ),
}
}
func TestSchemaConversion(t *testing.T) {
testCases := []struct {
schema Schema
bqSchema *bq.TableSchema
}{
{
// required
bqSchema: &bq.TableSchema{
Fields: []*bq.TableFieldSchema{
bqTableFieldSchema("desc", "name", "STRING", "REQUIRED"),
},
},
schema: Schema{
fieldSchema("desc", "name", "STRING", false, true),
},
},
{
// repeated
bqSchema: &bq.TableSchema{
Fields: []*bq.TableFieldSchema{
bqTableFieldSchema("desc", "name", "STRING", "REPEATED"),
},
},
schema: Schema{
fieldSchema("desc", "name", "STRING", true, false),
},
},
{
// nullable, string
bqSchema: &bq.TableSchema{
Fields: []*bq.TableFieldSchema{
bqTableFieldSchema("desc", "name", "STRING", ""),
},
},
schema: Schema{
fieldSchema("desc", "name", "STRING", false, false),
},
},
{
// integer
bqSchema: &bq.TableSchema{
Fields: []*bq.TableFieldSchema{
bqTableFieldSchema("desc", "name", "INTEGER", ""),
},
},
schema: Schema{
fieldSchema("desc", "name", "INTEGER", false, false),
},
},
{
// float
bqSchema: &bq.TableSchema{
Fields: []*bq.TableFieldSchema{
bqTableFieldSchema("desc", "name", "FLOAT", ""),
},
},
schema: Schema{
fieldSchema("desc", "name", "FLOAT", false, false),
},
},
{
// boolean
bqSchema: &bq.TableSchema{
Fields: []*bq.TableFieldSchema{
bqTableFieldSchema("desc", "name", "BOOLEAN", ""),
},
},
schema: Schema{
fieldSchema("desc", "name", "BOOLEAN", false, false),
},
},
{
// timestamp
bqSchema: &bq.TableSchema{
Fields: []*bq.TableFieldSchema{
bqTableFieldSchema("desc", "name", "TIMESTAMP", ""),
},
},
schema: Schema{
fieldSchema("desc", "name", "TIMESTAMP", false, false),
},
},
{
// civil times
bqSchema: &bq.TableSchema{
Fields: []*bq.TableFieldSchema{
bqTableFieldSchema("desc", "f1", "TIME", ""),
bqTableFieldSchema("desc", "f2", "DATE", ""),
bqTableFieldSchema("desc", "f3", "DATETIME", ""),
},
},
schema: Schema{
fieldSchema("desc", "f1", "TIME", false, false),
fieldSchema("desc", "f2", "DATE", false, false),
fieldSchema("desc", "f3", "DATETIME", false, false),
},
},
{
// numeric
bqSchema: &bq.TableSchema{
Fields: []*bq.TableFieldSchema{
bqTableFieldSchema("desc", "n", "NUMERIC", ""),
},
},
schema: Schema{
fieldSchema("desc", "n", "NUMERIC", false, false),
},
},
{
// nested
bqSchema: &bq.TableSchema{
Fields: []*bq.TableFieldSchema{
{
Description: "An outer schema wrapping a nested schema",
Name: "outer",
Mode: "REQUIRED",
Type: "RECORD",
Fields: []*bq.TableFieldSchema{
bqTableFieldSchema("inner field", "inner", "STRING", ""),
},
},
},
},
schema: Schema{
&FieldSchema{
Description: "An outer schema wrapping a nested schema",
Name: "outer",
Required: true,
Type: "RECORD",
Schema: Schema{
{
Description: "inner field",
Name: "inner",
Type: "STRING",
},
},
},
},
},
}
for _, tc := range testCases {
bqSchema := tc.schema.toBQ()
if !testutil.Equal(bqSchema, tc.bqSchema) {
t.Errorf("converting to TableSchema: got:\n%v\nwant:\n%v",
pretty.Value(bqSchema), pretty.Value(tc.bqSchema))
}
schema := bqToSchema(tc.bqSchema)
if !testutil.Equal(schema, tc.schema) {
t.Errorf("converting to Schema: got:\n%v\nwant:\n%v", schema, tc.schema)
}
}
}
type allStrings struct {
String string
ByteSlice []byte
}
type allSignedIntegers struct {
Int64 int64
Int32 int32
Int16 int16
Int8 int8
Int int
}
type allUnsignedIntegers struct {
Uint32 uint32
Uint16 uint16
Uint8 uint8
}
type allFloat struct {
Float64 float64
Float32 float32
// NOTE: Complex32 and Complex64 are unsupported by BigQuery
}
type allBoolean struct {
Bool bool
}
type allTime struct {
Timestamp time.Time
Time civil.Time
Date civil.Date
DateTime civil.DateTime
}
type allNumeric struct {
Numeric *big.Rat
}
func reqField(name, typ string) *FieldSchema {
return &FieldSchema{
Name: name,
Type: FieldType(typ),
Required: true,
}
}
func optField(name, typ string) *FieldSchema {
return &FieldSchema{
Name: name,
Type: FieldType(typ),
Required: false,
}
}
func TestSimpleInference(t *testing.T) {
testCases := []struct {
in interface{}
want Schema
}{
{
in: allSignedIntegers{},
want: Schema{
reqField("Int64", "INTEGER"),
reqField("Int32", "INTEGER"),
reqField("Int16", "INTEGER"),
reqField("Int8", "INTEGER"),
reqField("Int", "INTEGER"),
},
},
{
in: allUnsignedIntegers{},
want: Schema{
reqField("Uint32", "INTEGER"),
reqField("Uint16", "INTEGER"),
reqField("Uint8", "INTEGER"),
},
},
{
in: allFloat{},
want: Schema{
reqField("Float64", "FLOAT"),
reqField("Float32", "FLOAT"),
},
},
{
in: allBoolean{},
want: Schema{
reqField("Bool", "BOOLEAN"),
},
},
{
in: &allBoolean{},
want: Schema{
reqField("Bool", "BOOLEAN"),
},
},
{
in: allTime{},
want: Schema{
reqField("Timestamp", "TIMESTAMP"),
reqField("Time", "TIME"),
reqField("Date", "DATE"),
reqField("DateTime", "DATETIME"),
},
},
{
in: &allNumeric{},
want: Schema{
reqField("Numeric", "NUMERIC"),
},
},
{
in: allStrings{},
want: Schema{
reqField("String", "STRING"),
reqField("ByteSlice", "BYTES"),
},
},
}
for _, tc := range testCases {
got, err := InferSchema(tc.in)
if err != nil {
t.Fatalf("%T: error inferring TableSchema: %v", tc.in, err)
}
if !testutil.Equal(got, tc.want) {
t.Errorf("%T: inferring TableSchema: got:\n%#v\nwant:\n%#v", tc.in,
pretty.Value(got), pretty.Value(tc.want))
}
}
}
type containsNested struct {
hidden string
NotNested int
Nested struct {
Inside int
}
}
type containsDoubleNested struct {
NotNested int
Nested struct {
InsideNested struct {
Inside int
}
}
}
type ptrNested struct {
Ptr *struct{ Inside int }
}
type dup struct { // more than one field of the same struct type
A, B allBoolean
}
func TestNestedInference(t *testing.T) {
testCases := []struct {
in interface{}
want Schema
}{
{
in: containsNested{},
want: Schema{
reqField("NotNested", "INTEGER"),
&FieldSchema{
Name: "Nested",
Required: true,
Type: "RECORD",
Schema: Schema{reqField("Inside", "INTEGER")},
},
},
},
{
in: containsDoubleNested{},
want: Schema{
reqField("NotNested", "INTEGER"),
&FieldSchema{
Name: "Nested",
Required: true,
Type: "RECORD",
Schema: Schema{
{
Name: "InsideNested",
Required: true,
Type: "RECORD",
Schema: Schema{reqField("Inside", "INTEGER")},
},
},
},
},
},
{
in: ptrNested{},
want: Schema{
&FieldSchema{
Name: "Ptr",
Required: true,
Type: "RECORD",
Schema: Schema{reqField("Inside", "INTEGER")},
},
},
},
{
in: dup{},
want: Schema{
&FieldSchema{
Name: "A",
Required: true,
Type: "RECORD",
Schema: Schema{reqField("Bool", "BOOLEAN")},
},
&FieldSchema{
Name: "B",
Required: true,
Type: "RECORD",
Schema: Schema{reqField("Bool", "BOOLEAN")},
},
},
},
}
for _, tc := range testCases {
got, err := InferSchema(tc.in)
if err != nil {
t.Fatalf("%T: error inferring TableSchema: %v", tc.in, err)
}
if !testutil.Equal(got, tc.want) {
t.Errorf("%T: inferring TableSchema: got:\n%#v\nwant:\n%#v", tc.in,
pretty.Value(got), pretty.Value(tc.want))
}
}
}
type repeated struct {
NotRepeated []byte
RepeatedByteSlice [][]byte
Slice []int
Array [5]bool
}
type nestedRepeated struct {
NotRepeated int
Repeated []struct {
Inside int
}
RepeatedPtr []*struct{ Inside int }
}
func repField(name, typ string) *FieldSchema {
return &FieldSchema{
Name: name,
Type: FieldType(typ),
Repeated: true,
}
}
func TestRepeatedInference(t *testing.T) {
testCases := []struct {
in interface{}
want Schema
}{
{
in: repeated{},
want: Schema{
reqField("NotRepeated", "BYTES"),
repField("RepeatedByteSlice", "BYTES"),
repField("Slice", "INTEGER"),
repField("Array", "BOOLEAN"),
},
},
{
in: nestedRepeated{},
want: Schema{
reqField("NotRepeated", "INTEGER"),
{
Name: "Repeated",
Repeated: true,
Type: "RECORD",
Schema: Schema{reqField("Inside", "INTEGER")},
},
{
Name: "RepeatedPtr",
Repeated: true,
Type: "RECORD",
Schema: Schema{reqField("Inside", "INTEGER")},
},
},
},
}
for i, tc := range testCases {
got, err := InferSchema(tc.in)
if err != nil {
t.Fatalf("%d: error inferring TableSchema: %v", i, err)
}
if !testutil.Equal(got, tc.want) {
t.Errorf("%d: inferring TableSchema: got:\n%#v\nwant:\n%#v", i,
pretty.Value(got), pretty.Value(tc.want))
}
}
}
type allNulls struct {
A NullInt64
B NullFloat64
C NullBool
D NullString
E NullTimestamp
F NullTime
G NullDate
H NullDateTime
}
func TestNullInference(t *testing.T) {
got, err := InferSchema(allNulls{})
if err != nil {
t.Fatal(err)
}
want := Schema{
optField("A", "INTEGER"),
optField("B", "FLOAT"),
optField("C", "BOOLEAN"),
optField("D", "STRING"),
optField("E", "TIMESTAMP"),
optField("F", "TIME"),
optField("G", "DATE"),
optField("H", "DATETIME"),
}
if diff := testutil.Diff(got, want); diff != "" {
t.Error(diff)
}
}
type Embedded struct {
Embedded int
}
type embedded struct {
Embedded2 int
}
type nestedEmbedded struct {
Embedded
embedded
}
func TestEmbeddedInference(t *testing.T) {
got, err := InferSchema(nestedEmbedded{})
if err != nil {
t.Fatal(err)
}
want := Schema{
reqField("Embedded", "INTEGER"),
reqField("Embedded2", "INTEGER"),
}
if !testutil.Equal(got, want) {
t.Errorf("got %v, want %v", pretty.Value(got), pretty.Value(want))
}
}
func TestRecursiveInference(t *testing.T) {
type List struct {
Val int
Next *List
}
_, err := InferSchema(List{})
if err == nil {
t.Fatal("got nil, want error")
}
}
type withTags struct {
NoTag int
ExcludeTag int `bigquery:"-"`
SimpleTag int `bigquery:"simple_tag"`
UnderscoreTag int `bigquery:"_id"`
MixedCase int `bigquery:"MIXEDcase"`
Nullable []byte `bigquery:",nullable"`
NullNumeric *big.Rat `bigquery:",nullable"`
}
type withTagsNested struct {
Nested withTags `bigquery:"nested"`
NestedAnonymous struct {
ExcludeTag int `bigquery:"-"`
Inside int `bigquery:"inside"`
} `bigquery:"anon"`
PNested *struct{ X int } // not nullable, for backwards compatibility
PNestedNullable *struct{ X int } `bigquery:",nullable"`
}
type withTagsRepeated struct {
Repeated []withTags `bigquery:"repeated"`
RepeatedAnonymous []struct {
ExcludeTag int `bigquery:"-"`
Inside int `bigquery:"inside"`
} `bigquery:"anon"`
}
type withTagsEmbedded struct {
withTags
}
var withTagsSchema = Schema{
reqField("NoTag", "INTEGER"),
reqField("simple_tag", "INTEGER"),
reqField("_id", "INTEGER"),
reqField("MIXEDcase", "INTEGER"),
optField("Nullable", "BYTES"),
optField("NullNumeric", "NUMERIC"),
}
func TestTagInference(t *testing.T) {
testCases := []struct {
in interface{}
want Schema
}{
{
in: withTags{},
want: withTagsSchema,
},
{
in: withTagsNested{},
want: Schema{
&FieldSchema{
Name: "nested",
Required: true,
Type: "RECORD",
Schema: withTagsSchema,
},
&FieldSchema{
Name: "anon",
Required: true,
Type: "RECORD",
Schema: Schema{reqField("inside", "INTEGER")},
},
&FieldSchema{
Name: "PNested",
Required: true,
Type: "RECORD",
Schema: Schema{reqField("X", "INTEGER")},
},
&FieldSchema{
Name: "PNestedNullable",
Required: false,
Type: "RECORD",
Schema: Schema{reqField("X", "INTEGER")},
},
},
},
{
in: withTagsRepeated{},
want: Schema{
&FieldSchema{
Name: "repeated",
Repeated: true,
Type: "RECORD",
Schema: withTagsSchema,
},
&FieldSchema{
Name: "anon",
Repeated: true,
Type: "RECORD",
Schema: Schema{reqField("inside", "INTEGER")},
},
},
},
{
in: withTagsEmbedded{},
want: withTagsSchema,
},
}
for i, tc := range testCases {
got, err := InferSchema(tc.in)
if err != nil {
t.Fatalf("%d: error inferring TableSchema: %v", i, err)
}
if !testutil.Equal(got, tc.want) {
t.Errorf("%d: inferring TableSchema: got:\n%#v\nwant:\n%#v", i,
pretty.Value(got), pretty.Value(tc.want))
}
}
}
func TestTagInferenceErrors(t *testing.T) {
testCases := []struct {
in interface{}
err error
}{
{
in: struct {
LongTag int `bigquery:"abcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxy"`
}{},
err: errInvalidFieldName,
},
{
in: struct {
UnsupporedStartChar int `bigquery:"øab"`
}{},
err: errInvalidFieldName,
},
{
in: struct {
UnsupportedEndChar int `bigquery:"abø"`
}{},
err: errInvalidFieldName,
},
{
in: struct {
UnsupportedMiddleChar int `bigquery:"aøb"`
}{},
err: errInvalidFieldName,
},
{
in: struct {
StartInt int `bigquery:"1abc"`
}{},
err: errInvalidFieldName,
},
{
in: struct {
Hyphens int `bigquery:"a-b"`
}{},
err: errInvalidFieldName,
},
}
for i, tc := range testCases {
want := tc.err
_, got := InferSchema(tc.in)
if got != want {
t.Errorf("%d: inferring TableSchema: got:\n%#v\nwant:\n%#v", i, got, want)
}
}
_, err := InferSchema(struct {
X int `bigquery:",optional"`
}{})
if err == nil {
t.Error("got nil, want error")
}
}
func TestSchemaErrors(t *testing.T) {
testCases := []struct {
in interface{}
err error
}{
{
in: []byte{},
err: errNoStruct,
},
{
in: new(int),
err: errNoStruct,
},
{
in: struct{ Uint uint }{},
err: errUnsupportedFieldType,
},
{
in: struct{ Uint64 uint64 }{},
err: errUnsupportedFieldType,
},
{
in: struct{ Uintptr uintptr }{},
err: errUnsupportedFieldType,
},
{
in: struct{ Complex complex64 }{},
err: errUnsupportedFieldType,
},
{
in: struct{ Map map[string]int }{},
err: errUnsupportedFieldType,
},
{
in: struct{ Chan chan bool }{},
err: errUnsupportedFieldType,
},
{
in: struct{ Ptr *int }{},
err: errUnsupportedFieldType,
},
{
in: struct{ Interface interface{} }{},
err: errUnsupportedFieldType,
},
{
in: struct{ MultiDimensional [][]int }{},
err: errUnsupportedFieldType,
},
{
in: struct{ MultiDimensional [][][]byte }{},
err: errUnsupportedFieldType,
},
{
in: struct{ SliceOfPointer []*int }{},
err: errUnsupportedFieldType,
},
{
in: struct{ SliceOfNull []NullInt64 }{},
err: errUnsupportedFieldType,
},
{
in: struct{ ChanSlice []chan bool }{},
err: errUnsupportedFieldType,
},
{
in: struct{ NestedChan struct{ Chan []chan bool } }{},
err: errUnsupportedFieldType,
},
{
in: struct {
X int `bigquery:",nullable"`
}{},
err: errBadNullable,
},
{
in: struct {
X bool `bigquery:",nullable"`
}{},
err: errBadNullable,
},
{
in: struct {
X struct{ N int } `bigquery:",nullable"`
}{},
err: errBadNullable,
},
{
in: struct {
X []int `bigquery:",nullable"`
}{},
err: errBadNullable,
},
{
in: struct{ X *[]byte }{},
err: errUnsupportedFieldType,
},
{
in: struct{ X *[]int }{},
err: errUnsupportedFieldType,
},
{
in: struct{ X *int }{},
err: errUnsupportedFieldType,
},
}
for _, tc := range testCases {
want := tc.err
_, got := InferSchema(tc.in)
if got != want {
t.Errorf("%#v: got:\n%#v\nwant:\n%#v", tc.in, got, want)
}
}
}
func TestHasRecursiveType(t *testing.T) {
type (
nonStruct int
nonRec struct{ A string }
dup struct{ A, B nonRec }
rec struct {
A int
B *rec
}
recUnexported struct {
A int
b *rec
}
hasRec struct {
A int
R *rec
}
recSlicePointer struct {
A []*recSlicePointer
}
)
for _, test := range []struct {
in interface{}
want bool
}{
{nonStruct(0), false},
{nonRec{}, false},
{dup{}, false},
{rec{}, true},
{recUnexported{}, false},
{hasRec{}, true},
{&recSlicePointer{}, true},
} {
got, err := hasRecursiveType(reflect.TypeOf(test.in), nil)
if err != nil {
t.Fatal(err)
}
if got != test.want {
t.Errorf("%T: got %t, want %t", test.in, got, test.want)
}
}
}
func TestSchemaFromJSON(t *testing.T) {
testCasesExpectingSuccess := []struct {
bqSchemaJSON []byte
description string
expectedSchema Schema
}{
{
description: "Flat table with a mixture of NULLABLE and REQUIRED fields",
bqSchemaJSON: []byte(`
[
{"name":"flat_string","type":"STRING","mode":"NULLABLE","description":"Flat nullable string"},
{"name":"flat_bytes","type":"BYTES","mode":"REQUIRED","description":"Flat required BYTES"},
{"name":"flat_integer","type":"INTEGER","mode":"NULLABLE","description":"Flat nullable INTEGER"},
{"name":"flat_float","type":"FLOAT","mode":"REQUIRED","description":"Flat required FLOAT"},
{"name":"flat_boolean","type":"BOOLEAN","mode":"NULLABLE","description":"Flat nullable BOOLEAN"},
{"name":"flat_timestamp","type":"TIMESTAMP","mode":"REQUIRED","description":"Flat required TIMESTAMP"},
{"name":"flat_date","type":"DATE","mode":"NULLABLE","description":"Flat required DATE"},
{"name":"flat_time","type":"TIME","mode":"REQUIRED","description":"Flat nullable TIME"},
{"name":"flat_datetime","type":"DATETIME","mode":"NULLABLE","description":"Flat required DATETIME"},
{"name":"flat_numeric","type":"NUMERIC","mode":"REQUIRED","description":"Flat nullable NUMERIC"}
]`),
expectedSchema: Schema{
fieldSchema("Flat nullable string", "flat_string", "STRING", false, false),
fieldSchema("Flat required BYTES", "flat_bytes", "BYTES", false, true),
fieldSchema("Flat nullable INTEGER", "flat_integer", "INTEGER", false, false),
fieldSchema("Flat required FLOAT", "flat_float", "FLOAT", false, true),
fieldSchema("Flat nullable BOOLEAN", "flat_boolean", "BOOLEAN", false, false),
fieldSchema("Flat required TIMESTAMP", "flat_timestamp", "TIMESTAMP", false, true),
fieldSchema("Flat required DATE", "flat_date", "DATE", false, false),
fieldSchema("Flat nullable TIME", "flat_time", "TIME", false, true),
fieldSchema("Flat required DATETIME", "flat_datetime", "DATETIME", false, false),
fieldSchema("Flat nullable NUMERIC", "flat_numeric", "NUMERIC", false, true),
},
},
{
description: "Table with a nested RECORD",
bqSchemaJSON: []byte(`
[
{"name":"flat_string","type":"STRING","mode":"NULLABLE","description":"Flat nullable string"},
{"name":"nested_record","type":"RECORD","mode":"NULLABLE","description":"Nested nullable RECORD","fields":[{"name":"record_field_1","type":"STRING","mode":"NULLABLE","description":"First nested record field"},{"name":"record_field_2","type":"INTEGER","mode":"REQUIRED","description":"Second nested record field"}]}
]`),
expectedSchema: Schema{
fieldSchema("Flat nullable string", "flat_string", "STRING", false, false),
&FieldSchema{
Description: "Nested nullable RECORD",
Name: "nested_record",
Required: false,
Type: "RECORD",
Schema: Schema{
{
Description: "First nested record field",
Name: "record_field_1",
Required: false,
Type: "STRING",
},
{
Description: "Second nested record field",
Name: "record_field_2",
Required: true,
Type: "INTEGER",
},
},
},
},
},
{
description: "Table with a repeated RECORD",
bqSchemaJSON: []byte(`
[
{"name":"flat_string","type":"STRING","mode":"NULLABLE","description":"Flat nullable string"},
{"name":"nested_record","type":"RECORD","mode":"REPEATED","description":"Nested nullable RECORD","fields":[{"name":"record_field_1","type":"STRING","mode":"NULLABLE","description":"First nested record field"},{"name":"record_field_2","type":"INTEGER","mode":"REQUIRED","description":"Second nested record field"}]}
]`),
expectedSchema: Schema{
fieldSchema("Flat nullable string", "flat_string", "STRING", false, false),
&FieldSchema{
Description: "Nested nullable RECORD",
Name: "nested_record",
Repeated: true,
Required: false,
Type: "RECORD",
Schema: Schema{
{
Description: "First nested record field",
Name: "record_field_1",
Required: false,
Type: "STRING",
},
{
Description: "Second nested record field",
Name: "record_field_2",
Required: true,
Type: "INTEGER",
},
},
},
},
},
}
for _, tc := range testCasesExpectingSuccess {
convertedSchema, err := SchemaFromJSON(tc.bqSchemaJSON)
if err != nil {
t.Errorf("encountered an error when converting JSON table schema (%s): %v", tc.description, err)
continue
}
if !testutil.Equal(convertedSchema, tc.expectedSchema) {
t.Errorf("generated JSON table schema (%s) differs from the expected schema", tc.description)
}
}
testCasesExpectingFailure := []struct {
bqSchemaJSON []byte
description string
}{
{
description: "Schema with invalid JSON",
bqSchemaJSON: []byte(`This is not JSON`),
},
{
description: "Schema with unknown field type",
bqSchemaJSON: []byte(`[{"name":"strange_type","type":"STRANGE","description":"This type should not exist"}]`),
},
{
description: "Schema with zero length",
bqSchemaJSON: []byte(``),
},
}
for _, tc := range testCasesExpectingFailure {
_, err := SchemaFromJSON(tc.bqSchemaJSON)
if err == nil {
t.Errorf("converting this schema should have returned an error (%s): %v", tc.description, err)
continue
}
}
}
// Copyright 2015 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package bigquery
import (
"context"
"errors"
"fmt"
"time"
"cloud.google.com/go/internal/optional"
"cloud.google.com/go/internal/trace"
bq "google.golang.org/api/bigquery/v2"
)
// A Table is a reference to a BigQuery table.
type Table struct {
// ProjectID, DatasetID and TableID may be omitted if the Table is the destination for a query.
// In this case the result will be stored in an ephemeral table.
ProjectID string
DatasetID string
// TableID must contain only letters (a-z, A-Z), numbers (0-9), or underscores (_).
// The maximum length is 1,024 characters.
TableID string
c *Client
}
// TableMetadata contains information about a BigQuery table.
type TableMetadata struct {
// The following fields can be set when creating a table.
// The user-friendly name for the table.
Name string
// The user-friendly description of the table.
Description string
// The table schema. If provided on create, ViewQuery must be empty.
Schema Schema
// The query to use for a view. If provided on create, Schema must be nil.
ViewQuery string
// Use Legacy SQL for the view query.
// At most one of UseLegacySQL and UseStandardSQL can be true.
UseLegacySQL bool
// Use Legacy SQL for the view query. The default.
// At most one of UseLegacySQL and UseStandardSQL can be true.
// Deprecated: use UseLegacySQL.
UseStandardSQL bool
// If non-nil, the table is partitioned by time.
TimePartitioning *TimePartitioning
// Clustering specifies the data clustering configuration for the table.
Clustering *Clustering
// The time when this table expires. If set, this table will expire at the
// specified time. Expired tables will be deleted and their storage
// reclaimed. The zero value is ignored.
ExpirationTime time.Time
// User-provided labels.
Labels map[string]string
// Information about a table stored outside of BigQuery.
ExternalDataConfig *ExternalDataConfig
// Custom encryption configuration (e.g., Cloud KMS keys).
EncryptionConfig *EncryptionConfig
// All the fields below are read-only.
FullID string // An opaque ID uniquely identifying the table.
Type TableType
CreationTime time.Time
LastModifiedTime time.Time
// The size of the table in bytes.
// This does not include data that is being buffered during a streaming insert.
NumBytes int64
// The number of rows of data in this table.
// This does not include data that is being buffered during a streaming insert.
NumRows uint64
// Contains information regarding this table's streaming buffer, if one is
// present. This field will be nil if the table is not being streamed to or if
// there is no data in the streaming buffer.
StreamingBuffer *StreamingBuffer
// ETag is the ETag obtained when reading metadata. Pass it to Table.Update to
// ensure that the metadata hasn't changed since it was read.
ETag string
}
// TableCreateDisposition specifies the circumstances under which destination table will be created.
// Default is CreateIfNeeded.
type TableCreateDisposition string
const (
// CreateIfNeeded will create the table if it does not already exist.
// Tables are created atomically on successful completion of a job.
CreateIfNeeded TableCreateDisposition = "CREATE_IF_NEEDED"
// CreateNever ensures the table must already exist and will not be
// automatically created.
CreateNever TableCreateDisposition = "CREATE_NEVER"
)
// TableWriteDisposition specifies how existing data in a destination table is treated.
// Default is WriteAppend.
type TableWriteDisposition string
const (
// WriteAppend will append to any existing data in the destination table.
// Data is appended atomically on successful completion of a job.
WriteAppend TableWriteDisposition = "WRITE_APPEND"
// WriteTruncate overrides the existing data in the destination table.
// Data is overwritten atomically on successful completion of a job.
WriteTruncate TableWriteDisposition = "WRITE_TRUNCATE"
// WriteEmpty fails writes if the destination table already contains data.
WriteEmpty TableWriteDisposition = "WRITE_EMPTY"
)
// TableType is the type of table.
type TableType string
const (
// RegularTable is a regular table.
RegularTable TableType = "TABLE"
// ViewTable is a table type describing that the table is view. See more
// information at https://cloud.google.com/bigquery/docs/views.
ViewTable TableType = "VIEW"
// ExternalTable is a table type describing that the table is an external
// table (also known as a federated data source). See more information at
// https://cloud.google.com/bigquery/external-data-sources.
ExternalTable TableType = "EXTERNAL"
)
// TimePartitioning describes the time-based date partitioning on a table.
// For more information see: https://cloud.google.com/bigquery/docs/creating-partitioned-tables.
type TimePartitioning struct {
// The amount of time to keep the storage for a partition.
// If the duration is empty (0), the data in the partitions do not expire.
Expiration time.Duration
// If empty, the table is partitioned by pseudo column '_PARTITIONTIME'; if set, the
// table is partitioned by this field. The field must be a top-level TIMESTAMP or
// DATE field. Its mode must be NULLABLE or REQUIRED.
Field string
// If true, queries that reference this table must include a filter (e.g. a WHERE predicate)
// that can be used for partition elimination.
RequirePartitionFilter bool
}
func (p *TimePartitioning) toBQ() *bq.TimePartitioning {
if p == nil {
return nil
}
return &bq.TimePartitioning{
Type: "DAY",
ExpirationMs: int64(p.Expiration / time.Millisecond),
Field: p.Field,
RequirePartitionFilter: p.RequirePartitionFilter,
}
}
func bqToTimePartitioning(q *bq.TimePartitioning) *TimePartitioning {
if q == nil {
return nil
}
return &TimePartitioning{
Expiration: time.Duration(q.ExpirationMs) * time.Millisecond,
Field: q.Field,
RequirePartitionFilter: q.RequirePartitionFilter,
}
}
// Clustering governs the organization of data within a partitioned table.
// For more information, see https://cloud.google.com/bigquery/docs/clustered-tables
type Clustering struct {
Fields []string
}
func (c *Clustering) toBQ() *bq.Clustering {
if c == nil {
return nil
}
return &bq.Clustering{
Fields: c.Fields,
}
}
func bqToClustering(q *bq.Clustering) *Clustering {
if q == nil {
return nil
}
return &Clustering{
Fields: q.Fields,
}
}
// EncryptionConfig configures customer-managed encryption on tables.
type EncryptionConfig struct {
// Describes the Cloud KMS encryption key that will be used to protect
// destination BigQuery table. The BigQuery Service Account associated with your
// project requires access to this encryption key.
KMSKeyName string
}
func (e *EncryptionConfig) toBQ() *bq.EncryptionConfiguration {
if e == nil {
return nil
}
return &bq.EncryptionConfiguration{
KmsKeyName: e.KMSKeyName,
}
}
func bqToEncryptionConfig(q *bq.EncryptionConfiguration) *EncryptionConfig {
if q == nil {
return nil
}
return &EncryptionConfig{
KMSKeyName: q.KmsKeyName,
}
}
// StreamingBuffer holds information about the streaming buffer.
type StreamingBuffer struct {
// A lower-bound estimate of the number of bytes currently in the streaming
// buffer.
EstimatedBytes uint64
// A lower-bound estimate of the number of rows currently in the streaming
// buffer.
EstimatedRows uint64
// The time of the oldest entry in the streaming buffer.
OldestEntryTime time.Time
}
func (t *Table) toBQ() *bq.TableReference {
return &bq.TableReference{
ProjectId: t.ProjectID,
DatasetId: t.DatasetID,
TableId: t.TableID,
}
}
// FullyQualifiedName returns the ID of the table in projectID:datasetID.tableID format.
func (t *Table) FullyQualifiedName() string {
return fmt.Sprintf("%s:%s.%s", t.ProjectID, t.DatasetID, t.TableID)
}
// implicitTable reports whether Table is an empty placeholder, which signifies that a new table should be created with an auto-generated Table ID.
func (t *Table) implicitTable() bool {
return t.ProjectID == "" && t.DatasetID == "" && t.TableID == ""
}
// Create creates a table in the BigQuery service.
// Pass in a TableMetadata value to configure the table.
// If tm.View.Query is non-empty, the created table will be of type VIEW.
// If no ExpirationTime is specified, the table will never expire.
// After table creation, a view can be modified only if its table was initially created
// with a view.
func (t *Table) Create(ctx context.Context, tm *TableMetadata) (err error) {
ctx = trace.StartSpan(ctx, "cloud.google.com/go/bigquery.Table.Create")
defer func() { trace.EndSpan(ctx, err) }()
table, err := tm.toBQ()
if err != nil {
return err
}
table.TableReference = &bq.TableReference{
ProjectId: t.ProjectID,
DatasetId: t.DatasetID,
TableId: t.TableID,
}
req := t.c.bqs.Tables.Insert(t.ProjectID, t.DatasetID, table).Context(ctx)
setClientHeader(req.Header())
_, err = req.Do()
return err
}
func (tm *TableMetadata) toBQ() (*bq.Table, error) {
t := &bq.Table{}
if tm == nil {
return t, nil
}
if tm.Schema != nil && tm.ViewQuery != "" {
return nil, errors.New("bigquery: provide Schema or ViewQuery, not both")
}
t.FriendlyName = tm.Name
t.Description = tm.Description
t.Labels = tm.Labels
if tm.Schema != nil {
t.Schema = tm.Schema.toBQ()
}
if tm.ViewQuery != "" {
if tm.UseStandardSQL && tm.UseLegacySQL {
return nil, errors.New("bigquery: cannot provide both UseStandardSQL and UseLegacySQL")
}
t.View = &bq.ViewDefinition{Query: tm.ViewQuery}
if tm.UseLegacySQL {
t.View.UseLegacySql = true
} else {
t.View.UseLegacySql = false
t.View.ForceSendFields = append(t.View.ForceSendFields, "UseLegacySql")
}
} else if tm.UseLegacySQL || tm.UseStandardSQL {
return nil, errors.New("bigquery: UseLegacy/StandardSQL requires ViewQuery")
}
t.TimePartitioning = tm.TimePartitioning.toBQ()
t.Clustering = tm.Clustering.toBQ()
if !validExpiration(tm.ExpirationTime) {
return nil, fmt.Errorf("invalid expiration time: %v.\n"+
"Valid expiration times are after 1678 and before 2262", tm.ExpirationTime)
}
if !tm.ExpirationTime.IsZero() && tm.ExpirationTime != NeverExpire {
t.ExpirationTime = tm.ExpirationTime.UnixNano() / 1e6
}
if tm.ExternalDataConfig != nil {
edc := tm.ExternalDataConfig.toBQ()
t.ExternalDataConfiguration = &edc
}
t.EncryptionConfiguration = tm.EncryptionConfig.toBQ()
if tm.FullID != "" {
return nil, errors.New("cannot set FullID on create")
}
if tm.Type != "" {
return nil, errors.New("cannot set Type on create")
}
if !tm.CreationTime.IsZero() {
return nil, errors.New("cannot set CreationTime on create")
}
if !tm.LastModifiedTime.IsZero() {
return nil, errors.New("cannot set LastModifiedTime on create")
}
if tm.NumBytes != 0 {
return nil, errors.New("cannot set NumBytes on create")
}
if tm.NumRows != 0 {
return nil, errors.New("cannot set NumRows on create")
}
if tm.StreamingBuffer != nil {
return nil, errors.New("cannot set StreamingBuffer on create")
}
if tm.ETag != "" {
return nil, errors.New("cannot set ETag on create")
}
return t, nil
}
// Metadata fetches the metadata for the table.
func (t *Table) Metadata(ctx context.Context) (md *TableMetadata, err error) {
ctx = trace.StartSpan(ctx, "cloud.google.com/go/bigquery.Table.Metadata")
defer func() { trace.EndSpan(ctx, err) }()
req := t.c.bqs.Tables.Get(t.ProjectID, t.DatasetID, t.TableID).Context(ctx)
setClientHeader(req.Header())
var table *bq.Table
err = runWithRetry(ctx, func() (err error) {
table, err = req.Do()
return err
})
if err != nil {
return nil, err
}
return bqToTableMetadata(table)
}
func bqToTableMetadata(t *bq.Table) (*TableMetadata, error) {
md := &TableMetadata{
Description: t.Description,
Name: t.FriendlyName,
Type: TableType(t.Type),
FullID: t.Id,
Labels: t.Labels,
NumBytes: t.NumBytes,
NumRows: t.NumRows,
ExpirationTime: unixMillisToTime(t.ExpirationTime),
CreationTime: unixMillisToTime(t.CreationTime),
LastModifiedTime: unixMillisToTime(int64(t.LastModifiedTime)),
ETag: t.Etag,
EncryptionConfig: bqToEncryptionConfig(t.EncryptionConfiguration),
}
if t.Schema != nil {
md.Schema = bqToSchema(t.Schema)
}
if t.View != nil {
md.ViewQuery = t.View.Query
md.UseLegacySQL = t.View.UseLegacySql
}
md.TimePartitioning = bqToTimePartitioning(t.TimePartitioning)
md.Clustering = bqToClustering(t.Clustering)
if t.StreamingBuffer != nil {
md.StreamingBuffer = &StreamingBuffer{
EstimatedBytes: t.StreamingBuffer.EstimatedBytes,
EstimatedRows: t.StreamingBuffer.EstimatedRows,
OldestEntryTime: unixMillisToTime(int64(t.StreamingBuffer.OldestEntryTime)),
}
}
if t.ExternalDataConfiguration != nil {
edc, err := bqToExternalDataConfig(t.ExternalDataConfiguration)
if err != nil {
return nil, err
}
md.ExternalDataConfig = edc
}
return md, nil
}
// Delete deletes the table.
func (t *Table) Delete(ctx context.Context) (err error) {
ctx = trace.StartSpan(ctx, "cloud.google.com/go/bigquery.Table.Delete")
defer func() { trace.EndSpan(ctx, err) }()
req := t.c.bqs.Tables.Delete(t.ProjectID, t.DatasetID, t.TableID).Context(ctx)
setClientHeader(req.Header())
return req.Do()
}
// Read fetches the contents of the table.
func (t *Table) Read(ctx context.Context) *RowIterator {
return t.read(ctx, fetchPage)
}
func (t *Table) read(ctx context.Context, pf pageFetcher) *RowIterator {
return newRowIterator(ctx, t, pf)
}
// NeverExpire is a sentinel value used to remove a table'e expiration time.
var NeverExpire = time.Time{}.Add(-1)
// Update modifies specific Table metadata fields.
func (t *Table) Update(ctx context.Context, tm TableMetadataToUpdate, etag string) (md *TableMetadata, err error) {
ctx = trace.StartSpan(ctx, "cloud.google.com/go/bigquery.Table.Update")
defer func() { trace.EndSpan(ctx, err) }()
bqt, err := tm.toBQ()
if err != nil {
return nil, err
}
call := t.c.bqs.Tables.Patch(t.ProjectID, t.DatasetID, t.TableID, bqt).Context(ctx)
setClientHeader(call.Header())
if etag != "" {
call.Header().Set("If-Match", etag)
}
var res *bq.Table
if err := runWithRetry(ctx, func() (err error) {
res, err = call.Do()
return err
}); err != nil {
return nil, err
}
return bqToTableMetadata(res)
}
func (tm *TableMetadataToUpdate) toBQ() (*bq.Table, error) {
t := &bq.Table{}
forceSend := func(field string) {
t.ForceSendFields = append(t.ForceSendFields, field)
}
if tm.Description != nil {
t.Description = optional.ToString(tm.Description)
forceSend("Description")
}
if tm.Name != nil {
t.FriendlyName = optional.ToString(tm.Name)
forceSend("FriendlyName")
}
if tm.Schema != nil {
t.Schema = tm.Schema.toBQ()
forceSend("Schema")
}
if tm.EncryptionConfig != nil {
t.EncryptionConfiguration = tm.EncryptionConfig.toBQ()
}
if !validExpiration(tm.ExpirationTime) {
return nil, fmt.Errorf("invalid expiration time: %v.\n"+
"Valid expiration times are after 1678 and before 2262", tm.ExpirationTime)
}
if tm.ExpirationTime == NeverExpire {
t.NullFields = append(t.NullFields, "ExpirationTime")
} else if !tm.ExpirationTime.IsZero() {
t.ExpirationTime = tm.ExpirationTime.UnixNano() / 1e6
forceSend("ExpirationTime")
}
if tm.TimePartitioning != nil {
t.TimePartitioning = tm.TimePartitioning.toBQ()
t.TimePartitioning.ForceSendFields = []string{"Expiration", "RequirePartitionFilter"}
}
if tm.ViewQuery != nil {
t.View = &bq.ViewDefinition{
Query: optional.ToString(tm.ViewQuery),
ForceSendFields: []string{"Query"},
}
}
if tm.UseLegacySQL != nil {
if t.View == nil {
t.View = &bq.ViewDefinition{}
}
t.View.UseLegacySql = optional.ToBool(tm.UseLegacySQL)
t.View.ForceSendFields = append(t.View.ForceSendFields, "UseLegacySql")
}
labels, forces, nulls := tm.update()
t.Labels = labels
t.ForceSendFields = append(t.ForceSendFields, forces...)
t.NullFields = append(t.NullFields, nulls...)
return t, nil
}
// validExpiration ensures a specified time is either the sentinel NeverExpire,
// the zero value, or within the defined range of UnixNano. Internal
// represetations of expiration times are based upon Time.UnixNano. Any time
// before 1678 or after 2262 cannot be represented by an int64 and is therefore
// undefined and invalid. See https://godoc.org/time#Time.UnixNano.
func validExpiration(t time.Time) bool {
return t == NeverExpire || t.IsZero() || time.Unix(0, t.UnixNano()).Equal(t)
}
// TableMetadataToUpdate is used when updating a table's metadata.
// Only non-nil fields will be updated.
type TableMetadataToUpdate struct {
// The user-friendly description of this table.
Description optional.String
// The user-friendly name for this table.
Name optional.String
// The table's schema.
// When updating a schema, you can add columns but not remove them.
Schema Schema
// The table's encryption configuration. When calling Update, ensure that
// all mutable fields of EncryptionConfig are populated.
EncryptionConfig *EncryptionConfig
// The time when this table expires. To remove a table's expiration,
// set ExpirationTime to NeverExpire. The zero value is ignored.
ExpirationTime time.Time
// The query to use for a view.
ViewQuery optional.String
// Use Legacy SQL for the view query.
UseLegacySQL optional.Bool
// TimePartitioning allows modification of certain aspects of partition
// configuration such as partition expiration and whether partition
// filtration is required at query time. When calling Update, ensure
// that all mutable fields of TimePartitioning are populated.
TimePartitioning *TimePartitioning
labelUpdater
}
// labelUpdater contains common code for updating labels.
type labelUpdater struct {
setLabels map[string]string
deleteLabels map[string]bool
}
// SetLabel causes a label to be added or modified on a call to Update.
func (u *labelUpdater) SetLabel(name, value string) {
if u.setLabels == nil {
u.setLabels = map[string]string{}
}
u.setLabels[name] = value
}
// DeleteLabel causes a label to be deleted on a call to Update.
func (u *labelUpdater) DeleteLabel(name string) {
if u.deleteLabels == nil {
u.deleteLabels = map[string]bool{}
}
u.deleteLabels[name] = true
}
func (u *labelUpdater) update() (labels map[string]string, forces, nulls []string) {
if u.setLabels == nil && u.deleteLabels == nil {
return nil, nil, nil
}
labels = map[string]string{}
for k, v := range u.setLabels {
labels[k] = v
}
if len(labels) == 0 && len(u.deleteLabels) > 0 {
forces = []string{"Labels"}
}
for l := range u.deleteLabels {
nulls = append(nulls, "Labels."+l)
}
return labels, forces, nulls
}
// Copyright 2017 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package bigquery
import (
"testing"
"time"
"cloud.google.com/go/internal/testutil"
bq "google.golang.org/api/bigquery/v2"
)
func TestBQToTableMetadata(t *testing.T) {
aTime := time.Date(2017, 1, 26, 0, 0, 0, 0, time.Local)
aTimeMillis := aTime.UnixNano() / 1e6
for _, test := range []struct {
in *bq.Table
want *TableMetadata
}{
{&bq.Table{}, &TableMetadata{}}, // test minimal case
{
&bq.Table{
CreationTime: aTimeMillis,
Description: "desc",
Etag: "etag",
ExpirationTime: aTimeMillis,
FriendlyName: "fname",
Id: "id",
LastModifiedTime: uint64(aTimeMillis),
Location: "loc",
NumBytes: 123,
NumLongTermBytes: 23,
NumRows: 7,
StreamingBuffer: &bq.Streamingbuffer{
EstimatedBytes: 11,
EstimatedRows: 3,
OldestEntryTime: uint64(aTimeMillis),
},
TimePartitioning: &bq.TimePartitioning{
ExpirationMs: 7890,
Type: "DAY",
Field: "pfield",
},
Clustering: &bq.Clustering{
Fields: []string{"cfield1", "cfield2"},
},
EncryptionConfiguration: &bq.EncryptionConfiguration{KmsKeyName: "keyName"},
Type: "EXTERNAL",
View: &bq.ViewDefinition{Query: "view-query"},
Labels: map[string]string{"a": "b"},
ExternalDataConfiguration: &bq.ExternalDataConfiguration{
SourceFormat: "GOOGLE_SHEETS",
},
},
&TableMetadata{
Description: "desc",
Name: "fname",
ViewQuery: "view-query",
FullID: "id",
Type: ExternalTable,
Labels: map[string]string{"a": "b"},
ExternalDataConfig: &ExternalDataConfig{SourceFormat: GoogleSheets},
ExpirationTime: aTime.Truncate(time.Millisecond),
CreationTime: aTime.Truncate(time.Millisecond),
LastModifiedTime: aTime.Truncate(time.Millisecond),
NumBytes: 123,
NumRows: 7,
TimePartitioning: &TimePartitioning{
Expiration: 7890 * time.Millisecond,
Field: "pfield",
},
Clustering: &Clustering{
Fields: []string{"cfield1", "cfield2"},
},
StreamingBuffer: &StreamingBuffer{
EstimatedBytes: 11,
EstimatedRows: 3,
OldestEntryTime: aTime,
},
EncryptionConfig: &EncryptionConfig{KMSKeyName: "keyName"},
ETag: "etag",
},
},
} {
got, err := bqToTableMetadata(test.in)
if err != nil {
t.Fatal(err)
}
if diff := testutil.Diff(got, test.want); diff != "" {
t.Errorf("%+v:\n, -got, +want:\n%s", test.in, diff)
}
}
}
func TestTableMetadataToBQ(t *testing.T) {
aTime := time.Date(2017, 1, 26, 0, 0, 0, 0, time.Local)
aTimeMillis := aTime.UnixNano() / 1e6
sc := Schema{fieldSchema("desc", "name", "STRING", false, true)}
for _, test := range []struct {
in *TableMetadata
want *bq.Table
}{
{nil, &bq.Table{}},
{&TableMetadata{}, &bq.Table{}},
{
&TableMetadata{
Name: "n",
Description: "d",
Schema: sc,
ExpirationTime: aTime,
Labels: map[string]string{"a": "b"},
ExternalDataConfig: &ExternalDataConfig{SourceFormat: Bigtable},
EncryptionConfig: &EncryptionConfig{KMSKeyName: "keyName"},
},
&bq.Table{
FriendlyName: "n",
Description: "d",
Schema: &bq.TableSchema{
Fields: []*bq.TableFieldSchema{
bqTableFieldSchema("desc", "name", "STRING", "REQUIRED"),
},
},
ExpirationTime: aTimeMillis,
Labels: map[string]string{"a": "b"},
ExternalDataConfiguration: &bq.ExternalDataConfiguration{SourceFormat: "BIGTABLE"},
EncryptionConfiguration: &bq.EncryptionConfiguration{KmsKeyName: "keyName"},
},
},
{
&TableMetadata{ViewQuery: "q"},
&bq.Table{
View: &bq.ViewDefinition{
Query: "q",
UseLegacySql: false,
ForceSendFields: []string{"UseLegacySql"},
},
},
},
{
&TableMetadata{
ViewQuery: "q",
UseLegacySQL: true,
TimePartitioning: &TimePartitioning{},
},
&bq.Table{
View: &bq.ViewDefinition{
Query: "q",
UseLegacySql: true,
},
TimePartitioning: &bq.TimePartitioning{
Type: "DAY",
ExpirationMs: 0,
},
},
},
{
&TableMetadata{
ViewQuery: "q",
UseStandardSQL: true,
TimePartitioning: &TimePartitioning{
Expiration: time.Second,
Field: "ofDreams",
},
Clustering: &Clustering{
Fields: []string{"cfield1"},
},
},
&bq.Table{
View: &bq.ViewDefinition{
Query: "q",
UseLegacySql: false,
ForceSendFields: []string{"UseLegacySql"},
},
TimePartitioning: &bq.TimePartitioning{
Type: "DAY",
ExpirationMs: 1000,
Field: "ofDreams",
},
Clustering: &bq.Clustering{
Fields: []string{"cfield1"},
},
},
},
{
&TableMetadata{ExpirationTime: NeverExpire},
&bq.Table{ExpirationTime: 0},
},
} {
got, err := test.in.toBQ()
if err != nil {
t.Fatalf("%+v: %v", test.in, err)
}
if diff := testutil.Diff(got, test.want); diff != "" {
t.Errorf("%+v:\n-got, +want:\n%s", test.in, diff)
}
}
// Errors
for _, in := range []*TableMetadata{
{Schema: sc, ViewQuery: "q"}, // can't have both schema and query
{UseLegacySQL: true}, // UseLegacySQL without query
{UseStandardSQL: true}, // UseStandardSQL without query
// read-only fields
{FullID: "x"},
{Type: "x"},
{CreationTime: aTime},
{LastModifiedTime: aTime},
{NumBytes: 1},
{NumRows: 1},
{StreamingBuffer: &StreamingBuffer{}},
{ETag: "x"},
// expiration time outside allowable range is invalid
// See https://godoc.org/time#Time.UnixNano
{ExpirationTime: time.Date(1677, 9, 21, 0, 12, 43, 145224192, time.UTC).Add(-1)},
{ExpirationTime: time.Date(2262, 04, 11, 23, 47, 16, 854775807, time.UTC).Add(1)},
} {
_, err := in.toBQ()
if err == nil {
t.Errorf("%+v: got nil, want error", in)
}
}
}
func TestTableMetadataToUpdateToBQ(t *testing.T) {
aTime := time.Date(2017, 1, 26, 0, 0, 0, 0, time.Local)
for _, test := range []struct {
tm TableMetadataToUpdate
want *bq.Table
}{
{
tm: TableMetadataToUpdate{},
want: &bq.Table{},
},
{
tm: TableMetadataToUpdate{
Description: "d",
Name: "n",
},
want: &bq.Table{
Description: "d",
FriendlyName: "n",
ForceSendFields: []string{"Description", "FriendlyName"},
},
},
{
tm: TableMetadataToUpdate{
Schema: Schema{fieldSchema("desc", "name", "STRING", false, true)},
ExpirationTime: aTime,
},
want: &bq.Table{
Schema: &bq.TableSchema{
Fields: []*bq.TableFieldSchema{
bqTableFieldSchema("desc", "name", "STRING", "REQUIRED"),
},
},
ExpirationTime: aTime.UnixNano() / 1e6,
ForceSendFields: []string{"Schema", "ExpirationTime"},
},
},
{
tm: TableMetadataToUpdate{ViewQuery: "q"},
want: &bq.Table{
View: &bq.ViewDefinition{Query: "q", ForceSendFields: []string{"Query"}},
},
},
{
tm: TableMetadataToUpdate{UseLegacySQL: false},
want: &bq.Table{
View: &bq.ViewDefinition{
UseLegacySql: false,
ForceSendFields: []string{"UseLegacySql"},
},
},
},
{
tm: TableMetadataToUpdate{ViewQuery: "q", UseLegacySQL: true},
want: &bq.Table{
View: &bq.ViewDefinition{
Query: "q",
UseLegacySql: true,
ForceSendFields: []string{"Query", "UseLegacySql"},
},
},
},
{
tm: func() (tm TableMetadataToUpdate) {
tm.SetLabel("L", "V")
tm.DeleteLabel("D")
return tm
}(),
want: &bq.Table{
Labels: map[string]string{"L": "V"},
NullFields: []string{"Labels.D"},
},
},
{
tm: TableMetadataToUpdate{ExpirationTime: NeverExpire},
want: &bq.Table{
NullFields: []string{"ExpirationTime"},
},
},
} {
got, _ := test.tm.toBQ()
if !testutil.Equal(got, test.want) {
t.Errorf("%+v:\ngot %+v\nwant %+v", test.tm, got, test.want)
}
}
}
func TestTableMetadataToUpdateToBQErrors(t *testing.T) {
// See https://godoc.org/time#Time.UnixNano
start := time.Date(1677, 9, 21, 0, 12, 43, 145224192, time.UTC)
end := time.Date(2262, 04, 11, 23, 47, 16, 854775807, time.UTC)
for _, test := range []struct {
desc string
aTime time.Time
wantErr bool
}{
{desc: "ignored zero value", aTime: time.Time{}, wantErr: false},
{desc: "earliest valid time", aTime: start, wantErr: false},
{desc: "latested valid time", aTime: end, wantErr: false},
{desc: "invalid times before 1678", aTime: start.Add(-1), wantErr: true},
{desc: "invalid times after 2262", aTime: end.Add(1), wantErr: true},
{desc: "valid times after 1678", aTime: start.Add(1), wantErr: false},
{desc: "valid times before 2262", aTime: end.Add(-1), wantErr: false},
} {
tm := &TableMetadataToUpdate{ExpirationTime: test.aTime}
_, err := tm.toBQ()
if test.wantErr && err == nil {
t.Errorf("[%s] got no error, want error", test.desc)
}
if !test.wantErr && err != nil {
t.Errorf("[%s] got error, want no error", test.desc)
}
}
}
// Copyright 2015 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package bigquery
import (
"encoding/base64"
"errors"
"fmt"
"math"
"math/big"
"reflect"
"strconv"
"strings"
"time"
"cloud.google.com/go/civil"
bq "google.golang.org/api/bigquery/v2"
)
// Value stores the contents of a single cell from a BigQuery result.
type Value interface{}
// ValueLoader stores a slice of Values representing a result row from a Read operation.
// See RowIterator.Next for more information.
type ValueLoader interface {
Load(v []Value, s Schema) error
}
// valueList converts a []Value to implement ValueLoader.
type valueList []Value
// Load stores a sequence of values in a valueList.
// It resets the slice length to zero, then appends each value to it.
func (vs *valueList) Load(v []Value, _ Schema) error {
*vs = append((*vs)[:0], v...)
return nil
}
// valueMap converts a map[string]Value to implement ValueLoader.
type valueMap map[string]Value
// Load stores a sequence of values in a valueMap.
func (vm *valueMap) Load(v []Value, s Schema) error {
if *vm == nil {
*vm = map[string]Value{}
}
loadMap(*vm, v, s)
return nil
}
func loadMap(m map[string]Value, vals []Value, s Schema) {
for i, f := range s {
val := vals[i]
var v interface{}
switch {
case val == nil:
v = val
case f.Schema == nil:
v = val
case !f.Repeated:
m2 := map[string]Value{}
loadMap(m2, val.([]Value), f.Schema)
v = m2
default: // repeated and nested
sval := val.([]Value)
vs := make([]Value, len(sval))
for j, e := range sval {
m2 := map[string]Value{}
loadMap(m2, e.([]Value), f.Schema)
vs[j] = m2
}
v = vs
}
m[f.Name] = v
}
}
type structLoader struct {
typ reflect.Type // type of struct
err error
ops []structLoaderOp
vstructp reflect.Value // pointer to current struct value; changed by set
}
// A setFunc is a function that sets a struct field or slice/array
// element to a value.
type setFunc func(v reflect.Value, val interface{}) error
// A structLoaderOp instructs the loader to set a struct field to a row value.
type structLoaderOp struct {
fieldIndex []int
valueIndex int
setFunc setFunc
repeated bool
}
var errNoNulls = errors.New("bigquery: NULL values cannot be read into structs")
func setAny(v reflect.Value, x interface{}) error {
if x == nil {
return errNoNulls
}
v.Set(reflect.ValueOf(x))
return nil
}
func setInt(v reflect.Value, x interface{}) error {
if x == nil {
return errNoNulls
}
xx := x.(int64)
if v.OverflowInt(xx) {
return fmt.Errorf("bigquery: value %v overflows struct field of type %v", xx, v.Type())
}
v.SetInt(xx)
return nil
}
func setUint(v reflect.Value, x interface{}) error {
if x == nil {
return errNoNulls
}
xx := x.(int64)
if xx < 0 || v.OverflowUint(uint64(xx)) {
return fmt.Errorf("bigquery: value %v overflows struct field of type %v", xx, v.Type())
}
v.SetUint(uint64(xx))
return nil
}
func setFloat(v reflect.Value, x interface{}) error {
if x == nil {
return errNoNulls
}
xx := x.(float64)
if v.OverflowFloat(xx) {
return fmt.Errorf("bigquery: value %v overflows struct field of type %v", xx, v.Type())
}
v.SetFloat(xx)
return nil
}
func setBool(v reflect.Value, x interface{}) error {
if x == nil {
return errNoNulls
}
v.SetBool(x.(bool))
return nil
}
func setString(v reflect.Value, x interface{}) error {
if x == nil {
return errNoNulls
}
v.SetString(x.(string))
return nil
}
func setBytes(v reflect.Value, x interface{}) error {
if x == nil {
v.SetBytes(nil)
} else {
v.SetBytes(x.([]byte))
}
return nil
}
func setNull(v reflect.Value, x interface{}, build func() interface{}) error {
if x == nil {
v.Set(reflect.Zero(v.Type()))
} else {
n := build()
v.Set(reflect.ValueOf(n))
}
return nil
}
// set remembers a value for the next call to Load. The value must be
// a pointer to a struct. (This is checked in RowIterator.Next.)
func (sl *structLoader) set(structp interface{}, schema Schema) error {
if sl.err != nil {
return sl.err
}
sl.vstructp = reflect.ValueOf(structp)
typ := sl.vstructp.Type().Elem()
if sl.typ == nil {
// First call: remember the type and compile the schema.
sl.typ = typ
ops, err := compileToOps(typ, schema)
if err != nil {
sl.err = err
return err
}
sl.ops = ops
} else if sl.typ != typ {
return fmt.Errorf("bigquery: struct type changed from %s to %s", sl.typ, typ)
}
return nil
}
// compileToOps produces a sequence of operations that will set the fields of a
// value of structType to the contents of a row with schema.
func compileToOps(structType reflect.Type, schema Schema) ([]structLoaderOp, error) {
var ops []structLoaderOp
fields, err := fieldCache.Fields(structType)
if err != nil {
return nil, err
}
for i, schemaField := range schema {
// Look for an exported struct field with the same name as the schema
// field, ignoring case (BigQuery column names are case-insensitive,
// and we want to act like encoding/json anyway).
structField := fields.Match(schemaField.Name)
if structField == nil {
// Ignore schema fields with no corresponding struct field.
continue
}
op := structLoaderOp{
fieldIndex: structField.Index,
valueIndex: i,
}
t := structField.Type
if schemaField.Repeated {
if t.Kind() != reflect.Slice && t.Kind() != reflect.Array {
return nil, fmt.Errorf("bigquery: repeated schema field %s requires slice or array, but struct field %s has type %s",
schemaField.Name, structField.Name, t)
}
t = t.Elem()
op.repeated = true
}
if schemaField.Type == RecordFieldType {
// Field can be a struct or a pointer to a struct.
if t.Kind() == reflect.Ptr {
t = t.Elem()
}
if t.Kind() != reflect.Struct {
return nil, fmt.Errorf("bigquery: field %s has type %s, expected struct or *struct",
structField.Name, structField.Type)
}
nested, err := compileToOps(t, schemaField.Schema)
if err != nil {
return nil, err
}
op.setFunc = func(v reflect.Value, val interface{}) error {
return setNested(nested, v, val)
}
} else {
op.setFunc = determineSetFunc(t, schemaField.Type)
if op.setFunc == nil {
return nil, fmt.Errorf("bigquery: schema field %s of type %s is not assignable to struct field %s of type %s",
schemaField.Name, schemaField.Type, structField.Name, t)
}
}
ops = append(ops, op)
}
return ops, nil
}
// determineSetFunc chooses the best function for setting a field of type ftype
// to a value whose schema field type is stype. It returns nil if stype
// is not assignable to ftype.
// determineSetFunc considers only basic types. See compileToOps for
// handling of repetition and nesting.
func determineSetFunc(ftype reflect.Type, stype FieldType) setFunc {
switch stype {
case StringFieldType:
if ftype.Kind() == reflect.String {
return setString
}
if ftype == typeOfNullString {
return func(v reflect.Value, x interface{}) error {
return setNull(v, x, func() interface{} {
return NullString{StringVal: x.(string), Valid: true}
})
}
}
case BytesFieldType:
if ftype == typeOfByteSlice {
return setBytes
}
case IntegerFieldType:
if isSupportedUintType(ftype) {
return setUint
} else if isSupportedIntType(ftype) {
return setInt
}
if ftype == typeOfNullInt64 {
return func(v reflect.Value, x interface{}) error {
return setNull(v, x, func() interface{} {
return NullInt64{Int64: x.(int64), Valid: true}
})
}
}
case FloatFieldType:
switch ftype.Kind() {
case reflect.Float32, reflect.Float64:
return setFloat
}
if ftype == typeOfNullFloat64 {
return func(v reflect.Value, x interface{}) error {
return setNull(v, x, func() interface{} {
return NullFloat64{Float64: x.(float64), Valid: true}
})
}
}
case BooleanFieldType:
if ftype.Kind() == reflect.Bool {
return setBool
}
if ftype == typeOfNullBool {
return func(v reflect.Value, x interface{}) error {
return setNull(v, x, func() interface{} {
return NullBool{Bool: x.(bool), Valid: true}
})
}
}
case TimestampFieldType:
if ftype == typeOfGoTime {
return setAny
}
if ftype == typeOfNullTimestamp {
return func(v reflect.Value, x interface{}) error {
return setNull(v, x, func() interface{} {
return NullTimestamp{Timestamp: x.(time.Time), Valid: true}
})
}
}
case DateFieldType:
if ftype == typeOfDate {
return setAny
}
if ftype == typeOfNullDate {
return func(v reflect.Value, x interface{}) error {
return setNull(v, x, func() interface{} {
return NullDate{Date: x.(civil.Date), Valid: true}
})
}
}
case TimeFieldType:
if ftype == typeOfTime {
return setAny
}
if ftype == typeOfNullTime {
return func(v reflect.Value, x interface{}) error {
return setNull(v, x, func() interface{} {
return NullTime{Time: x.(civil.Time), Valid: true}
})
}
}
case DateTimeFieldType:
if ftype == typeOfDateTime {
return setAny
}
if ftype == typeOfNullDateTime {
return func(v reflect.Value, x interface{}) error {
return setNull(v, x, func() interface{} {
return NullDateTime{DateTime: x.(civil.DateTime), Valid: true}
})
}
}
case NumericFieldType:
if ftype == typeOfRat {
return func(v reflect.Value, x interface{}) error {
return setNull(v, x, func() interface{} { return x.(*big.Rat) })
}
}
}
return nil
}
func (sl *structLoader) Load(values []Value, _ Schema) error {
if sl.err != nil {
return sl.err
}
return runOps(sl.ops, sl.vstructp.Elem(), values)
}
// runOps executes a sequence of ops, setting the fields of vstruct to the
// supplied values.
func runOps(ops []structLoaderOp, vstruct reflect.Value, values []Value) error {
for _, op := range ops {
field := vstruct.FieldByIndex(op.fieldIndex)
var err error
if op.repeated {
err = setRepeated(field, values[op.valueIndex].([]Value), op.setFunc)
} else {
err = op.setFunc(field, values[op.valueIndex])
}
if err != nil {
return err
}
}
return nil
}
func setNested(ops []structLoaderOp, v reflect.Value, val interface{}) error {
// v is either a struct or a pointer to a struct.
if v.Kind() == reflect.Ptr {
// If the value is nil, set the pointer to nil.
if val == nil {
v.Set(reflect.Zero(v.Type()))
return nil
}
// If the pointer is nil, set it to a zero struct value.
if v.IsNil() {
v.Set(reflect.New(v.Type().Elem()))
}
v = v.Elem()
}
return runOps(ops, v, val.([]Value))
}
func setRepeated(field reflect.Value, vslice []Value, setElem setFunc) error {
vlen := len(vslice)
var flen int
switch field.Type().Kind() {
case reflect.Slice:
// Make a slice of the right size, avoiding allocation if possible.
switch {
case field.Len() < vlen:
field.Set(reflect.MakeSlice(field.Type(), vlen, vlen))
case field.Len() > vlen:
field.SetLen(vlen)
}
flen = vlen
case reflect.Array:
flen = field.Len()
if flen > vlen {
// Set extra elements to their zero value.
z := reflect.Zero(field.Type().Elem())
for i := vlen; i < flen; i++ {
field.Index(i).Set(z)
}
}
default:
return fmt.Errorf("bigquery: impossible field type %s", field.Type())
}
for i, val := range vslice {
if i < flen { // avoid writing past the end of a short array
if err := setElem(field.Index(i), val); err != nil {
return err
}
}
}
return nil
}
// A ValueSaver returns a row of data to be inserted into a table.
type ValueSaver interface {
// Save returns a row to be inserted into a BigQuery table, represented
// as a map from field name to Value.
// If insertID is non-empty, BigQuery will use it to de-duplicate
// insertions of this row on a best-effort basis.
Save() (row map[string]Value, insertID string, err error)
}
// ValuesSaver implements ValueSaver for a slice of Values.
type ValuesSaver struct {
Schema Schema
// If non-empty, BigQuery will use InsertID to de-duplicate insertions
// of this row on a best-effort basis.
InsertID string
Row []Value
}
// Save implements ValueSaver.
func (vls *ValuesSaver) Save() (map[string]Value, string, error) {
m, err := valuesToMap(vls.Row, vls.Schema)
return m, vls.InsertID, err
}
func valuesToMap(vs []Value, schema Schema) (map[string]Value, error) {
if len(vs) != len(schema) {
return nil, errors.New("Schema does not match length of row to be inserted")
}
m := make(map[string]Value)
for i, fieldSchema := range schema {
if vs[i] == nil {
m[fieldSchema.Name] = nil
continue
}
if fieldSchema.Type != RecordFieldType {
m[fieldSchema.Name] = toUploadValue(vs[i], fieldSchema)
continue
}
// Nested record, possibly repeated.
vals, ok := vs[i].([]Value)
if !ok {
return nil, errors.New("nested record is not a []Value")
}
if !fieldSchema.Repeated {
value, err := valuesToMap(vals, fieldSchema.Schema)
if err != nil {
return nil, err
}
m[fieldSchema.Name] = value
continue
}
// A repeated nested field is converted into a slice of maps.
var maps []Value
for _, v := range vals {
sv, ok := v.([]Value)
if !ok {
return nil, errors.New("nested record in slice is not a []Value")
}
value, err := valuesToMap(sv, fieldSchema.Schema)
if err != nil {
return nil, err
}
maps = append(maps, value)
}
m[fieldSchema.Name] = maps
}
return m, nil
}
// StructSaver implements ValueSaver for a struct.
// The struct is converted to a map of values by using the values of struct
// fields corresponding to schema fields. Additional and missing
// fields are ignored, as are nested struct pointers that are nil.
type StructSaver struct {
// Schema determines what fields of the struct are uploaded. It should
// match the table's schema.
// Schema is optional for StructSavers that are passed to Uploader.Put.
Schema Schema
// If non-empty, BigQuery will use InsertID to de-duplicate insertions
// of this row on a best-effort basis.
InsertID string
// Struct should be a struct or a pointer to a struct.
Struct interface{}
}
// Save implements ValueSaver.
func (ss *StructSaver) Save() (row map[string]Value, insertID string, err error) {
vstruct := reflect.ValueOf(ss.Struct)
row, err = structToMap(vstruct, ss.Schema)
if err != nil {
return nil, "", err
}
return row, ss.InsertID, nil
}
func structToMap(vstruct reflect.Value, schema Schema) (map[string]Value, error) {
if vstruct.Kind() == reflect.Ptr {
vstruct = vstruct.Elem()
}
if !vstruct.IsValid() {
return nil, nil
}
m := map[string]Value{}
if vstruct.Kind() != reflect.Struct {
return nil, fmt.Errorf("bigquery: type is %s, need struct or struct pointer", vstruct.Type())
}
fields, err := fieldCache.Fields(vstruct.Type())
if err != nil {
return nil, err
}
for _, schemaField := range schema {
// Look for an exported struct field with the same name as the schema
// field, ignoring case.
structField := fields.Match(schemaField.Name)
if structField == nil {
continue
}
val, err := structFieldToUploadValue(vstruct.FieldByIndex(structField.Index), schemaField)
if err != nil {
return nil, err
}
// Add the value to the map, unless it is nil.
if val != nil {
m[schemaField.Name] = val
}
}
return m, nil
}
// structFieldToUploadValue converts a struct field to a value suitable for ValueSaver.Save, using
// the schemaField as a guide.
// structFieldToUploadValue is careful to return a true nil interface{} when needed, so its
// caller can easily identify a nil value.
func structFieldToUploadValue(vfield reflect.Value, schemaField *FieldSchema) (interface{}, error) {
if schemaField.Repeated && (vfield.Kind() != reflect.Slice && vfield.Kind() != reflect.Array) {
return nil, fmt.Errorf("bigquery: repeated schema field %s requires slice or array, but value has type %s",
schemaField.Name, vfield.Type())
}
// A non-nested field can be represented by its Go value, except for some types.
if schemaField.Type != RecordFieldType {
return toUploadValueReflect(vfield, schemaField), nil
}
// A non-repeated nested field is converted into a map[string]Value.
if !schemaField.Repeated {
m, err := structToMap(vfield, schemaField.Schema)
if err != nil {
return nil, err
}
if m == nil {
return nil, nil
}
return m, nil
}
// A repeated nested field is converted into a slice of maps.
// If the field is zero-length (but not nil), we return a zero-length []Value.
if vfield.IsNil() {
return nil, nil
}
vals := []Value{}
for i := 0; i < vfield.Len(); i++ {
m, err := structToMap(vfield.Index(i), schemaField.Schema)
if err != nil {
return nil, err
}
vals = append(vals, m)
}
return vals, nil
}
func toUploadValue(val interface{}, fs *FieldSchema) interface{} {
if fs.Type == TimeFieldType || fs.Type == DateTimeFieldType || fs.Type == NumericFieldType {
return toUploadValueReflect(reflect.ValueOf(val), fs)
}
return val
}
func toUploadValueReflect(v reflect.Value, fs *FieldSchema) interface{} {
switch fs.Type {
case TimeFieldType:
if v.Type() == typeOfNullTime {
return v.Interface()
}
return formatUploadValue(v, fs, func(v reflect.Value) string {
return CivilTimeString(v.Interface().(civil.Time))
})
case DateTimeFieldType:
if v.Type() == typeOfNullDateTime {
return v.Interface()
}
return formatUploadValue(v, fs, func(v reflect.Value) string {
return CivilDateTimeString(v.Interface().(civil.DateTime))
})
case NumericFieldType:
if r, ok := v.Interface().(*big.Rat); ok && r == nil {
return nil
}
return formatUploadValue(v, fs, func(v reflect.Value) string {
return NumericString(v.Interface().(*big.Rat))
})
default:
if !fs.Repeated || v.Len() > 0 {
return v.Interface()
}
// The service treats a null repeated field as an error. Return
// nil to omit the field entirely.
return nil
}
}
func formatUploadValue(v reflect.Value, fs *FieldSchema, cvt func(reflect.Value) string) interface{} {
if !fs.Repeated {
return cvt(v)
}
if v.Len() == 0 {
return nil
}
s := make([]string, v.Len())
for i := 0; i < v.Len(); i++ {
s[i] = cvt(v.Index(i))
}
return s
}
// CivilTimeString returns a string representing a civil.Time in a format compatible
// with BigQuery SQL. It rounds the time to the nearest microsecond and returns a
// string with six digits of sub-second precision.
//
// Use CivilTimeString when using civil.Time in DML, for example in INSERT
// statements.
func CivilTimeString(t civil.Time) string {
if t.Nanosecond == 0 {
return t.String()
}
micro := (t.Nanosecond + 500) / 1000 // round to nearest microsecond
t.Nanosecond = 0
return t.String() + fmt.Sprintf(".%06d", micro)
}
// CivilDateTimeString returns a string representing a civil.DateTime in a format compatible
// with BigQuery SQL. It separate the date and time with a space, and formats the time
// with CivilTimeString.
//
// Use CivilDateTimeString when using civil.DateTime in DML, for example in INSERT
// statements.
func CivilDateTimeString(dt civil.DateTime) string {
return dt.Date.String() + " " + CivilTimeString(dt.Time)
}
// parseCivilDateTime parses a date-time represented in a BigQuery SQL
// compatible format and returns a civil.DateTime.
func parseCivilDateTime(s string) (civil.DateTime, error) {
parts := strings.Fields(s)
if len(parts) != 2 {
return civil.DateTime{}, fmt.Errorf("bigquery: bad DATETIME value %q", s)
}
return civil.ParseDateTime(parts[0] + "T" + parts[1])
}
const (
// NumericPrecisionDigits is the maximum number of digits in a NUMERIC value.
NumericPrecisionDigits = 38
// NumericScaleDigits is the maximum number of digits after the decimal point in a NUMERIC value.
NumericScaleDigits = 9
)
// NumericString returns a string representing a *big.Rat in a format compatible
// with BigQuery SQL. It returns a floating-point literal with 9 digits
// after the decimal point.
func NumericString(r *big.Rat) string {
return r.FloatString(NumericScaleDigits)
}
// convertRows converts a series of TableRows into a series of Value slices.
// schema is used to interpret the data from rows; its length must match the
// length of each row.
func convertRows(rows []*bq.TableRow, schema Schema) ([][]Value, error) {
var rs [][]Value
for _, r := range rows {
row, err := convertRow(r, schema)
if err != nil {
return nil, err
}
rs = append(rs, row)
}
return rs, nil
}
func convertRow(r *bq.TableRow, schema Schema) ([]Value, error) {
if len(schema) != len(r.F) {
return nil, errors.New("schema length does not match row length")
}
var values []Value
for i, cell := range r.F {
fs := schema[i]
v, err := convertValue(cell.V, fs.Type, fs.Schema)
if err != nil {
return nil, err
}
values = append(values, v)
}
return values, nil
}
func convertValue(val interface{}, typ FieldType, schema Schema) (Value, error) {
switch val := val.(type) {
case nil:
return nil, nil
case []interface{}:
return convertRepeatedRecord(val, typ, schema)
case map[string]interface{}:
return convertNestedRecord(val, schema)
case string:
return convertBasicType(val, typ)
default:
return nil, fmt.Errorf("got value %v; expected a value of type %s", val, typ)
}
}
func convertRepeatedRecord(vals []interface{}, typ FieldType, schema Schema) (Value, error) {
var values []Value
for _, cell := range vals {
// each cell contains a single entry, keyed by "v"
val := cell.(map[string]interface{})["v"]
v, err := convertValue(val, typ, schema)
if err != nil {
return nil, err
}
values = append(values, v)
}
return values, nil
}
func convertNestedRecord(val map[string]interface{}, schema Schema) (Value, error) {
// convertNestedRecord is similar to convertRow, as a record has the same structure as a row.
// Nested records are wrapped in a map with a single key, "f".
record := val["f"].([]interface{})
if len(record) != len(schema) {
return nil, errors.New("schema length does not match record length")
}
var values []Value
for i, cell := range record {
// each cell contains a single entry, keyed by "v"
val := cell.(map[string]interface{})["v"]
fs := schema[i]
v, err := convertValue(val, fs.Type, fs.Schema)
if err != nil {
return nil, err
}
values = append(values, v)
}
return values, nil
}
// convertBasicType returns val as an interface with a concrete type specified by typ.
func convertBasicType(val string, typ FieldType) (Value, error) {
switch typ {
case StringFieldType:
return val, nil
case BytesFieldType:
return base64.StdEncoding.DecodeString(val)
case IntegerFieldType:
return strconv.ParseInt(val, 10, 64)
case FloatFieldType:
return strconv.ParseFloat(val, 64)
case BooleanFieldType:
return strconv.ParseBool(val)
case TimestampFieldType:
f, err := strconv.ParseFloat(val, 64)
if err != nil {
return nil, err
}
secs := math.Trunc(f)
nanos := (f - secs) * 1e9
return Value(time.Unix(int64(secs), int64(nanos)).UTC()), nil
case DateFieldType:
return civil.ParseDate(val)
case TimeFieldType:
return civil.ParseTime(val)
case DateTimeFieldType:
return civil.ParseDateTime(val)
case NumericFieldType:
r, ok := (&big.Rat{}).SetString(val)
if !ok {
return nil, fmt.Errorf("bigquery: invalid NUMERIC value %q", val)
}
return Value(r), nil
default:
return nil, fmt.Errorf("unrecognized type: %s", typ)
}
}
// Copyright 2015 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package bigquery
import (
"encoding/base64"
"fmt"
"math"
"math/big"
"testing"
"time"
"cloud.google.com/go/civil"
"cloud.google.com/go/internal/testutil"
"github.com/google/go-cmp/cmp"
bq "google.golang.org/api/bigquery/v2"
)
func TestConvertBasicValues(t *testing.T) {
schema := Schema{
{Type: StringFieldType},
{Type: IntegerFieldType},
{Type: FloatFieldType},
{Type: BooleanFieldType},
{Type: BytesFieldType},
{Type: NumericFieldType},
}
row := &bq.TableRow{
F: []*bq.TableCell{
{V: "a"},
{V: "1"},
{V: "1.2"},
{V: "true"},
{V: base64.StdEncoding.EncodeToString([]byte("foo"))},
{V: "123.123456789"},
},
}
got, err := convertRow(row, schema)
if err != nil {
t.Fatalf("error converting: %v", err)
}
want := []Value{"a", int64(1), 1.2, true, []byte("foo"), big.NewRat(123123456789, 1e9)}
if !testutil.Equal(got, want) {
t.Errorf("converting basic values: got:\n%v\nwant:\n%v", got, want)
}
}
func TestConvertTime(t *testing.T) {
schema := Schema{
{Type: TimestampFieldType},
{Type: DateFieldType},
{Type: TimeFieldType},
{Type: DateTimeFieldType},
}
ts := testTimestamp.Round(time.Millisecond)
row := &bq.TableRow{
F: []*bq.TableCell{
{V: fmt.Sprintf("%.10f", float64(ts.UnixNano())/1e9)},
{V: testDate.String()},
{V: testTime.String()},
{V: testDateTime.String()},
},
}
got, err := convertRow(row, schema)
if err != nil {
t.Fatalf("error converting: %v", err)
}
want := []Value{ts, testDate, testTime, testDateTime}
for i, g := range got {
w := want[i]
if !testutil.Equal(g, w) {
t.Errorf("#%d: got:\n%v\nwant:\n%v", i, g, w)
}
}
if got[0].(time.Time).Location() != time.UTC {
t.Errorf("expected time zone UTC: got:\n%v", got)
}
}
func TestConvertSmallTimes(t *testing.T) {
for _, year := range []int{1600, 1066, 1} {
want := time.Date(year, time.January, 1, 0, 0, 0, 0, time.UTC)
s := fmt.Sprintf("%.10f", float64(want.Unix()))
got, err := convertBasicType(s, TimestampFieldType)
if err != nil {
t.Fatal(err)
}
if !got.(time.Time).Equal(want) {
t.Errorf("got %v, want %v", got, want)
}
}
}
func TestConvertNullValues(t *testing.T) {
schema := Schema{{Type: StringFieldType}}
row := &bq.TableRow{
F: []*bq.TableCell{
{V: nil},
},
}
got, err := convertRow(row, schema)
if err != nil {
t.Fatalf("error converting: %v", err)
}
want := []Value{nil}
if !testutil.Equal(got, want) {
t.Errorf("converting null values: got:\n%v\nwant:\n%v", got, want)
}
}
func TestBasicRepetition(t *testing.T) {
schema := Schema{
{Type: IntegerFieldType, Repeated: true},
}
row := &bq.TableRow{
F: []*bq.TableCell{
{
V: []interface{}{
map[string]interface{}{
"v": "1",
},
map[string]interface{}{
"v": "2",
},
map[string]interface{}{
"v": "3",
},
},
},
},
}
got, err := convertRow(row, schema)
if err != nil {
t.Fatalf("error converting: %v", err)
}
want := []Value{[]Value{int64(1), int64(2), int64(3)}}
if !testutil.Equal(got, want) {
t.Errorf("converting basic repeated values: got:\n%v\nwant:\n%v", got, want)
}
}
func TestNestedRecordContainingRepetition(t *testing.T) {
schema := Schema{
{
Type: RecordFieldType,
Schema: Schema{
{Type: IntegerFieldType, Repeated: true},
},
},
}
row := &bq.TableRow{
F: []*bq.TableCell{
{
V: map[string]interface{}{
"f": []interface{}{
map[string]interface{}{
"v": []interface{}{
map[string]interface{}{"v": "1"},
map[string]interface{}{"v": "2"},
map[string]interface{}{"v": "3"},
},
},
},
},
},
},
}
got, err := convertRow(row, schema)
if err != nil {
t.Fatalf("error converting: %v", err)
}
want := []Value{[]Value{[]Value{int64(1), int64(2), int64(3)}}}
if !testutil.Equal(got, want) {
t.Errorf("converting basic repeated values: got:\n%v\nwant:\n%v", got, want)
}
}
func TestRepeatedRecordContainingRepetition(t *testing.T) {
schema := Schema{
{
Type: RecordFieldType,
Repeated: true,
Schema: Schema{
{Type: IntegerFieldType, Repeated: true},
},
},
}
row := &bq.TableRow{F: []*bq.TableCell{
{
V: []interface{}{ // repeated records.
map[string]interface{}{ // first record.
"v": map[string]interface{}{ // pointless single-key-map wrapper.
"f": []interface{}{ // list of record fields.
map[string]interface{}{ // only record (repeated ints)
"v": []interface{}{ // pointless wrapper.
map[string]interface{}{
"v": "1",
},
map[string]interface{}{
"v": "2",
},
map[string]interface{}{
"v": "3",
},
},
},
},
},
},
map[string]interface{}{ // second record.
"v": map[string]interface{}{
"f": []interface{}{
map[string]interface{}{
"v": []interface{}{
map[string]interface{}{
"v": "4",
},
map[string]interface{}{
"v": "5",
},
map[string]interface{}{
"v": "6",
},
},
},
},
},
},
},
},
}}
got, err := convertRow(row, schema)
if err != nil {
t.Fatalf("error converting: %v", err)
}
want := []Value{ // the row is a list of length 1, containing an entry for the repeated record.
[]Value{ // the repeated record is a list of length 2, containing an entry for each repetition.
[]Value{ // the record is a list of length 1, containing an entry for the repeated integer field.
[]Value{int64(1), int64(2), int64(3)}, // the repeated integer field is a list of length 3.
},
[]Value{ // second record
[]Value{int64(4), int64(5), int64(6)},
},
},
}
if !testutil.Equal(got, want) {
t.Errorf("converting repeated records with repeated values: got:\n%v\nwant:\n%v", got, want)
}
}
func TestRepeatedRecordContainingRecord(t *testing.T) {
schema := Schema{
{
Type: RecordFieldType,
Repeated: true,
Schema: Schema{
{
Type: StringFieldType,
},
{
Type: RecordFieldType,
Schema: Schema{
{Type: IntegerFieldType},
{Type: StringFieldType},
},
},
},
},
}
row := &bq.TableRow{F: []*bq.TableCell{
{
V: []interface{}{ // repeated records.
map[string]interface{}{ // first record.
"v": map[string]interface{}{ // pointless single-key-map wrapper.
"f": []interface{}{ // list of record fields.
map[string]interface{}{ // first record field (name)
"v": "first repeated record",
},
map[string]interface{}{ // second record field (nested record).
"v": map[string]interface{}{ // pointless single-key-map wrapper.
"f": []interface{}{ // nested record fields
map[string]interface{}{
"v": "1",
},
map[string]interface{}{
"v": "two",
},
},
},
},
},
},
},
map[string]interface{}{ // second record.
"v": map[string]interface{}{
"f": []interface{}{
map[string]interface{}{
"v": "second repeated record",
},
map[string]interface{}{
"v": map[string]interface{}{
"f": []interface{}{
map[string]interface{}{
"v": "3",
},
map[string]interface{}{
"v": "four",
},
},
},
},
},
},
},
},
},
}}
got, err := convertRow(row, schema)
if err != nil {
t.Fatalf("error converting: %v", err)
}
// TODO: test with flattenresults.
want := []Value{ // the row is a list of length 1, containing an entry for the repeated record.
[]Value{ // the repeated record is a list of length 2, containing an entry for each repetition.
[]Value{ // record contains a string followed by a nested record.
"first repeated record",
[]Value{
int64(1),
"two",
},
},
[]Value{ // second record.
"second repeated record",
[]Value{
int64(3),
"four",
},
},
},
}
if !testutil.Equal(got, want) {
t.Errorf("converting repeated records containing record : got:\n%v\nwant:\n%v", got, want)
}
}
func TestConvertRowErrors(t *testing.T) {
// mismatched lengths
if _, err := convertRow(&bq.TableRow{F: []*bq.TableCell{{V: ""}}}, Schema{}); err == nil {
t.Error("got nil, want error")
}
v3 := map[string]interface{}{"v": 3}
for _, test := range []struct {
value interface{}
fs FieldSchema
}{
{3, FieldSchema{Type: IntegerFieldType}}, // not a string
{[]interface{}{v3}, // not a string, repeated
FieldSchema{Type: IntegerFieldType, Repeated: true}},
{map[string]interface{}{"f": []interface{}{v3}}, // not a string, nested
FieldSchema{Type: RecordFieldType, Schema: Schema{{Type: IntegerFieldType}}}},
{map[string]interface{}{"f": []interface{}{v3}}, // wrong length, nested
FieldSchema{Type: RecordFieldType, Schema: Schema{}}},
} {
_, err := convertRow(
&bq.TableRow{F: []*bq.TableCell{{V: test.value}}},
Schema{&test.fs})
if err == nil {
t.Errorf("value %v, fs %v: got nil, want error", test.value, test.fs)
}
}
// bad field type
if _, err := convertBasicType("", FieldType("BAD")); err == nil {
t.Error("got nil, want error")
}
}
func TestValuesSaverConvertsToMap(t *testing.T) {
testCases := []struct {
vs ValuesSaver
wantInsertID string
wantRow map[string]Value
}{
{
vs: ValuesSaver{
Schema: Schema{
{Name: "intField", Type: IntegerFieldType},
{Name: "strField", Type: StringFieldType},
{Name: "dtField", Type: DateTimeFieldType},
{Name: "nField", Type: NumericFieldType},
},
InsertID: "iid",
Row: []Value{1, "a",
civil.DateTime{
Date: civil.Date{Year: 1, Month: 2, Day: 3},
Time: civil.Time{Hour: 4, Minute: 5, Second: 6, Nanosecond: 7000}},
big.NewRat(123456789000, 1e9),
},
},
wantInsertID: "iid",
wantRow: map[string]Value{
"intField": 1,
"strField": "a",
"dtField": "0001-02-03 04:05:06.000007",
"nField": "123.456789000",
},
},
{
vs: ValuesSaver{
Schema: Schema{
{Name: "intField", Type: IntegerFieldType},
{
Name: "recordField",
Type: RecordFieldType,
Schema: Schema{
{Name: "nestedInt", Type: IntegerFieldType, Repeated: true},
},
},
},
InsertID: "iid",
Row: []Value{1, []Value{[]Value{2, 3}}},
},
wantInsertID: "iid",
wantRow: map[string]Value{
"intField": 1,
"recordField": map[string]Value{
"nestedInt": []Value{2, 3},
},
},
},
{ // repeated nested field
vs: ValuesSaver{
Schema: Schema{
{
Name: "records",
Type: RecordFieldType,
Schema: Schema{
{Name: "x", Type: IntegerFieldType},
{Name: "y", Type: IntegerFieldType},
},
Repeated: true,
},
},
InsertID: "iid",
Row: []Value{ // a row is a []Value
[]Value{ // repeated field's value is a []Value
[]Value{1, 2}, // first record of the repeated field
[]Value{3, 4}, // second record
},
},
},
wantInsertID: "iid",
wantRow: map[string]Value{
"records": []Value{
map[string]Value{"x": 1, "y": 2},
map[string]Value{"x": 3, "y": 4},
},
},
},
}
for _, tc := range testCases {
gotRow, gotInsertID, err := tc.vs.Save()
if err != nil {
t.Errorf("Expected successful save; got: %v", err)
continue
}
if !testutil.Equal(gotRow, tc.wantRow) {
t.Errorf("%v row:\ngot:\n%+v\nwant:\n%+v", tc.vs, gotRow, tc.wantRow)
}
if !testutil.Equal(gotInsertID, tc.wantInsertID) {
t.Errorf("%v ID:\ngot:\n%+v\nwant:\n%+v", tc.vs, gotInsertID, tc.wantInsertID)
}
}
}
func TestValuesToMapErrors(t *testing.T) {
for _, test := range []struct {
values []Value
schema Schema
}{
{ // mismatched length
[]Value{1},
Schema{},
},
{ // nested record not a slice
[]Value{1},
Schema{{Type: RecordFieldType}},
},
{ // nested record mismatched length
[]Value{[]Value{1}},
Schema{{Type: RecordFieldType}},
},
{ // nested repeated record not a slice
[]Value{[]Value{1}},
Schema{{Type: RecordFieldType, Repeated: true}},
},
{ // nested repeated record mismatched length
[]Value{[]Value{[]Value{1}}},
Schema{{Type: RecordFieldType, Repeated: true}},
},
} {
_, err := valuesToMap(test.values, test.schema)
if err == nil {
t.Errorf("%v, %v: got nil, want error", test.values, test.schema)
}
}
}
func TestStructSaver(t *testing.T) {
schema := Schema{
{Name: "s", Type: StringFieldType},
{Name: "r", Type: IntegerFieldType, Repeated: true},
{Name: "t", Type: TimeFieldType},
{Name: "tr", Type: TimeFieldType, Repeated: true},
{Name: "nested", Type: RecordFieldType, Schema: Schema{
{Name: "b", Type: BooleanFieldType},
}},
{Name: "rnested", Type: RecordFieldType, Repeated: true, Schema: Schema{
{Name: "b", Type: BooleanFieldType},
}},
{Name: "p", Type: IntegerFieldType, Required: false},
{Name: "n", Type: NumericFieldType, Required: false},
{Name: "nr", Type: NumericFieldType, Repeated: true},
}
type (
N struct{ B bool }
T struct {
S string
R []int
T civil.Time
TR []civil.Time
Nested *N
Rnested []*N
P NullInt64
N *big.Rat
NR []*big.Rat
}
)
check := func(msg string, in interface{}, want map[string]Value) {
ss := StructSaver{
Schema: schema,
InsertID: "iid",
Struct: in,
}
got, gotIID, err := ss.Save()
if err != nil {
t.Fatalf("%s: %v", msg, err)
}
if wantIID := "iid"; gotIID != wantIID {
t.Errorf("%s: InsertID: got %q, want %q", msg, gotIID, wantIID)
}
if diff := testutil.Diff(got, want); diff != "" {
t.Errorf("%s: %s", msg, diff)
}
}
ct1 := civil.Time{Hour: 1, Minute: 2, Second: 3, Nanosecond: 4000}
ct2 := civil.Time{Hour: 5, Minute: 6, Second: 7, Nanosecond: 8000}
in := T{
S: "x",
R: []int{1, 2},
T: ct1,
TR: []civil.Time{ct1, ct2},
Nested: &N{B: true},
Rnested: []*N{{true}, {false}},
P: NullInt64{Valid: true, Int64: 17},
N: big.NewRat(123456, 1000),
NR: []*big.Rat{big.NewRat(3, 1), big.NewRat(56789, 1e5)},
}
want := map[string]Value{
"s": "x",
"r": []int{1, 2},
"t": "01:02:03.000004",
"tr": []string{"01:02:03.000004", "05:06:07.000008"},
"nested": map[string]Value{"b": true},
"rnested": []Value{map[string]Value{"b": true}, map[string]Value{"b": false}},
"p": NullInt64{Valid: true, Int64: 17},
"n": "123.456000000",
"nr": []string{"3.000000000", "0.567890000"},
}
check("all values", in, want)
check("all values, ptr", &in, want)
check("empty struct", T{}, map[string]Value{"s": "", "t": "00:00:00", "p": NullInt64{}})
// Missing and extra fields ignored.
type T2 struct {
S string
// missing R, Nested, RNested
Extra int
}
check("missing and extra", T2{S: "x"}, map[string]Value{"s": "x"})
check("nils in slice", T{Rnested: []*N{{true}, nil, {false}}},
map[string]Value{
"s": "",
"t": "00:00:00",
"p": NullInt64{},
"rnested": []Value{map[string]Value{"b": true}, map[string]Value(nil), map[string]Value{"b": false}},
})
check("zero-length repeated", T{Rnested: []*N{}},
map[string]Value{
"rnested": []Value{},
"s": "",
"t": "00:00:00",
"p": NullInt64{},
})
}
func TestStructSaverErrors(t *testing.T) {
type (
badField struct {
I int `bigquery:"@"`
}
badR struct{ R int }
badRN struct{ R []int }
)
for i, test := range []struct {
inputStruct interface{}
schema Schema
}{
{0, nil}, // not a struct
{&badField{}, nil}, // bad field name
{&badR{}, Schema{{Name: "r", Repeated: true}}}, // repeated field has bad type
{&badR{}, Schema{{Name: "r", Type: RecordFieldType}}}, // nested field has bad type
{&badRN{[]int{0}}, // nested repeated field has bad type
Schema{{Name: "r", Type: RecordFieldType, Repeated: true}}},
} {
ss := &StructSaver{Struct: test.inputStruct, Schema: test.schema}
_, _, err := ss.Save()
if err == nil {
t.Errorf("#%d, %v, %v: got nil, want error", i, test.inputStruct, test.schema)
}
}
}
func TestNumericString(t *testing.T) {
for _, test := range []struct {
in *big.Rat
want string
}{
{big.NewRat(2, 3), "0.666666667"}, // round to 9 places
{big.NewRat(1, 2), "0.500000000"},
{big.NewRat(1, 2*1e8), "0.000000005"},
{big.NewRat(5, 1e10), "0.000000001"}, // round up the 5 in the 10th decimal place
{big.NewRat(-5, 1e10), "-0.000000001"}, // round half away from zero
} {
got := NumericString(test.in)
if got != test.want {
t.Errorf("%v: got %q, want %q", test.in, got, test.want)
}
}
}
func TestConvertRows(t *testing.T) {
schema := Schema{
{Type: StringFieldType},
{Type: IntegerFieldType},
{Type: FloatFieldType},
{Type: BooleanFieldType},
}
rows := []*bq.TableRow{
{F: []*bq.TableCell{
{V: "a"},
{V: "1"},
{V: "1.2"},
{V: "true"},
}},
{F: []*bq.TableCell{
{V: "b"},
{V: "2"},
{V: "2.2"},
{V: "false"},
}},
}
want := [][]Value{
{"a", int64(1), 1.2, true},
{"b", int64(2), 2.2, false},
}
got, err := convertRows(rows, schema)
if err != nil {
t.Fatalf("got %v, want nil", err)
}
if !testutil.Equal(got, want) {
t.Errorf("\ngot %v\nwant %v", got, want)
}
rows[0].F[0].V = 1
_, err = convertRows(rows, schema)
if err == nil {
t.Error("got nil, want error")
}
}
func TestValueList(t *testing.T) {
schema := Schema{
{Name: "s", Type: StringFieldType},
{Name: "i", Type: IntegerFieldType},
{Name: "f", Type: FloatFieldType},
{Name: "b", Type: BooleanFieldType},
}
want := []Value{"x", 7, 3.14, true}
var got []Value
vl := (*valueList)(&got)
if err := vl.Load(want, schema); err != nil {
t.Fatal(err)
}
if !testutil.Equal(got, want) {
t.Errorf("got %+v, want %+v", got, want)
}
// Load truncates, not appends.
// https://github.com/GoogleCloudPlatform/google-cloud-go/issues/437
if err := vl.Load(want, schema); err != nil {
t.Fatal(err)
}
if !testutil.Equal(got, want) {
t.Errorf("got %+v, want %+v", got, want)
}
}
func TestValueMap(t *testing.T) {
ns := Schema{
{Name: "x", Type: IntegerFieldType},
{Name: "y", Type: IntegerFieldType},
}
schema := Schema{
{Name: "s", Type: StringFieldType},
{Name: "i", Type: IntegerFieldType},
{Name: "f", Type: FloatFieldType},
{Name: "b", Type: BooleanFieldType},
{Name: "n", Type: RecordFieldType, Schema: ns},
{Name: "rn", Type: RecordFieldType, Schema: ns, Repeated: true},
}
in := []Value{"x", 7, 3.14, true,
[]Value{1, 2},
[]Value{[]Value{3, 4}, []Value{5, 6}},
}
var vm valueMap
if err := vm.Load(in, schema); err != nil {
t.Fatal(err)
}
want := map[string]Value{
"s": "x",
"i": 7,
"f": 3.14,
"b": true,
"n": map[string]Value{"x": 1, "y": 2},
"rn": []Value{
map[string]Value{"x": 3, "y": 4},
map[string]Value{"x": 5, "y": 6},
},
}
if !testutil.Equal(vm, valueMap(want)) {
t.Errorf("got\n%+v\nwant\n%+v", vm, want)
}
in = make([]Value, len(schema))
want = map[string]Value{
"s": nil,
"i": nil,
"f": nil,
"b": nil,
"n": nil,
"rn": nil,
}
var vm2 valueMap
if err := vm2.Load(in, schema); err != nil {
t.Fatal(err)
}
if !testutil.Equal(vm2, valueMap(want)) {
t.Errorf("got\n%+v\nwant\n%+v", vm2, want)
}
}
var (
// For testing StructLoader
schema2 = Schema{
{Name: "s", Type: StringFieldType},
{Name: "s2", Type: StringFieldType},
{Name: "by", Type: BytesFieldType},
{Name: "I", Type: IntegerFieldType},
{Name: "U", Type: IntegerFieldType},
{Name: "F", Type: FloatFieldType},
{Name: "B", Type: BooleanFieldType},
{Name: "TS", Type: TimestampFieldType},
{Name: "D", Type: DateFieldType},
{Name: "T", Type: TimeFieldType},
{Name: "DT", Type: DateTimeFieldType},
{Name: "N", Type: NumericFieldType},
{Name: "nested", Type: RecordFieldType, Schema: Schema{
{Name: "nestS", Type: StringFieldType},
{Name: "nestI", Type: IntegerFieldType},
}},
{Name: "t", Type: StringFieldType},
}
testTimestamp = time.Date(2016, 11, 5, 7, 50, 22, 8, time.UTC)
testDate = civil.Date{Year: 2016, Month: 11, Day: 5}
testTime = civil.Time{Hour: 7, Minute: 50, Second: 22, Nanosecond: 8}
testDateTime = civil.DateTime{Date: testDate, Time: testTime}
testNumeric = big.NewRat(123, 456)
testValues = []Value{"x", "y", []byte{1, 2, 3}, int64(7), int64(8), 3.14, true,
testTimestamp, testDate, testTime, testDateTime, testNumeric,
[]Value{"nested", int64(17)}, "z"}
)
type testStruct1 struct {
B bool
I int
U uint16
times
S string
S2 String
By []byte
s string
F float64
N *big.Rat
Nested nested
Tagged string `bigquery:"t"`
}
type String string
type nested struct {
NestS string
NestI int
}
type times struct {
TS time.Time
T civil.Time
D civil.Date
DT civil.DateTime
}
func TestStructLoader(t *testing.T) {
var ts1 testStruct1
mustLoad(t, &ts1, schema2, testValues)
// Note: the schema field named "s" gets matched to the exported struct
// field "S", not the unexported "s".
want := &testStruct1{
B: true,
I: 7,
U: 8,
F: 3.14,
times: times{TS: testTimestamp, T: testTime, D: testDate, DT: testDateTime},
S: "x",
S2: "y",
By: []byte{1, 2, 3},
N: big.NewRat(123, 456),
Nested: nested{NestS: "nested", NestI: 17},
Tagged: "z",
}
if diff := testutil.Diff(&ts1, want, cmp.AllowUnexported(testStruct1{})); diff != "" {
t.Error(diff)
}
// Test pointers to nested structs.
type nestedPtr struct{ Nested *nested }
var np nestedPtr
mustLoad(t, &np, schema2, testValues)
want2 := &nestedPtr{Nested: &nested{NestS: "nested", NestI: 17}}
if diff := testutil.Diff(&np, want2); diff != "" {
t.Error(diff)
}
// Existing values should be reused.
nst := &nested{NestS: "x", NestI: -10}
np = nestedPtr{Nested: nst}
mustLoad(t, &np, schema2, testValues)
if diff := testutil.Diff(&np, want2); diff != "" {
t.Error(diff)
}
if np.Nested != nst {
t.Error("nested struct pointers not equal")
}
}
type repStruct struct {
Nums []int
ShortNums [2]int // to test truncation
LongNums [5]int // to test padding with zeroes
Nested []*nested
}
var (
repSchema = Schema{
{Name: "nums", Type: IntegerFieldType, Repeated: true},
{Name: "shortNums", Type: IntegerFieldType, Repeated: true},
{Name: "longNums", Type: IntegerFieldType, Repeated: true},
{Name: "nested", Type: RecordFieldType, Repeated: true, Schema: Schema{
{Name: "nestS", Type: StringFieldType},
{Name: "nestI", Type: IntegerFieldType},
}},
}
v123 = []Value{int64(1), int64(2), int64(3)}
repValues = []Value{v123, v123, v123,
[]Value{
[]Value{"x", int64(1)},
[]Value{"y", int64(2)},
},
}
)
func TestStructLoaderRepeated(t *testing.T) {
var r1 repStruct
mustLoad(t, &r1, repSchema, repValues)
want := repStruct{
Nums: []int{1, 2, 3},
ShortNums: [...]int{1, 2}, // extra values discarded
LongNums: [...]int{1, 2, 3, 0, 0},
Nested: []*nested{{"x", 1}, {"y", 2}},
}
if diff := testutil.Diff(r1, want); diff != "" {
t.Error(diff)
}
r2 := repStruct{
Nums: []int{-1, -2, -3, -4, -5}, // truncated to zero and appended to
LongNums: [...]int{-1, -2, -3, -4, -5}, // unset elements are zeroed
}
mustLoad(t, &r2, repSchema, repValues)
if diff := testutil.Diff(r2, want); diff != "" {
t.Error(diff)
}
if got, want := cap(r2.Nums), 5; got != want {
t.Errorf("cap(r2.Nums) = %d, want %d", got, want)
}
// Short slice case.
r3 := repStruct{Nums: []int{-1}}
mustLoad(t, &r3, repSchema, repValues)
if diff := testutil.Diff(r3, want); diff != "" {
t.Error(diff)
}
if got, want := cap(r3.Nums), 3; got != want {
t.Errorf("cap(r3.Nums) = %d, want %d", got, want)
}
}
type testStructNullable struct {
String NullString
Bytes []byte
Integer NullInt64
Float NullFloat64
Boolean NullBool
Timestamp NullTimestamp
Date NullDate
Time NullTime
DateTime NullDateTime
Numeric *big.Rat
Record *subNullable
}
type subNullable struct {
X NullInt64
}
var testStructNullableSchema = Schema{
{Name: "String", Type: StringFieldType, Required: false},
{Name: "Bytes", Type: BytesFieldType, Required: false},
{Name: "Integer", Type: IntegerFieldType, Required: false},
{Name: "Float", Type: FloatFieldType, Required: false},
{Name: "Boolean", Type: BooleanFieldType, Required: false},
{Name: "Timestamp", Type: TimestampFieldType, Required: false},
{Name: "Date", Type: DateFieldType, Required: false},
{Name: "Time", Type: TimeFieldType, Required: false},
{Name: "DateTime", Type: DateTimeFieldType, Required: false},
{Name: "Numeric", Type: NumericFieldType, Required: false},
{Name: "Record", Type: RecordFieldType, Required: false, Schema: Schema{
{Name: "X", Type: IntegerFieldType, Required: false},
}},
}
func TestStructLoaderNullable(t *testing.T) {
var ts testStructNullable
nilVals := make([]Value, len(testStructNullableSchema))
mustLoad(t, &ts, testStructNullableSchema, nilVals)
want := testStructNullable{}
if diff := testutil.Diff(ts, want); diff != "" {
t.Error(diff)
}
nonnilVals := []Value{"x", []byte{1, 2, 3}, int64(1), 2.3, true, testTimestamp, testDate, testTime,
testDateTime, big.NewRat(1, 2), []Value{int64(4)}}
// All ts fields are nil. Loading non-nil values will cause them all to
// be allocated.
mustLoad(t, &ts, testStructNullableSchema, nonnilVals)
want = testStructNullable{
String: NullString{StringVal: "x", Valid: true},
Bytes: []byte{1, 2, 3},
Integer: NullInt64{Int64: 1, Valid: true},
Float: NullFloat64{Float64: 2.3, Valid: true},
Boolean: NullBool{Bool: true, Valid: true},
Timestamp: NullTimestamp{Timestamp: testTimestamp, Valid: true},
Date: NullDate{Date: testDate, Valid: true},
Time: NullTime{Time: testTime, Valid: true},
DateTime: NullDateTime{DateTime: testDateTime, Valid: true},
Numeric: big.NewRat(1, 2),
Record: &subNullable{X: NullInt64{Int64: 4, Valid: true}},
}
if diff := testutil.Diff(ts, want); diff != "" {
t.Error(diff)
}
// Struct pointers are reused, byte slices are not.
want = ts
want.Bytes = []byte{17}
vals2 := []Value{nil, []byte{17}, nil, nil, nil, nil, nil, nil, nil, nil, []Value{int64(7)}}
mustLoad(t, &ts, testStructNullableSchema, vals2)
if ts.Record != want.Record {
t.Error("record pointers not identical")
}
}
func TestStructLoaderOverflow(t *testing.T) {
type S struct {
I int16
U uint16
F float32
}
schema := Schema{
{Name: "I", Type: IntegerFieldType},
{Name: "U", Type: IntegerFieldType},
{Name: "F", Type: FloatFieldType},
}
var s S
z64 := int64(0)
for _, vals := range [][]Value{
{int64(math.MaxInt16 + 1), z64, 0},
{z64, int64(math.MaxInt32), 0},
{z64, int64(-1), 0},
{z64, z64, math.MaxFloat32 * 2},
} {
if err := load(&s, schema, vals); err == nil {
t.Errorf("%+v: got nil, want error", vals)
}
}
}
func TestStructLoaderFieldOverlap(t *testing.T) {
// It's OK if the struct has fields that the schema does not, and vice versa.
type S1 struct {
I int
X [][]int // not in the schema; does not even correspond to a valid BigQuery type
// many schema fields missing
}
var s1 S1
if err := load(&s1, schema2, testValues); err != nil {
t.Fatal(err)
}
want1 := S1{I: 7}
if diff := testutil.Diff(s1, want1); diff != "" {
t.Error(diff)
}
// It's even valid to have no overlapping fields at all.
type S2 struct{ Z int }
var s2 S2
mustLoad(t, &s2, schema2, testValues)
want2 := S2{}
if diff := testutil.Diff(s2, want2); diff != "" {
t.Error(diff)
}
}
func TestStructLoaderErrors(t *testing.T) {
check := func(sp interface{}) {
var sl structLoader
err := sl.set(sp, schema2)
if err == nil {
t.Errorf("%T: got nil, want error", sp)
}
}
type bad1 struct{ F int32 } // wrong type for FLOAT column
check(&bad1{})
type bad2 struct{ I uint } // unsupported integer type
check(&bad2{})
type bad3 struct {
I int `bigquery:"@"`
} // bad field name
check(&bad3{})
type bad4 struct{ Nested int } // non-struct for nested field
check(&bad4{})
type bad5 struct{ Nested struct{ NestS int } } // bad nested struct
check(&bad5{})
bad6 := &struct{ Nums int }{} // non-slice for repeated field
sl := structLoader{}
err := sl.set(bad6, repSchema)
if err == nil {
t.Errorf("%T: got nil, want error", bad6)
}
// sl.set's error is sticky, even with good input.
err2 := sl.set(&repStruct{}, repSchema)
if err2 != err {
t.Errorf("%v != %v, expected equal", err2, err)
}
// sl.Load is similarly sticky
err2 = sl.Load(nil, nil)
if err2 != err {
t.Errorf("%v != %v, expected equal", err2, err)
}
// Null values.
schema := Schema{
{Name: "i", Type: IntegerFieldType},
{Name: "f", Type: FloatFieldType},
{Name: "b", Type: BooleanFieldType},
{Name: "s", Type: StringFieldType},
{Name: "d", Type: DateFieldType},
{Name: "r", Type: RecordFieldType, Schema: Schema{{Name: "X", Type: IntegerFieldType}}},
}
type s struct {
I int
F float64
B bool
S string
D civil.Date
}
vals := []Value{int64(0), 0.0, false, "", testDate}
mustLoad(t, &s{}, schema, vals)
for i, e := range vals {
vals[i] = nil
got := load(&s{}, schema, vals)
if got != errNoNulls {
t.Errorf("#%d: got %v, want %v", i, got, errNoNulls)
}
vals[i] = e
}
// Using more than one struct type with the same structLoader.
type different struct {
B bool
I int
times
S string
s string
Nums []int
}
sl = structLoader{}
if err := sl.set(&testStruct1{}, schema2); err != nil {
t.Fatal(err)
}
err = sl.set(&different{}, schema2)
if err == nil {
t.Error("different struct types: got nil, want error")
}
}
func mustLoad(t *testing.T, pval interface{}, schema Schema, vals []Value) {
if err := load(pval, schema, vals); err != nil {
t.Fatalf("loading: %v", err)
}
}
func load(pval interface{}, schema Schema, vals []Value) error {
var sl structLoader
if err := sl.set(pval, schema); err != nil {
return err
}
return sl.Load(vals, nil)
}
func BenchmarkStructLoader_NoCompile(b *testing.B) {
benchmarkStructLoader(b, false)
}
func BenchmarkStructLoader_Compile(b *testing.B) {
benchmarkStructLoader(b, true)
}
func benchmarkStructLoader(b *testing.B, compile bool) {
var ts1 testStruct1
for i := 0; i < b.N; i++ {
var sl structLoader
for j := 0; j < 10; j++ {
if err := load(&ts1, schema2, testValues); err != nil {
b.Fatal(err)
}
if !compile {
sl.typ = nil
}
}
}
}
/*
Copyright 2015 Google LLC
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
package bigtable
import (
"context"
"errors"
"fmt"
"math"
"regexp"
"strings"
"time"
"cloud.google.com/go/bigtable/internal/gax"
btopt "cloud.google.com/go/bigtable/internal/option"
"cloud.google.com/go/iam"
"cloud.google.com/go/internal/optional"
"cloud.google.com/go/longrunning"
lroauto "cloud.google.com/go/longrunning/autogen"
"github.com/golang/protobuf/ptypes"
durpb "github.com/golang/protobuf/ptypes/duration"
"google.golang.org/api/cloudresourcemanager/v1"
"google.golang.org/api/iterator"
"google.golang.org/api/option"
gtransport "google.golang.org/api/transport/grpc"
btapb "google.golang.org/genproto/googleapis/bigtable/admin/v2"
"google.golang.org/genproto/protobuf/field_mask"
"google.golang.org/grpc"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/metadata"
"google.golang.org/grpc/status"
)
const adminAddr = "bigtableadmin.googleapis.com:443"
// AdminClient is a client type for performing admin operations within a specific instance.
type AdminClient struct {
conn *grpc.ClientConn
tClient btapb.BigtableTableAdminClient
lroClient *lroauto.OperationsClient
project, instance string
// Metadata to be sent with each request.
md metadata.MD
}
// NewAdminClient creates a new AdminClient for a given project and instance.
func NewAdminClient(ctx context.Context, project, instance string, opts ...option.ClientOption) (*AdminClient, error) {
o, err := btopt.DefaultClientOptions(adminAddr, AdminScope, clientUserAgent)
if err != nil {
return nil, err
}
// Need to add scopes for long running operations (for create table & snapshots)
o = append(o, option.WithScopes(cloudresourcemanager.CloudPlatformScope))
o = append(o, opts...)
conn, err := gtransport.Dial(ctx, o...)
if err != nil {
return nil, fmt.Errorf("dialing: %v", err)
}
lroClient, err := lroauto.NewOperationsClient(ctx, option.WithGRPCConn(conn))
if err != nil {
// This error "should not happen", since we are just reusing old connection
// and never actually need to dial.
// If this does happen, we could leak conn. However, we cannot close conn:
// If the user invoked the function with option.WithGRPCConn,
// we would close a connection that's still in use.
// TODO(pongad): investigate error conditions.
return nil, err
}
return &AdminClient{
conn: conn,
tClient: btapb.NewBigtableTableAdminClient(conn),
lroClient: lroClient,
project: project,
instance: instance,
md: metadata.Pairs(resourcePrefixHeader, fmt.Sprintf("projects/%s/instances/%s", project, instance)),
}, nil
}
// Close closes the AdminClient.
func (ac *AdminClient) Close() error {
return ac.conn.Close()
}
func (ac *AdminClient) instancePrefix() string {
return fmt.Sprintf("projects/%s/instances/%s", ac.project, ac.instance)
}
// Tables returns a list of the tables in the instance.
func (ac *AdminClient) Tables(ctx context.Context) ([]string, error) {
ctx = mergeOutgoingMetadata(ctx, ac.md)
prefix := ac.instancePrefix()
req := &btapb.ListTablesRequest{
Parent: prefix,
}
var res *btapb.ListTablesResponse
err := gax.Invoke(ctx, func(ctx context.Context) error {
var err error
res, err = ac.tClient.ListTables(ctx, req)
return err
}, retryOptions...)
if err != nil {
return nil, err
}
names := make([]string, 0, len(res.Tables))
for _, tbl := range res.Tables {
names = append(names, strings.TrimPrefix(tbl.Name, prefix+"/tables/"))
}
return names, nil
}
// TableConf contains all of the information necessary to create a table with column families.
type TableConf struct {
TableID string
SplitKeys []string
// Families is a map from family name to GCPolicy
Families map[string]GCPolicy
}
// CreateTable creates a new table in the instance.
// This method may return before the table's creation is complete.
func (ac *AdminClient) CreateTable(ctx context.Context, table string) error {
return ac.CreateTableFromConf(ctx, &TableConf{TableID: table})
}
// CreatePresplitTable creates a new table in the instance.
// The list of row keys will be used to initially split the table into multiple tablets.
// Given two split keys, "s1" and "s2", three tablets will be created,
// spanning the key ranges: [, s1), [s1, s2), [s2, ).
// This method may return before the table's creation is complete.
func (ac *AdminClient) CreatePresplitTable(ctx context.Context, table string, splitKeys []string) error {
return ac.CreateTableFromConf(ctx, &TableConf{TableID: table, SplitKeys: splitKeys})
}
// CreateTableFromConf creates a new table in the instance from the given configuration.
func (ac *AdminClient) CreateTableFromConf(ctx context.Context, conf *TableConf) error {
ctx = mergeOutgoingMetadata(ctx, ac.md)
var reqSplits []*btapb.CreateTableRequest_Split
for _, split := range conf.SplitKeys {
reqSplits = append(reqSplits, &btapb.CreateTableRequest_Split{Key: []byte(split)})
}
var tbl btapb.Table
if conf.Families != nil {
tbl.ColumnFamilies = make(map[string]*btapb.ColumnFamily)
for fam, policy := range conf.Families {
tbl.ColumnFamilies[fam] = &btapb.ColumnFamily{GcRule: policy.proto()}
}
}
prefix := ac.instancePrefix()
req := &btapb.CreateTableRequest{
Parent: prefix,
TableId: conf.TableID,
Table: &tbl,
InitialSplits: reqSplits,
}
_, err := ac.tClient.CreateTable(ctx, req)
return err
}
// CreateColumnFamily creates a new column family in a table.
func (ac *AdminClient) CreateColumnFamily(ctx context.Context, table, family string) error {
// TODO(dsymonds): Permit specifying gcexpr and any other family settings.
ctx = mergeOutgoingMetadata(ctx, ac.md)
prefix := ac.instancePrefix()
req := &btapb.ModifyColumnFamiliesRequest{
Name: prefix + "/tables/" + table,
Modifications: []*btapb.ModifyColumnFamiliesRequest_Modification{{
Id: family,
Mod: &btapb.ModifyColumnFamiliesRequest_Modification_Create{Create: &btapb.ColumnFamily{}},
}},
}
_, err := ac.tClient.ModifyColumnFamilies(ctx, req)
return err
}
// DeleteTable deletes a table and all of its data.
func (ac *AdminClient) DeleteTable(ctx context.Context, table string) error {
ctx = mergeOutgoingMetadata(ctx, ac.md)
prefix := ac.instancePrefix()
req := &btapb.DeleteTableRequest{
Name: prefix + "/tables/" + table,
}
_, err := ac.tClient.DeleteTable(ctx, req)
return err
}
// DeleteColumnFamily deletes a column family in a table and all of its data.
func (ac *AdminClient) DeleteColumnFamily(ctx context.Context, table, family string) error {
ctx = mergeOutgoingMetadata(ctx, ac.md)
prefix := ac.instancePrefix()
req := &btapb.ModifyColumnFamiliesRequest{
Name: prefix + "/tables/" + table,
Modifications: []*btapb.ModifyColumnFamiliesRequest_Modification{{
Id: family,
Mod: &btapb.ModifyColumnFamiliesRequest_Modification_Drop{Drop: true},
}},
}
_, err := ac.tClient.ModifyColumnFamilies(ctx, req)
return err
}
// TableInfo represents information about a table.
type TableInfo struct {
// DEPRECATED - This field is deprecated. Please use FamilyInfos instead.
Families []string
FamilyInfos []FamilyInfo
}
// FamilyInfo represents information about a column family.
type FamilyInfo struct {
Name string
GCPolicy string
}
// TableInfo retrieves information about a table.
func (ac *AdminClient) TableInfo(ctx context.Context, table string) (*TableInfo, error) {
ctx = mergeOutgoingMetadata(ctx, ac.md)
prefix := ac.instancePrefix()
req := &btapb.GetTableRequest{
Name: prefix + "/tables/" + table,
}
var res *btapb.Table
err := gax.Invoke(ctx, func(ctx context.Context) error {
var err error
res, err = ac.tClient.GetTable(ctx, req)
return err
}, retryOptions...)
if err != nil {
return nil, err
}
ti := &TableInfo{}
for name, fam := range res.ColumnFamilies {
ti.Families = append(ti.Families, name)
ti.FamilyInfos = append(ti.FamilyInfos, FamilyInfo{Name: name, GCPolicy: GCRuleToString(fam.GcRule)})
}
return ti, nil
}
// SetGCPolicy specifies which cells in a column family should be garbage collected.
// GC executes opportunistically in the background; table reads may return data
// matching the GC policy.
func (ac *AdminClient) SetGCPolicy(ctx context.Context, table, family string, policy GCPolicy) error {
ctx = mergeOutgoingMetadata(ctx, ac.md)
prefix := ac.instancePrefix()
req := &btapb.ModifyColumnFamiliesRequest{
Name: prefix + "/tables/" + table,
Modifications: []*btapb.ModifyColumnFamiliesRequest_Modification{{
Id: family,
Mod: &btapb.ModifyColumnFamiliesRequest_Modification_Update{Update: &btapb.ColumnFamily{GcRule: policy.proto()}},
}},
}
_, err := ac.tClient.ModifyColumnFamilies(ctx, req)
return err
}
// DropRowRange permanently deletes a row range from the specified table.
func (ac *AdminClient) DropRowRange(ctx context.Context, table, rowKeyPrefix string) error {
ctx = mergeOutgoingMetadata(ctx, ac.md)
prefix := ac.instancePrefix()
req := &btapb.DropRowRangeRequest{
Name: prefix + "/tables/" + table,
Target: &btapb.DropRowRangeRequest_RowKeyPrefix{RowKeyPrefix: []byte(rowKeyPrefix)},
}
_, err := ac.tClient.DropRowRange(ctx, req)
return err
}
// CreateTableFromSnapshot creates a table from snapshot.
// The table will be created in the same cluster as the snapshot.
//
// This is a private alpha release of Cloud Bigtable snapshots. This feature
// is not currently available to most Cloud Bigtable customers. This feature
// might be changed in backward-incompatible ways and is not recommended for
// production use. It is not subject to any SLA or deprecation policy.
func (ac *AdminClient) CreateTableFromSnapshot(ctx context.Context, table, cluster, snapshot string) error {
ctx = mergeOutgoingMetadata(ctx, ac.md)
prefix := ac.instancePrefix()
snapshotPath := prefix + "/clusters/" + cluster + "/snapshots/" + snapshot
req := &btapb.CreateTableFromSnapshotRequest{
Parent: prefix,
TableId: table,
SourceSnapshot: snapshotPath,
}
op, err := ac.tClient.CreateTableFromSnapshot(ctx, req)
if err != nil {
return err
}
resp := btapb.Table{}
return longrunning.InternalNewOperation(ac.lroClient, op).Wait(ctx, &resp)
}
// DefaultSnapshotDuration is the default TTL for a snapshot.
const DefaultSnapshotDuration time.Duration = 0
// SnapshotTable creates a new snapshot in the specified cluster from the
// specified source table. Setting the TTL to `DefaultSnapshotDuration` will
// use the server side default for the duration.
//
// This is a private alpha release of Cloud Bigtable snapshots. This feature
// is not currently available to most Cloud Bigtable customers. This feature
// might be changed in backward-incompatible ways and is not recommended for
// production use. It is not subject to any SLA or deprecation policy.
func (ac *AdminClient) SnapshotTable(ctx context.Context, table, cluster, snapshot string, ttl time.Duration) error {
ctx = mergeOutgoingMetadata(ctx, ac.md)
prefix := ac.instancePrefix()
var ttlProto *durpb.Duration
if ttl > 0 {
ttlProto = ptypes.DurationProto(ttl)
}
req := &btapb.SnapshotTableRequest{
Name: prefix + "/tables/" + table,
Cluster: prefix + "/clusters/" + cluster,
SnapshotId: snapshot,
Ttl: ttlProto,
}
op, err := ac.tClient.SnapshotTable(ctx, req)
if err != nil {
return err
}
resp := btapb.Snapshot{}
return longrunning.InternalNewOperation(ac.lroClient, op).Wait(ctx, &resp)
}
// Snapshots returns a SnapshotIterator for iterating over the snapshots in a cluster.
// To list snapshots across all of the clusters in the instance specify "-" as the cluster.
//
// This is a private alpha release of Cloud Bigtable snapshots. This feature is not
// currently available to most Cloud Bigtable customers. This feature might be
// changed in backward-incompatible ways and is not recommended for production use.
// It is not subject to any SLA or deprecation policy.
func (ac *AdminClient) Snapshots(ctx context.Context, cluster string) *SnapshotIterator {
ctx = mergeOutgoingMetadata(ctx, ac.md)
prefix := ac.instancePrefix()
clusterPath := prefix + "/clusters/" + cluster
it := &SnapshotIterator{}
req := &btapb.ListSnapshotsRequest{
Parent: clusterPath,
}
fetch := func(pageSize int, pageToken string) (string, error) {
req.PageToken = pageToken
if pageSize > math.MaxInt32 {
req.PageSize = math.MaxInt32
} else {
req.PageSize = int32(pageSize)
}
var resp *btapb.ListSnapshotsResponse
err := gax.Invoke(ctx, func(ctx context.Context) error {
var err error
resp, err = ac.tClient.ListSnapshots(ctx, req)
return err
}, retryOptions...)
if err != nil {
return "", err
}
for _, s := range resp.Snapshots {
snapshotInfo, err := newSnapshotInfo(s)
if err != nil {
return "", fmt.Errorf("Failed to parse snapshot proto %v", err)
}
it.items = append(it.items, snapshotInfo)
}
return resp.NextPageToken, nil
}
bufLen := func() int { return len(it.items) }
takeBuf := func() interface{} { b := it.items; it.items = nil; return b }
it.pageInfo, it.nextFunc = iterator.NewPageInfo(fetch, bufLen, takeBuf)
return it
}
func newSnapshotInfo(snapshot *btapb.Snapshot) (*SnapshotInfo, error) {
nameParts := strings.Split(snapshot.Name, "/")
name := nameParts[len(nameParts)-1]
tablePathParts := strings.Split(snapshot.SourceTable.Name, "/")
tableID := tablePathParts[len(tablePathParts)-1]
createTime, err := ptypes.Timestamp(snapshot.CreateTime)
if err != nil {
return nil, fmt.Errorf("Invalid createTime: %v", err)
}
deleteTime, err := ptypes.Timestamp(snapshot.DeleteTime)
if err != nil {
return nil, fmt.Errorf("Invalid deleteTime: %v", err)
}
return &SnapshotInfo{
Name: name,
SourceTable: tableID,
DataSize: snapshot.DataSizeBytes,
CreateTime: createTime,
DeleteTime: deleteTime,
}, nil
}
// SnapshotIterator is an EntryIterator that iterates over log entries.
//
// This is a private alpha release of Cloud Bigtable snapshots. This feature
// is not currently available to most Cloud Bigtable customers. This feature
// might be changed in backward-incompatible ways and is not recommended for
// production use. It is not subject to any SLA or deprecation policy.
type SnapshotIterator struct {
items []*SnapshotInfo
pageInfo *iterator.PageInfo
nextFunc func() error
}
// PageInfo supports pagination. See https://godoc.org/google.golang.org/api/iterator package for details.
func (it *SnapshotIterator) PageInfo() *iterator.PageInfo {
return it.pageInfo
}
// Next returns the next result. Its second return value is iterator.Done
// (https://godoc.org/google.golang.org/api/iterator) if there are no more
// results. Once Next returns Done, all subsequent calls will return Done.
func (it *SnapshotIterator) Next() (*SnapshotInfo, error) {
if err := it.nextFunc(); err != nil {
return nil, err
}
item := it.items[0]
it.items = it.items[1:]
return item, nil
}
// SnapshotInfo contains snapshot metadata.
type SnapshotInfo struct {
Name string
SourceTable string
DataSize int64
CreateTime time.Time
DeleteTime time.Time
}
// SnapshotInfo gets snapshot metadata.
//
// This is a private alpha release of Cloud Bigtable snapshots. This feature
// is not currently available to most Cloud Bigtable customers. This feature
// might be changed in backward-incompatible ways and is not recommended for
// production use. It is not subject to any SLA or deprecation policy.
func (ac *AdminClient) SnapshotInfo(ctx context.Context, cluster, snapshot string) (*SnapshotInfo, error) {
ctx = mergeOutgoingMetadata(ctx, ac.md)
prefix := ac.instancePrefix()
clusterPath := prefix + "/clusters/" + cluster
snapshotPath := clusterPath + "/snapshots/" + snapshot
req := &btapb.GetSnapshotRequest{
Name: snapshotPath,
}
var resp *btapb.Snapshot
err := gax.Invoke(ctx, func(ctx context.Context) error {
var err error
resp, err = ac.tClient.GetSnapshot(ctx, req)
return err
}, retryOptions...)
if err != nil {
return nil, err
}
return newSnapshotInfo(resp)
}
// DeleteSnapshot deletes a snapshot in a cluster.
//
// This is a private alpha release of Cloud Bigtable snapshots. This feature
// is not currently available to most Cloud Bigtable customers. This feature
// might be changed in backward-incompatible ways and is not recommended for
// production use. It is not subject to any SLA or deprecation policy.
func (ac *AdminClient) DeleteSnapshot(ctx context.Context, cluster, snapshot string) error {
ctx = mergeOutgoingMetadata(ctx, ac.md)
prefix := ac.instancePrefix()
clusterPath := prefix + "/clusters/" + cluster
snapshotPath := clusterPath + "/snapshots/" + snapshot
req := &btapb.DeleteSnapshotRequest{
Name: snapshotPath,
}
_, err := ac.tClient.DeleteSnapshot(ctx, req)
return err
}
// getConsistencyToken gets the consistency token for a table.
func (ac *AdminClient) getConsistencyToken(ctx context.Context, tableName string) (string, error) {
req := &btapb.GenerateConsistencyTokenRequest{
Name: tableName,
}
resp, err := ac.tClient.GenerateConsistencyToken(ctx, req)
if err != nil {
return "", err
}
return resp.GetConsistencyToken(), nil
}
// isConsistent checks if a token is consistent for a table.
func (ac *AdminClient) isConsistent(ctx context.Context, tableName, token string) (bool, error) {
req := &btapb.CheckConsistencyRequest{
Name: tableName,
ConsistencyToken: token,
}
var resp *btapb.CheckConsistencyResponse
// Retry calls on retryable errors to avoid losing the token gathered before.
err := gax.Invoke(ctx, func(ctx context.Context) error {
var err error
resp, err = ac.tClient.CheckConsistency(ctx, req)
return err
}, retryOptions...)
if err != nil {
return false, err
}
return resp.GetConsistent(), nil
}
// WaitForReplication waits until all the writes committed before the call started have been propagated to all the clusters in the instance via replication.
func (ac *AdminClient) WaitForReplication(ctx context.Context, table string) error {
// Get the token.
prefix := ac.instancePrefix()
tableName := prefix + "/tables/" + table
token, err := ac.getConsistencyToken(ctx, tableName)
if err != nil {
return err
}
// Periodically check if the token is consistent.
timer := time.NewTicker(time.Second * 10)
defer timer.Stop()
for {
consistent, err := ac.isConsistent(ctx, tableName, token)
if err != nil {
return err
}
if consistent {
return nil
}
// Sleep for a bit or until the ctx is cancelled.
select {
case <-ctx.Done():
return ctx.Err()
case <-timer.C:
}
}
}
const instanceAdminAddr = "bigtableadmin.googleapis.com:443"
// InstanceAdminClient is a client type for performing admin operations on instances.
// These operations can be substantially more dangerous than those provided by AdminClient.
type InstanceAdminClient struct {
conn *grpc.ClientConn
iClient btapb.BigtableInstanceAdminClient
lroClient *lroauto.OperationsClient
project string
// Metadata to be sent with each request.
md metadata.MD
}
// NewInstanceAdminClient creates a new InstanceAdminClient for a given project.
func NewInstanceAdminClient(ctx context.Context, project string, opts ...option.ClientOption) (*InstanceAdminClient, error) {
o, err := btopt.DefaultClientOptions(instanceAdminAddr, InstanceAdminScope, clientUserAgent)
if err != nil {
return nil, err
}
o = append(o, opts...)
conn, err := gtransport.Dial(ctx, o...)
if err != nil {
return nil, fmt.Errorf("dialing: %v", err)
}
lroClient, err := lroauto.NewOperationsClient(ctx, option.WithGRPCConn(conn))
if err != nil {
// This error "should not happen", since we are just reusing old connection
// and never actually need to dial.
// If this does happen, we could leak conn. However, we cannot close conn:
// If the user invoked the function with option.WithGRPCConn,
// we would close a connection that's still in use.
// TODO(pongad): investigate error conditions.
return nil, err
}
return &InstanceAdminClient{
conn: conn,
iClient: btapb.NewBigtableInstanceAdminClient(conn),
lroClient: lroClient,
project: project,
md: metadata.Pairs(resourcePrefixHeader, "projects/"+project),
}, nil
}
// Close closes the InstanceAdminClient.
func (iac *InstanceAdminClient) Close() error {
return iac.conn.Close()
}
// StorageType is the type of storage used for all tables in an instance
type StorageType int
const (
SSD StorageType = iota
HDD
)
func (st StorageType) proto() btapb.StorageType {
if st == HDD {
return btapb.StorageType_HDD
}
return btapb.StorageType_SSD
}
// InstanceType is the type of the instance
type InstanceType int32
const (
PRODUCTION InstanceType = InstanceType(btapb.Instance_PRODUCTION)
DEVELOPMENT = InstanceType(btapb.Instance_DEVELOPMENT)
)
// InstanceInfo represents information about an instance
type InstanceInfo struct {
Name string // name of the instance
DisplayName string // display name for UIs
}
// InstanceConf contains the information necessary to create an Instance
type InstanceConf struct {
InstanceId, DisplayName, ClusterId, Zone string
// NumNodes must not be specified for DEVELOPMENT instance types
NumNodes int32
StorageType StorageType
InstanceType InstanceType
}
// InstanceWithClustersConfig contains the information necessary to create an Instance
type InstanceWithClustersConfig struct {
InstanceID, DisplayName string
Clusters []ClusterConfig
InstanceType InstanceType
}
var instanceNameRegexp = regexp.MustCompile(`^projects/([^/]+)/instances/([a-z][-a-z0-9]*)$`)
// CreateInstance creates a new instance in the project.
// This method will return when the instance has been created or when an error occurs.
func (iac *InstanceAdminClient) CreateInstance(ctx context.Context, conf *InstanceConf) error {
newConfig := InstanceWithClustersConfig{
InstanceID: conf.InstanceId,
DisplayName: conf.DisplayName,
InstanceType: conf.InstanceType,
Clusters: []ClusterConfig{
{
InstanceID: conf.InstanceId,
ClusterID: conf.ClusterId,
Zone: conf.Zone,
NumNodes: conf.NumNodes,
StorageType: conf.StorageType,
},
},
}
return iac.CreateInstanceWithClusters(ctx, &newConfig)
}
// CreateInstanceWithClusters creates a new instance with configured clusters in the project.
// This method will return when the instance has been created or when an error occurs.
func (iac *InstanceAdminClient) CreateInstanceWithClusters(ctx context.Context, conf *InstanceWithClustersConfig) error {
ctx = mergeOutgoingMetadata(ctx, iac.md)
clusters := make(map[string]*btapb.Cluster)
for _, cluster := range conf.Clusters {
clusters[cluster.ClusterID] = cluster.proto(iac.project)
}
req := &btapb.CreateInstanceRequest{
Parent: "projects/" + iac.project,
InstanceId: conf.InstanceID,
Instance: &btapb.Instance{DisplayName: conf.DisplayName, Type: btapb.Instance_Type(conf.InstanceType)},
Clusters: clusters,
}
lro, err := iac.iClient.CreateInstance(ctx, req)
if err != nil {
return err
}
resp := btapb.Instance{}
return longrunning.InternalNewOperation(iac.lroClient, lro).Wait(ctx, &resp)
}
// DeleteInstance deletes an instance from the project.
func (iac *InstanceAdminClient) DeleteInstance(ctx context.Context, instanceID string) error {
ctx = mergeOutgoingMetadata(ctx, iac.md)
req := &btapb.DeleteInstanceRequest{Name: "projects/" + iac.project + "/instances/" + instanceID}
_, err := iac.iClient.DeleteInstance(ctx, req)
return err
}
// Instances returns a list of instances in the project.
func (iac *InstanceAdminClient) Instances(ctx context.Context) ([]*InstanceInfo, error) {
ctx = mergeOutgoingMetadata(ctx, iac.md)
req := &btapb.ListInstancesRequest{
Parent: "projects/" + iac.project,
}
var res *btapb.ListInstancesResponse
err := gax.Invoke(ctx, func(ctx context.Context) error {
var err error
res, err = iac.iClient.ListInstances(ctx, req)
return err
}, retryOptions...)
if err != nil {
return nil, err
}
if len(res.FailedLocations) > 0 {
// We don't have a good way to return a partial result in the face of some zones being unavailable.
// Fail the entire request.
return nil, status.Errorf(codes.Unavailable, "Failed locations: %v", res.FailedLocations)
}
var is []*InstanceInfo
for _, i := range res.Instances {
m := instanceNameRegexp.FindStringSubmatch(i.Name)
if m == nil {
return nil, fmt.Errorf("malformed instance name %q", i.Name)
}
is = append(is, &InstanceInfo{
Name: m[2],
DisplayName: i.DisplayName,
})
}
return is, nil
}
// InstanceInfo returns information about an instance.
func (iac *InstanceAdminClient) InstanceInfo(ctx context.Context, instanceID string) (*InstanceInfo, error) {
ctx = mergeOutgoingMetadata(ctx, iac.md)
req := &btapb.GetInstanceRequest{
Name: "projects/" + iac.project + "/instances/" + instanceID,
}
var res *btapb.Instance
err := gax.Invoke(ctx, func(ctx context.Context) error {
var err error
res, err = iac.iClient.GetInstance(ctx, req)
return err
}, retryOptions...)
if err != nil {
return nil, err
}
m := instanceNameRegexp.FindStringSubmatch(res.Name)
if m == nil {
return nil, fmt.Errorf("malformed instance name %q", res.Name)
}
return &InstanceInfo{
Name: m[2],
DisplayName: res.DisplayName,
}, nil
}
// ClusterConfig contains the information necessary to create a cluster
type ClusterConfig struct {
InstanceID, ClusterID, Zone string
NumNodes int32
StorageType StorageType
}
func (cc *ClusterConfig) proto(project string) *btapb.Cluster {
return &btapb.Cluster{
ServeNodes: cc.NumNodes,
DefaultStorageType: cc.StorageType.proto(),
Location: "projects/" + project + "/locations/" + cc.Zone,
}
}
// ClusterInfo represents information about a cluster.
type ClusterInfo struct {
Name string // name of the cluster
Zone string // GCP zone of the cluster (e.g. "us-central1-a")
ServeNodes int // number of allocated serve nodes
State string // state of the cluster
}
// CreateCluster creates a new cluster in an instance.
// This method will return when the cluster has been created or when an error occurs.
func (iac *InstanceAdminClient) CreateCluster(ctx context.Context, conf *ClusterConfig) error {
ctx = mergeOutgoingMetadata(ctx, iac.md)
req := &btapb.CreateClusterRequest{
Parent: "projects/" + iac.project + "/instances/" + conf.InstanceID,
ClusterId: conf.ClusterID,
Cluster: conf.proto(iac.project),
}
lro, err := iac.iClient.CreateCluster(ctx, req)
if err != nil {
return err
}
resp := btapb.Cluster{}
return longrunning.InternalNewOperation(iac.lroClient, lro).Wait(ctx, &resp)
}
// DeleteCluster deletes a cluster from an instance.
func (iac *InstanceAdminClient) DeleteCluster(ctx context.Context, instanceID, clusterID string) error {
ctx = mergeOutgoingMetadata(ctx, iac.md)
req := &btapb.DeleteClusterRequest{Name: "projects/" + iac.project + "/instances/" + instanceID + "/clusters/" + clusterID}
_, err := iac.iClient.DeleteCluster(ctx, req)
return err
}
// UpdateCluster updates attributes of a cluster
func (iac *InstanceAdminClient) UpdateCluster(ctx context.Context, instanceID, clusterID string, serveNodes int32) error {
ctx = mergeOutgoingMetadata(ctx, iac.md)
cluster := &btapb.Cluster{
Name: "projects/" + iac.project + "/instances/" + instanceID + "/clusters/" + clusterID,
ServeNodes: serveNodes}
lro, err := iac.iClient.UpdateCluster(ctx, cluster)
if err != nil {
return err
}
return longrunning.InternalNewOperation(iac.lroClient, lro).Wait(ctx, nil)
}
// Clusters lists the clusters in an instance.
func (iac *InstanceAdminClient) Clusters(ctx context.Context, instanceID string) ([]*ClusterInfo, error) {
ctx = mergeOutgoingMetadata(ctx, iac.md)
req := &btapb.ListClustersRequest{Parent: "projects/" + iac.project + "/instances/" + instanceID}
var res *btapb.ListClustersResponse
err := gax.Invoke(ctx, func(ctx context.Context) error {
var err error
res, err = iac.iClient.ListClusters(ctx, req)
return err
}, retryOptions...)
if err != nil {
return nil, err
}
// TODO(garyelliott): Deal with failed_locations.
var cis []*ClusterInfo
for _, c := range res.Clusters {
nameParts := strings.Split(c.Name, "/")
locParts := strings.Split(c.Location, "/")
cis = append(cis, &ClusterInfo{
Name: nameParts[len(nameParts)-1],
Zone: locParts[len(locParts)-1],
ServeNodes: int(c.ServeNodes),
State: c.State.String(),
})
}
return cis, nil
}
// GetCluster fetches a cluster in an instance
func (iac *InstanceAdminClient) GetCluster(ctx context.Context, instanceID, clusterID string) (*ClusterInfo, error) {
ctx = mergeOutgoingMetadata(ctx, iac.md)
req := &btapb.GetClusterRequest{Name: "projects/" + iac.project + "/instances/" + instanceID + "/clusters/" + clusterID}
var c *btapb.Cluster
err := gax.Invoke(ctx, func(ctx context.Context) error {
var err error
c, err = iac.iClient.GetCluster(ctx, req)
return err
}, retryOptions...)
if err != nil {
return nil, err
}
nameParts := strings.Split(c.Name, "/")
locParts := strings.Split(c.Location, "/")
cis := &ClusterInfo{
Name: nameParts[len(nameParts)-1],
Zone: locParts[len(locParts)-1],
ServeNodes: int(c.ServeNodes),
State: c.State.String(),
}
return cis, nil
}
// InstanceIAM returns the instance's IAM handle.
func (iac *InstanceAdminClient) InstanceIAM(instanceID string) *iam.Handle {
return iam.InternalNewHandleGRPCClient(iac.iClient, "projects/"+iac.project+"/instances/"+instanceID)
}
// Routing policies.
const (
// MultiClusterRouting is a policy that allows read/write requests to be
// routed to any cluster in the instance. Requests will will fail over to
// another cluster in the event of transient errors or delays. Choosing
// this option sacrifices read-your-writes consistency to improve
// availability.
MultiClusterRouting = "multi_cluster_routing_use_any"
// SingleClusterRouting is a policy that unconditionally routes all
// read/write requests to a specific cluster. This option preserves
// read-your-writes consistency, but does not improve availability.
SingleClusterRouting = "single_cluster_routing"
)
// ProfileConf contains the information necessary to create an profile
type ProfileConf struct {
Name string
ProfileID string
InstanceID string
Etag string
Description string
RoutingPolicy string
ClusterID string
AllowTransactionalWrites bool
// If true, warnings are ignored
IgnoreWarnings bool
}
// ProfileIterator iterates over profiles.
type ProfileIterator struct {
items []*btapb.AppProfile
pageInfo *iterator.PageInfo
nextFunc func() error
}
// ProfileAttrsToUpdate define addrs to update during an Update call. If unset, no fields will be replaced.
type ProfileAttrsToUpdate struct {
// If set, updates the description.
Description optional.String
//If set, updates the routing policy.
RoutingPolicy optional.String
//If RoutingPolicy is updated to SingleClusterRouting, set these fields as well.
ClusterID string
AllowTransactionalWrites bool
// If true, warnings are ignored
IgnoreWarnings bool
}
// GetFieldMaskPath returns the field mask path.
func (p *ProfileAttrsToUpdate) GetFieldMaskPath() []string {
path := make([]string, 0)
if p.Description != nil {
path = append(path, "description")
}
if p.RoutingPolicy != nil {
path = append(path, optional.ToString(p.RoutingPolicy))
}
return path
}
// PageInfo supports pagination. See https://godoc.org/google.golang.org/api/iterator package for details.
func (it *ProfileIterator) PageInfo() *iterator.PageInfo {
return it.pageInfo
}
// Next returns the next result. Its second return value is iterator.Done
// (https://godoc.org/google.golang.org/api/iterator) if there are no more
// results. Once Next returns Done, all subsequent calls will return Done.
func (it *ProfileIterator) Next() (*btapb.AppProfile, error) {
if err := it.nextFunc(); err != nil {
return nil, err
}
item := it.items[0]
it.items = it.items[1:]
return item, nil
}
// CreateAppProfile creates an app profile within an instance.
func (iac *InstanceAdminClient) CreateAppProfile(ctx context.Context, profile ProfileConf) (*btapb.AppProfile, error) {
ctx = mergeOutgoingMetadata(ctx, iac.md)
parent := "projects/" + iac.project + "/instances/" + profile.InstanceID
appProfile := &btapb.AppProfile{
Etag: profile.Etag,
Description: profile.Description,
}
if profile.RoutingPolicy == "" {
return nil, errors.New("invalid routing policy")
}
switch profile.RoutingPolicy {
case MultiClusterRouting:
appProfile.RoutingPolicy = &btapb.AppProfile_MultiClusterRoutingUseAny_{
MultiClusterRoutingUseAny: &btapb.AppProfile_MultiClusterRoutingUseAny{},
}
case SingleClusterRouting:
appProfile.RoutingPolicy = &btapb.AppProfile_SingleClusterRouting_{
SingleClusterRouting: &btapb.AppProfile_SingleClusterRouting{
ClusterId: profile.ClusterID,
AllowTransactionalWrites: profile.AllowTransactionalWrites,
},
}
default:
return nil, errors.New("invalid routing policy")
}
return iac.iClient.CreateAppProfile(ctx, &btapb.CreateAppProfileRequest{
Parent: parent,
AppProfile: appProfile,
AppProfileId: profile.ProfileID,
IgnoreWarnings: profile.IgnoreWarnings,
})
}
// GetAppProfile gets information about an app profile.
func (iac *InstanceAdminClient) GetAppProfile(ctx context.Context, instanceID, name string) (*btapb.AppProfile, error) {
ctx = mergeOutgoingMetadata(ctx, iac.md)
profileRequest := &btapb.GetAppProfileRequest{
Name: "projects/" + iac.project + "/instances/" + instanceID + "/appProfiles/" + name,
}
var ap *btapb.AppProfile
err := gax.Invoke(ctx, func(ctx context.Context) error {
var err error
ap, err = iac.iClient.GetAppProfile(ctx, profileRequest)
return err
}, retryOptions...)
if err != nil {
return nil, err
}
return ap, err
}
// ListAppProfiles lists information about app profiles in an instance.
func (iac *InstanceAdminClient) ListAppProfiles(ctx context.Context, instanceID string) *ProfileIterator {
ctx = mergeOutgoingMetadata(ctx, iac.md)
listRequest := &btapb.ListAppProfilesRequest{
Parent: "projects/" + iac.project + "/instances/" + instanceID,
}
pit := &ProfileIterator{}
fetch := func(pageSize int, pageToken string) (string, error) {
listRequest.PageToken = pageToken
var profileRes *btapb.ListAppProfilesResponse
err := gax.Invoke(ctx, func(ctx context.Context) error {
var err error
profileRes, err = iac.iClient.ListAppProfiles(ctx, listRequest)
return err
}, retryOptions...)
if err != nil {
return "", err
}
for _, a := range profileRes.AppProfiles {
pit.items = append(pit.items, a)
}
return profileRes.NextPageToken, nil
}
bufLen := func() int { return len(pit.items) }
takeBuf := func() interface{} { b := pit.items; pit.items = nil; return b }
pit.pageInfo, pit.nextFunc = iterator.NewPageInfo(fetch, bufLen, takeBuf)
return pit
}
// UpdateAppProfile updates an app profile within an instance.
// updateAttrs should be set. If unset, all fields will be replaced.
func (iac *InstanceAdminClient) UpdateAppProfile(ctx context.Context, instanceID, profileID string, updateAttrs ProfileAttrsToUpdate) error {
ctx = mergeOutgoingMetadata(ctx, iac.md)
profile := &btapb.AppProfile{
Name: "projects/" + iac.project + "/instances/" + instanceID + "/appProfiles/" + profileID,
}
if updateAttrs.Description != nil {
profile.Description = optional.ToString(updateAttrs.Description)
}
if updateAttrs.RoutingPolicy != nil {
switch optional.ToString(updateAttrs.RoutingPolicy) {
case MultiClusterRouting:
profile.RoutingPolicy = &btapb.AppProfile_MultiClusterRoutingUseAny_{
MultiClusterRoutingUseAny: &btapb.AppProfile_MultiClusterRoutingUseAny{},
}
case SingleClusterRouting:
profile.RoutingPolicy = &btapb.AppProfile_SingleClusterRouting_{
SingleClusterRouting: &btapb.AppProfile_SingleClusterRouting{
ClusterId: updateAttrs.ClusterID,
AllowTransactionalWrites: updateAttrs.AllowTransactionalWrites,
},
}
default:
return errors.New("invalid routing policy")
}
}
patchRequest := &btapb.UpdateAppProfileRequest{
AppProfile: profile,
UpdateMask: &field_mask.FieldMask{
Paths: updateAttrs.GetFieldMaskPath(),
},
IgnoreWarnings: updateAttrs.IgnoreWarnings,
}
updateRequest, err := iac.iClient.UpdateAppProfile(ctx, patchRequest)
if err != nil {
return err
}
return longrunning.InternalNewOperation(iac.lroClient, updateRequest).Wait(ctx, nil)
}
// DeleteAppProfile deletes an app profile from an instance.
func (iac *InstanceAdminClient) DeleteAppProfile(ctx context.Context, instanceID, name string) error {
ctx = mergeOutgoingMetadata(ctx, iac.md)
deleteProfileRequest := &btapb.DeleteAppProfileRequest{
Name: "projects/" + iac.project + "/instances/" + instanceID + "/appProfiles/" + name,
IgnoreWarnings: true,
}
_, err := iac.iClient.DeleteAppProfile(ctx, deleteProfileRequest)
return err
}
// Copyright 2015 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package bigtable
import (
"context"
"fmt"
"math"
"sort"
"strings"
"testing"
"time"
"cloud.google.com/go/internal/testutil"
"github.com/golang/protobuf/proto"
"google.golang.org/api/iterator"
btapb "google.golang.org/genproto/googleapis/bigtable/admin/v2"
)
func TestAdminIntegration(t *testing.T) {
testEnv, err := NewIntegrationEnv()
if err != nil {
t.Fatalf("IntegrationEnv: %v", err)
}
defer testEnv.Close()
timeout := 2 * time.Second
if testEnv.Config().UseProd {
timeout = 5 * time.Minute
}
ctx, _ := context.WithTimeout(context.Background(), timeout)
adminClient, err := testEnv.NewAdminClient()
if err != nil {
t.Fatalf("NewAdminClient: %v", err)
}
defer adminClient.Close()
iAdminClient, err := testEnv.NewInstanceAdminClient()
if err != nil {
t.Fatalf("NewInstanceAdminClient: %v", err)
}
if iAdminClient != nil {
defer iAdminClient.Close()
iInfo, err := iAdminClient.InstanceInfo(ctx, adminClient.instance)
if err != nil {
t.Errorf("InstanceInfo: %v", err)
}
if iInfo.Name != adminClient.instance {
t.Errorf("InstanceInfo returned name %#v, want %#v", iInfo.Name, adminClient.instance)
}
}
list := func() []string {
tbls, err := adminClient.Tables(ctx)
if err != nil {
t.Fatalf("Fetching list of tables: %v", err)
}
sort.Strings(tbls)
return tbls
}
containsAll := func(got, want []string) bool {
gotSet := make(map[string]bool)
for _, s := range got {
gotSet[s] = true
}
for _, s := range want {
if !gotSet[s] {
return false
}
}
return true
}
defer adminClient.DeleteTable(ctx, "mytable")
if err := adminClient.CreateTable(ctx, "mytable"); err != nil {
t.Fatalf("Creating table: %v", err)
}
defer adminClient.DeleteTable(ctx, "myothertable")
if err := adminClient.CreateTable(ctx, "myothertable"); err != nil {
t.Fatalf("Creating table: %v", err)
}
if got, want := list(), []string{"myothertable", "mytable"}; !containsAll(got, want) {
t.Errorf("adminClient.Tables returned %#v, want %#v", got, want)
}
must(adminClient.WaitForReplication(ctx, "mytable"))
if err := adminClient.DeleteTable(ctx, "myothertable"); err != nil {
t.Fatalf("Deleting table: %v", err)
}
tables := list()
if got, want := tables, []string{"mytable"}; !containsAll(got, want) {
t.Errorf("adminClient.Tables returned %#v, want %#v", got, want)
}
if got, unwanted := tables, []string{"myothertable"}; containsAll(got, unwanted) {
t.Errorf("adminClient.Tables return %#v. unwanted %#v", got, unwanted)
}
tblConf := TableConf{
TableID: "conftable",
Families: map[string]GCPolicy{
"fam1": MaxVersionsPolicy(1),
"fam2": MaxVersionsPolicy(2),
},
}
if err := adminClient.CreateTableFromConf(ctx, &tblConf); err != nil {
t.Fatalf("Creating table from TableConf: %v", err)
}
defer adminClient.DeleteTable(ctx, tblConf.TableID)
tblInfo, err := adminClient.TableInfo(ctx, tblConf.TableID)
if err != nil {
t.Fatalf("Getting table info: %v", err)
}
sort.Strings(tblInfo.Families)
wantFams := []string{"fam1", "fam2"}
if !testutil.Equal(tblInfo.Families, wantFams) {
t.Errorf("Column family mismatch, got %v, want %v", tblInfo.Families, wantFams)
}
// Populate mytable and drop row ranges
if err = adminClient.CreateColumnFamily(ctx, "mytable", "cf"); err != nil {
t.Fatalf("Creating column family: %v", err)
}
client, err := testEnv.NewClient()
if err != nil {
t.Fatalf("NewClient: %v", err)
}
defer client.Close()
tbl := client.Open("mytable")
prefixes := []string{"a", "b", "c"}
for _, prefix := range prefixes {
for i := 0; i < 5; i++ {
mut := NewMutation()
mut.Set("cf", "col", 1000, []byte("1"))
if err := tbl.Apply(ctx, fmt.Sprintf("%v-%v", prefix, i), mut); err != nil {
t.Fatalf("Mutating row: %v", err)
}
}
}
if err = adminClient.DropRowRange(ctx, "mytable", "a"); err != nil {
t.Errorf("DropRowRange a: %v", err)
}
if err = adminClient.DropRowRange(ctx, "mytable", "c"); err != nil {
t.Errorf("DropRowRange c: %v", err)
}
if err = adminClient.DropRowRange(ctx, "mytable", "x"); err != nil {
t.Errorf("DropRowRange x: %v", err)
}
var gotRowCount int
must(tbl.ReadRows(ctx, RowRange{}, func(row Row) bool {
gotRowCount++
if !strings.HasPrefix(row.Key(), "b") {
t.Errorf("Invalid row after dropping range: %v", row)
}
return true
}))
if gotRowCount != 5 {
t.Errorf("Invalid row count after dropping range: got %v, want %v", gotRowCount, 5)
}
}
func TestInstanceUpdate(t *testing.T) {
testEnv, err := NewIntegrationEnv()
if err != nil {
t.Fatalf("IntegrationEnv: %v", err)
}
defer testEnv.Close()
timeout := 2 * time.Second
if testEnv.Config().UseProd {
timeout = 5 * time.Minute
}
ctx, cancel := context.WithTimeout(context.Background(), timeout)
defer cancel()
adminClient, err := testEnv.NewAdminClient()
if err != nil {
t.Fatalf("NewAdminClient: %v", err)
}
defer adminClient.Close()
iAdminClient, err := testEnv.NewInstanceAdminClient()
if err != nil {
t.Fatalf("NewInstanceAdminClient: %v", err)
}
if iAdminClient == nil {
return
}
defer iAdminClient.Close()
iInfo, err := iAdminClient.InstanceInfo(ctx, adminClient.instance)
if err != nil {
t.Errorf("InstanceInfo: %v", err)
}
if iInfo.Name != adminClient.instance {
t.Errorf("InstanceInfo returned name %#v, want %#v", iInfo.Name, adminClient.instance)
}
if iInfo.DisplayName != adminClient.instance {
t.Errorf("InstanceInfo returned name %#v, want %#v", iInfo.Name, adminClient.instance)
}
const numNodes = 4
// update cluster nodes
if err := iAdminClient.UpdateCluster(ctx, adminClient.instance, testEnv.Config().Cluster, int32(numNodes)); err != nil {
t.Errorf("UpdateCluster: %v", err)
}
// get cluster after updating
cis, err := iAdminClient.GetCluster(ctx, adminClient.instance, testEnv.Config().Cluster)
if err != nil {
t.Errorf("GetCluster %v", err)
}
if cis.ServeNodes != int(numNodes) {
t.Errorf("ServeNodes returned %d, want %d", cis.ServeNodes, int(numNodes))
}
}
func TestAdminSnapshotIntegration(t *testing.T) {
testEnv, err := NewIntegrationEnv()
if err != nil {
t.Fatalf("IntegrationEnv: %v", err)
}
defer testEnv.Close()
if !testEnv.Config().UseProd {
t.Skip("emulator doesn't support snapshots")
}
timeout := 2 * time.Second
if testEnv.Config().UseProd {
timeout = 5 * time.Minute
}
ctx, _ := context.WithTimeout(context.Background(), timeout)
adminClient, err := testEnv.NewAdminClient()
if err != nil {
t.Fatalf("NewAdminClient: %v", err)
}
defer adminClient.Close()
table := testEnv.Config().Table
cluster := testEnv.Config().Cluster
list := func(cluster string) ([]*SnapshotInfo, error) {
infos := []*SnapshotInfo(nil)
it := adminClient.Snapshots(ctx, cluster)
for {
s, err := it.Next()
if err == iterator.Done {
break
}
if err != nil {
return nil, err
}
infos = append(infos, s)
}
return infos, err
}
// Delete the table at the end of the test. Schedule ahead of time
// in case the client fails
defer adminClient.DeleteTable(ctx, table)
if err := adminClient.CreateTable(ctx, table); err != nil {
t.Fatalf("Creating table: %v", err)
}
// Precondition: no snapshots
snapshots, err := list(cluster)
if err != nil {
t.Fatalf("Initial snapshot list: %v", err)
}
if got, want := len(snapshots), 0; got != want {
t.Fatalf("Initial snapshot list len: %d, want: %d", got, want)
}
// Create snapshot
defer adminClient.DeleteSnapshot(ctx, cluster, "mysnapshot")
if err = adminClient.SnapshotTable(ctx, table, cluster, "mysnapshot", 5*time.Hour); err != nil {
t.Fatalf("Creating snaphot: %v", err)
}
// List snapshot
snapshots, err = list(cluster)
if err != nil {
t.Fatalf("Listing snapshots: %v", err)
}
if got, want := len(snapshots), 1; got != want {
t.Fatalf("Listing snapshot count: %d, want: %d", got, want)
}
if got, want := snapshots[0].Name, "mysnapshot"; got != want {
t.Fatalf("Snapshot name: %s, want: %s", got, want)
}
if got, want := snapshots[0].SourceTable, table; got != want {
t.Fatalf("Snapshot SourceTable: %s, want: %s", got, want)
}
if got, want := snapshots[0].DeleteTime, snapshots[0].CreateTime.Add(5*time.Hour); math.Abs(got.Sub(want).Minutes()) > 1 {
t.Fatalf("Snapshot DeleteTime: %s, want: %s", got, want)
}
// Get snapshot
snapshot, err := adminClient.SnapshotInfo(ctx, cluster, "mysnapshot")
if err != nil {
t.Fatalf("SnapshotInfo: %v", snapshot)
}
if got, want := *snapshot, *snapshots[0]; got != want {
t.Fatalf("SnapshotInfo: %v, want: %v", got, want)
}
// Restore
restoredTable := table + "-restored"
defer adminClient.DeleteTable(ctx, restoredTable)
if err = adminClient.CreateTableFromSnapshot(ctx, restoredTable, cluster, "mysnapshot"); err != nil {
t.Fatalf("CreateTableFromSnapshot: %v", err)
}
if _, err := adminClient.TableInfo(ctx, restoredTable); err != nil {
t.Fatalf("Restored TableInfo: %v", err)
}
// Delete snapshot
if err = adminClient.DeleteSnapshot(ctx, cluster, "mysnapshot"); err != nil {
t.Fatalf("DeleteSnapshot: %v", err)
}
snapshots, err = list(cluster)
if err != nil {
t.Fatalf("List after Delete: %v", err)
}
if got, want := len(snapshots), 0; got != want {
t.Fatalf("List after delete len: %d, want: %d", got, want)
}
}
func TestGranularity(t *testing.T) {
testEnv, err := NewIntegrationEnv()
if err != nil {
t.Fatalf("IntegrationEnv: %v", err)
}
defer testEnv.Close()
timeout := 2 * time.Second
if testEnv.Config().UseProd {
timeout = 5 * time.Minute
}
ctx, _ := context.WithTimeout(context.Background(), timeout)
adminClient, err := testEnv.NewAdminClient()
if err != nil {
t.Fatalf("NewAdminClient: %v", err)
}
defer adminClient.Close()
list := func() []string {
tbls, err := adminClient.Tables(ctx)
if err != nil {
t.Fatalf("Fetching list of tables: %v", err)
}
sort.Strings(tbls)
return tbls
}
containsAll := func(got, want []string) bool {
gotSet := make(map[string]bool)
for _, s := range got {
gotSet[s] = true
}
for _, s := range want {
if !gotSet[s] {
return false
}
}
return true
}
defer adminClient.DeleteTable(ctx, "mytable")
if err := adminClient.CreateTable(ctx, "mytable"); err != nil {
t.Fatalf("Creating table: %v", err)
}
tables := list()
if got, want := tables, []string{"mytable"}; !containsAll(got, want) {
t.Errorf("adminClient.Tables returned %#v, want %#v", got, want)
}
// calling ModifyColumnFamilies to check the granularity of table
prefix := adminClient.instancePrefix()
req := &btapb.ModifyColumnFamiliesRequest{
Name: prefix + "/tables/" + "mytable",
Modifications: []*btapb.ModifyColumnFamiliesRequest_Modification{{
Id: "cf",
Mod: &btapb.ModifyColumnFamiliesRequest_Modification_Create{&btapb.ColumnFamily{}},
}},
}
table, err := adminClient.tClient.ModifyColumnFamilies(ctx, req)
if err != nil {
t.Fatalf("Creating column family: %v", err)
}
if table.Granularity != btapb.Table_TimestampGranularity(btapb.Table_MILLIS) {
t.Errorf("ModifyColumnFamilies returned granularity %#v, want %#v", table.Granularity, btapb.Table_TimestampGranularity(btapb.Table_MILLIS))
}
}
func TestInstanceAdminClient_AppProfile(t *testing.T) {
testEnv, err := NewIntegrationEnv()
if err != nil {
t.Fatalf("IntegrationEnv: %v", err)
}
defer testEnv.Close()
timeout := 2 * time.Second
if testEnv.Config().UseProd {
timeout = 5 * time.Minute
}
ctx, cancel := context.WithTimeout(context.Background(), timeout)
defer cancel()
adminClient, err := testEnv.NewAdminClient()
if err != nil {
t.Fatalf("NewAdminClient: %v", err)
}
defer adminClient.Close()
iAdminClient, err := testEnv.NewInstanceAdminClient()
if err != nil {
t.Fatalf("NewInstanceAdminClient: %v", err)
}
if iAdminClient == nil {
return
}
defer iAdminClient.Close()
profile := ProfileConf{
ProfileID: "app_profile1",
InstanceID: adminClient.instance,
ClusterID: testEnv.Config().Cluster,
Description: "creating new app profile 1",
RoutingPolicy: SingleClusterRouting,
}
createdProfile, err := iAdminClient.CreateAppProfile(ctx, profile)
if err != nil {
t.Fatalf("Creating app profile: %v", err)
}
gotProfile, err := iAdminClient.GetAppProfile(ctx, adminClient.instance, "app_profile1")
if err != nil {
t.Fatalf("Get app profile: %v", err)
}
if !proto.Equal(createdProfile, gotProfile) {
t.Fatalf("created profile: %s, got profile: %s", createdProfile.Name, gotProfile.Name)
}
list := func(instanceID string) ([]*btapb.AppProfile, error) {
profiles := []*btapb.AppProfile(nil)
it := iAdminClient.ListAppProfiles(ctx, instanceID)
for {
s, err := it.Next()
if err == iterator.Done {
break
}
if err != nil {
return nil, err
}
profiles = append(profiles, s)
}
return profiles, err
}
profiles, err := list(adminClient.instance)
if err != nil {
t.Fatalf("List app profile: %v", err)
}
if got, want := len(profiles), 1; got != want {
t.Fatalf("Initial app profile list len: %d, want: %d", got, want)
}
for _, test := range []struct {
desc string
uattrs ProfileAttrsToUpdate
want *btapb.AppProfile // nil means error
}{
{
desc: "empty update",
uattrs: ProfileAttrsToUpdate{},
want: nil,
},
{
desc: "empty description update",
uattrs: ProfileAttrsToUpdate{Description: ""},
want: &btapb.AppProfile{
Name: gotProfile.Name,
Description: "",
RoutingPolicy: gotProfile.RoutingPolicy,
Etag: gotProfile.Etag},
},
{
desc: "routing update",
uattrs: ProfileAttrsToUpdate{
RoutingPolicy: SingleClusterRouting,
ClusterID: testEnv.Config().Cluster,
},
want: &btapb.AppProfile{
Name: gotProfile.Name,
Description: "",
Etag: gotProfile.Etag,
RoutingPolicy: &btapb.AppProfile_SingleClusterRouting_{
SingleClusterRouting: &btapb.AppProfile_SingleClusterRouting{
ClusterId: testEnv.Config().Cluster,
}},
},
},
} {
err = iAdminClient.UpdateAppProfile(ctx, adminClient.instance, "app_profile1", test.uattrs)
if err != nil {
if test.want != nil {
t.Errorf("%s: %v", test.desc, err)
}
continue
}
if err == nil && test.want == nil {
t.Errorf("%s: got nil, want error", test.desc)
continue
}
got, _ := iAdminClient.GetAppProfile(ctx, adminClient.instance, "app_profile1")
if !proto.Equal(got, test.want) {
t.Fatalf("%s : got profile : %v, want profile: %v", test.desc, gotProfile, test.want)
}
}
err = iAdminClient.DeleteAppProfile(ctx, adminClient.instance, "app_profile1")
if err != nil {
t.Fatalf("Delete app profile: %v", err)
}
}
/*
Copyright 2015 Google LLC
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
package bigtable // import "cloud.google.com/go/bigtable"
import (
"context"
"errors"
"fmt"
"io"
"strconv"
"time"
"cloud.google.com/go/bigtable/internal/gax"
btopt "cloud.google.com/go/bigtable/internal/option"
"github.com/golang/protobuf/proto"
"google.golang.org/api/option"
gtransport "google.golang.org/api/transport/grpc"
btpb "google.golang.org/genproto/googleapis/bigtable/v2"
"google.golang.org/grpc"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/metadata"
"google.golang.org/grpc/status"
)
const prodAddr = "bigtable.googleapis.com:443"
// Client is a client for reading and writing data to tables in an instance.
//
// A Client is safe to use concurrently, except for its Close method.
type Client struct {
conn *grpc.ClientConn
client btpb.BigtableClient
project, instance string
appProfile string
}
// ClientConfig has configurations for the client.
type ClientConfig struct {
// The id of the app profile to associate with all data operations sent from this client.
// If unspecified, the default app profile for the instance will be used.
AppProfile string
}
// NewClient creates a new Client for a given project and instance.
// The default ClientConfig will be used.
func NewClient(ctx context.Context, project, instance string, opts ...option.ClientOption) (*Client, error) {
return NewClientWithConfig(ctx, project, instance, ClientConfig{}, opts...)
}
// NewClientWithConfig creates a new client with the given config.
func NewClientWithConfig(ctx context.Context, project, instance string, config ClientConfig, opts ...option.ClientOption) (*Client, error) {
o, err := btopt.DefaultClientOptions(prodAddr, Scope, clientUserAgent)
if err != nil {
return nil, err
}
// Default to a small connection pool that can be overridden.
o = append(o,
option.WithGRPCConnectionPool(4),
// Set the max size to correspond to server-side limits.
option.WithGRPCDialOption(grpc.WithDefaultCallOptions(grpc.MaxCallSendMsgSize(100<<20), grpc.MaxCallRecvMsgSize(100<<20))),
// TODO(grpc/grpc-go#1388) using connection pool without WithBlock
// can cause RPCs to fail randomly. We can delete this after the issue is fixed.
option.WithGRPCDialOption(grpc.WithBlock()))
o = append(o, opts...)
conn, err := gtransport.Dial(ctx, o...)
if err != nil {
return nil, fmt.Errorf("dialing: %v", err)
}
return &Client{
conn: conn,
client: btpb.NewBigtableClient(conn),
project: project,
instance: instance,
appProfile: config.AppProfile,
}, nil
}
// Close closes the Client.
func (c *Client) Close() error {
return c.conn.Close()
}
var (
idempotentRetryCodes = []codes.Code{codes.DeadlineExceeded, codes.Unavailable, codes.Aborted}
isIdempotentRetryCode = make(map[codes.Code]bool)
retryOptions = []gax.CallOption{
gax.WithDelayTimeoutSettings(100*time.Millisecond, 2000*time.Millisecond, 1.2),
gax.WithRetryCodes(idempotentRetryCodes),
}
)
func init() {
for _, code := range idempotentRetryCodes {
isIdempotentRetryCode[code] = true
}
}
func (c *Client) fullTableName(table string) string {
return fmt.Sprintf("projects/%s/instances/%s/tables/%s", c.project, c.instance, table)
}
// A Table refers to a table.
//
// A Table is safe to use concurrently.
type Table struct {
c *Client
table string
// Metadata to be sent with each request.
md metadata.MD
}
// Open opens a table.
func (c *Client) Open(table string) *Table {
return &Table{
c: c,
table: table,
md: metadata.Pairs(resourcePrefixHeader, c.fullTableName(table)),
}
}
// TODO(dsymonds): Read method that returns a sequence of ReadItems.
// ReadRows reads rows from a table. f is called for each row.
// If f returns false, the stream is shut down and ReadRows returns.
// f owns its argument, and f is called serially in order by row key.
//
// By default, the yielded rows will contain all values in all cells.
// Use RowFilter to limit the cells returned.
func (t *Table) ReadRows(ctx context.Context, arg RowSet, f func(Row) bool, opts ...ReadOption) error {
ctx = mergeOutgoingMetadata(ctx, t.md)
var prevRowKey string
var err error
ctx = traceStartSpan(ctx, "cloud.google.com/go/bigtable.ReadRows")
defer func() { traceEndSpan(ctx, err) }()
attrMap := make(map[string]interface{})
err = gax.Invoke(ctx, func(ctx context.Context) error {
if !arg.valid() {
// Empty row set, no need to make an API call.
// NOTE: we must return early if arg == RowList{} because reading
// an empty RowList from bigtable returns all rows from that table.
return nil
}
req := &btpb.ReadRowsRequest{
TableName: t.c.fullTableName(t.table),
AppProfileId: t.c.appProfile,
Rows: arg.proto(),
}
for _, opt := range opts {
opt.set(req)
}
ctx, cancel := context.WithCancel(ctx) // for aborting the stream
defer cancel()
startTime := time.Now()
stream, err := t.c.client.ReadRows(ctx, req)
if err != nil {
return err
}
cr := newChunkReader()
for {
res, err := stream.Recv()
if err == io.EOF {
break
}
if err != nil {
// Reset arg for next Invoke call.
arg = arg.retainRowsAfter(prevRowKey)
attrMap["rowKey"] = prevRowKey
attrMap["error"] = err.Error()
attrMap["time_secs"] = time.Since(startTime).Seconds()
tracePrintf(ctx, attrMap, "Retry details in ReadRows")
return err
}
attrMap["time_secs"] = time.Since(startTime).Seconds()
attrMap["rowCount"] = len(res.Chunks)
tracePrintf(ctx, attrMap, "Details in ReadRows")
for _, cc := range res.Chunks {
row, err := cr.Process(cc)
if err != nil {
// No need to prepare for a retry, this is an unretryable error.
return err
}
if row == nil {
continue
}
prevRowKey = row.Key()
if !f(row) {
// Cancel and drain stream.
cancel()
for {
if _, err := stream.Recv(); err != nil {
// The stream has ended. We don't return an error
// because the caller has intentionally interrupted the scan.
return nil
}
}
}
}
if err := cr.Close(); err != nil {
// No need to prepare for a retry, this is an unretryable error.
return err
}
}
return err
}, retryOptions...)
return err
}
// ReadRow is a convenience implementation of a single-row reader.
// A missing row will return a zero-length map and a nil error.
func (t *Table) ReadRow(ctx context.Context, row string, opts ...ReadOption) (Row, error) {
var r Row
err := t.ReadRows(ctx, SingleRow(row), func(rr Row) bool {
r = rr
return true
}, opts...)
return r, err
}
// decodeFamilyProto adds the cell data from f to the given row.
func decodeFamilyProto(r Row, row string, f *btpb.Family) {
fam := f.Name // does not have colon
for _, col := range f.Columns {
for _, cell := range col.Cells {
ri := ReadItem{
Row: row,
Column: fam + ":" + string(col.Qualifier),
Timestamp: Timestamp(cell.TimestampMicros),
Value: cell.Value,
}
r[fam] = append(r[fam], ri)
}
}
}
// RowSet is a set of rows to be read. It is satisfied by RowList, RowRange and RowRangeList.
// The serialized size of the RowSet must be no larger than 1MiB.
type RowSet interface {
proto() *btpb.RowSet
// retainRowsAfter returns a new RowSet that does not include the
// given row key or any row key lexicographically less than it.
retainRowsAfter(lastRowKey string) RowSet
// Valid reports whether this set can cover at least one row.
valid() bool
}
// RowList is a sequence of row keys.
type RowList []string
func (r RowList) proto() *btpb.RowSet {
keys := make([][]byte, len(r))
for i, row := range r {
keys[i] = []byte(row)
}
return &btpb.RowSet{RowKeys: keys}
}
func (r RowList) retainRowsAfter(lastRowKey string) RowSet {
var retryKeys RowList
for _, key := range r {
if key > lastRowKey {
retryKeys = append(retryKeys, key)
}
}
return retryKeys
}
func (r RowList) valid() bool {
return len(r) > 0
}
// A RowRange is a half-open interval [Start, Limit) encompassing
// all the rows with keys at least as large as Start, and less than Limit.
// (Bigtable string comparison is the same as Go's.)
// A RowRange can be unbounded, encompassing all keys at least as large as Start.
type RowRange struct {
start string
limit string
}
// NewRange returns the new RowRange [begin, end).
func NewRange(begin, end string) RowRange {
return RowRange{
start: begin,
limit: end,
}
}
// Unbounded tests whether a RowRange is unbounded.
func (r RowRange) Unbounded() bool {
return r.limit == ""
}
// Contains says whether the RowRange contains the key.
func (r RowRange) Contains(row string) bool {
return r.start <= row && (r.limit == "" || r.limit > row)
}
// String provides a printable description of a RowRange.
func (r RowRange) String() string {
a := strconv.Quote(r.start)
if r.Unbounded() {
return fmt.Sprintf("[%s,∞)", a)
}
return fmt.Sprintf("[%s,%q)", a, r.limit)
}
func (r RowRange) proto() *btpb.RowSet {
rr := &btpb.RowRange{
StartKey: &btpb.RowRange_StartKeyClosed{StartKeyClosed: []byte(r.start)},
}
if !r.Unbounded() {
rr.EndKey = &btpb.RowRange_EndKeyOpen{EndKeyOpen: []byte(r.limit)}
}
return &btpb.RowSet{RowRanges: []*btpb.RowRange{rr}}
}
func (r RowRange) retainRowsAfter(lastRowKey string) RowSet {
if lastRowKey == "" || lastRowKey < r.start {
return r
}
// Set the beginning of the range to the row after the last scanned.
start := lastRowKey + "\x00"
if r.Unbounded() {
return InfiniteRange(start)
}
return NewRange(start, r.limit)
}
func (r RowRange) valid() bool {
return r.Unbounded() || r.start < r.limit
}
// RowRangeList is a sequence of RowRanges representing the union of the ranges.
type RowRangeList []RowRange
func (r RowRangeList) proto() *btpb.RowSet {
ranges := make([]*btpb.RowRange, len(r))
for i, rr := range r {
// RowRange.proto() returns a RowSet with a single element RowRange array
ranges[i] = rr.proto().RowRanges[0]
}
return &btpb.RowSet{RowRanges: ranges}
}
func (r RowRangeList) retainRowsAfter(lastRowKey string) RowSet {
if lastRowKey == "" {
return r
}
// Return a list of any range that has not yet been completely processed
var ranges RowRangeList
for _, rr := range r {
retained := rr.retainRowsAfter(lastRowKey)
if retained.valid() {
ranges = append(ranges, retained.(RowRange))
}
}
return ranges
}
func (r RowRangeList) valid() bool {
for _, rr := range r {
if rr.valid() {
return true
}
}
return false
}
// SingleRow returns a RowSet for reading a single row.
func SingleRow(row string) RowSet {
return RowList{row}
}
// PrefixRange returns a RowRange consisting of all keys starting with the prefix.
func PrefixRange(prefix string) RowRange {
return RowRange{
start: prefix,
limit: prefixSuccessor(prefix),
}
}
// InfiniteRange returns the RowRange consisting of all keys at least as
// large as start.
func InfiniteRange(start string) RowRange {
return RowRange{
start: start,
limit: "",
}
}
// prefixSuccessor returns the lexically smallest string greater than the
// prefix, if it exists, or "" otherwise. In either case, it is the string
// needed for the Limit of a RowRange.
func prefixSuccessor(prefix string) string {
if prefix == "" {
return "" // infinite range
}
n := len(prefix)
for n--; n >= 0 && prefix[n] == '\xff'; n-- {
}
if n == -1 {
return ""
}
ans := []byte(prefix[:n])
ans = append(ans, prefix[n]+1)
return string(ans)
}
// A ReadOption is an optional argument to ReadRows.
type ReadOption interface {
set(req *btpb.ReadRowsRequest)
}
// RowFilter returns a ReadOption that applies f to the contents of read rows.
//
// If multiple RowFilters are provided, only the last is used. To combine filters,
// use ChainFilters or InterleaveFilters instead.
func RowFilter(f Filter) ReadOption { return rowFilter{f} }
type rowFilter struct{ f Filter }
func (rf rowFilter) set(req *btpb.ReadRowsRequest) { req.Filter = rf.f.proto() }
// LimitRows returns a ReadOption that will limit the number of rows to be read.
func LimitRows(limit int64) ReadOption { return limitRows{limit} }
type limitRows struct{ limit int64 }
func (lr limitRows) set(req *btpb.ReadRowsRequest) { req.RowsLimit = lr.limit }
// mutationsAreRetryable returns true if all mutations are idempotent
// and therefore retryable. A mutation is idempotent iff all cell timestamps
// have an explicit timestamp set and do not rely on the timestamp being set on the server.
func mutationsAreRetryable(muts []*btpb.Mutation) bool {
serverTime := int64(ServerTime)
for _, mut := range muts {
setCell := mut.GetSetCell()
if setCell != nil && setCell.TimestampMicros == serverTime {
return false
}
}
return true
}
const maxMutations = 100000
// Apply mutates a row atomically. A mutation must contain at least one
// operation and at most 100000 operations.
func (t *Table) Apply(ctx context.Context, row string, m *Mutation, opts ...ApplyOption) error {
ctx = mergeOutgoingMetadata(ctx, t.md)
after := func(res proto.Message) {
for _, o := range opts {
o.after(res)
}
}
var err error
ctx = traceStartSpan(ctx, "cloud.google.com/go/bigtable/Apply")
defer func() { traceEndSpan(ctx, err) }()
var callOptions []gax.CallOption
if m.cond == nil {
req := &btpb.MutateRowRequest{
TableName: t.c.fullTableName(t.table),
AppProfileId: t.c.appProfile,
RowKey: []byte(row),
Mutations: m.ops,
}
if mutationsAreRetryable(m.ops) {
callOptions = retryOptions
}
var res *btpb.MutateRowResponse
err := gax.Invoke(ctx, func(ctx context.Context) error {
var err error
res, err = t.c.client.MutateRow(ctx, req)
return err
}, callOptions...)
if err == nil {
after(res)
}
return err
}
req := &btpb.CheckAndMutateRowRequest{
TableName: t.c.fullTableName(t.table),
AppProfileId: t.c.appProfile,
RowKey: []byte(row),
PredicateFilter: m.cond.proto(),
}
if m.mtrue != nil {
if m.mtrue.cond != nil {
return errors.New("bigtable: conditional mutations cannot be nested")
}
req.TrueMutations = m.mtrue.ops
}
if m.mfalse != nil {
if m.mfalse.cond != nil {
return errors.New("bigtable: conditional mutations cannot be nested")
}
req.FalseMutations = m.mfalse.ops
}
if mutationsAreRetryable(req.TrueMutations) && mutationsAreRetryable(req.FalseMutations) {
callOptions = retryOptions
}
var cmRes *btpb.CheckAndMutateRowResponse
err = gax.Invoke(ctx, func(ctx context.Context) error {
var err error
cmRes, err = t.c.client.CheckAndMutateRow(ctx, req)
return err
}, callOptions...)
if err == nil {
after(cmRes)
}
return err
}
// An ApplyOption is an optional argument to Apply.
type ApplyOption interface {
after(res proto.Message)
}
type applyAfterFunc func(res proto.Message)
func (a applyAfterFunc) after(res proto.Message) { a(res) }
// GetCondMutationResult returns an ApplyOption that reports whether the conditional
// mutation's condition matched.
func GetCondMutationResult(matched *bool) ApplyOption {
return applyAfterFunc(func(res proto.Message) {
if res, ok := res.(*btpb.CheckAndMutateRowResponse); ok {
*matched = res.PredicateMatched
}
})
}
// Mutation represents a set of changes for a single row of a table.
type Mutation struct {
ops []*btpb.Mutation
// for conditional mutations
cond Filter
mtrue, mfalse *Mutation
}
// NewMutation returns a new mutation.
func NewMutation() *Mutation {
return new(Mutation)
}
// NewCondMutation returns a conditional mutation.
// The given row filter determines which mutation is applied:
// If the filter matches any cell in the row, mtrue is applied;
// otherwise, mfalse is applied.
// Either given mutation may be nil.
func NewCondMutation(cond Filter, mtrue, mfalse *Mutation) *Mutation {
return &Mutation{cond: cond, mtrue: mtrue, mfalse: mfalse}
}
// Set sets a value in a specified column, with the given timestamp.
// The timestamp will be truncated to millisecond granularity.
// A timestamp of ServerTime means to use the server timestamp.
func (m *Mutation) Set(family, column string, ts Timestamp, value []byte) {
m.ops = append(m.ops, &btpb.Mutation{Mutation: &btpb.Mutation_SetCell_{SetCell: &btpb.Mutation_SetCell{
FamilyName: family,
ColumnQualifier: []byte(column),
TimestampMicros: int64(ts.TruncateToMilliseconds()),
Value: value,
}}})
}
// DeleteCellsInColumn will delete all the cells whose columns are family:column.
func (m *Mutation) DeleteCellsInColumn(family, column string) {
m.ops = append(m.ops, &btpb.Mutation{Mutation: &btpb.Mutation_DeleteFromColumn_{DeleteFromColumn: &btpb.Mutation_DeleteFromColumn{
FamilyName: family,
ColumnQualifier: []byte(column),
}}})
}
// DeleteTimestampRange deletes all cells whose columns are family:column
// and whose timestamps are in the half-open interval [start, end).
// If end is zero, it will be interpreted as infinity.
// The timestamps will be truncated to millisecond granularity.
func (m *Mutation) DeleteTimestampRange(family, column string, start, end Timestamp) {
m.ops = append(m.ops, &btpb.Mutation{Mutation: &btpb.Mutation_DeleteFromColumn_{DeleteFromColumn: &btpb.Mutation_DeleteFromColumn{
FamilyName: family,
ColumnQualifier: []byte(column),
TimeRange: &btpb.TimestampRange{
StartTimestampMicros: int64(start.TruncateToMilliseconds()),
EndTimestampMicros: int64(end.TruncateToMilliseconds()),
},
}}})
}
// DeleteCellsInFamily will delete all the cells whose columns are family:*.
func (m *Mutation) DeleteCellsInFamily(family string) {
m.ops = append(m.ops, &btpb.Mutation{Mutation: &btpb.Mutation_DeleteFromFamily_{DeleteFromFamily: &btpb.Mutation_DeleteFromFamily{
FamilyName: family,
}}})
}
// DeleteRow deletes the entire row.
func (m *Mutation) DeleteRow() {
m.ops = append(m.ops, &btpb.Mutation{Mutation: &btpb.Mutation_DeleteFromRow_{DeleteFromRow: &btpb.Mutation_DeleteFromRow{}}})
}
// entryErr is a container that combines an entry with the error that was returned for it.
// Err may be nil if no error was returned for the Entry, or if the Entry has not yet been processed.
type entryErr struct {
Entry *btpb.MutateRowsRequest_Entry
Err error
}
// ApplyBulk applies multiple Mutations, up to a maximum of 100,000.
// Each mutation is individually applied atomically,
// but the set of mutations may be applied in any order.
//
// Two types of failures may occur. If the entire process
// fails, (nil, err) will be returned. If specific mutations
// fail to apply, ([]err, nil) will be returned, and the errors
// will correspond to the relevant rowKeys/muts arguments.
//
// Conditional mutations cannot be applied in bulk and providing one will result in an error.
func (t *Table) ApplyBulk(ctx context.Context, rowKeys []string, muts []*Mutation, opts ...ApplyOption) ([]error, error) {
ctx = mergeOutgoingMetadata(ctx, t.md)
if len(rowKeys) != len(muts) {
return nil, fmt.Errorf("mismatched rowKeys and mutation array lengths: %d, %d", len(rowKeys), len(muts))
}
origEntries := make([]*entryErr, len(rowKeys))
for i, key := range rowKeys {
mut := muts[i]
if mut.cond != nil {
return nil, errors.New("conditional mutations cannot be applied in bulk")
}
origEntries[i] = &entryErr{Entry: &btpb.MutateRowsRequest_Entry{RowKey: []byte(key), Mutations: mut.ops}}
}
var err error
ctx = traceStartSpan(ctx, "cloud.google.com/go/bigtable/ApplyBulk")
defer func() { traceEndSpan(ctx, err) }()
for _, group := range groupEntries(origEntries, maxMutations) {
attrMap := make(map[string]interface{})
err = gax.Invoke(ctx, func(ctx context.Context) error {
attrMap["rowCount"] = len(group)
tracePrintf(ctx, attrMap, "Row count in ApplyBulk")
err := t.doApplyBulk(ctx, group, opts...)
if err != nil {
// We want to retry the entire request with the current group
return err
}
group = t.getApplyBulkRetries(group)
if len(group) > 0 && len(idempotentRetryCodes) > 0 {
// We have at least one mutation that needs to be retried.
// Return an arbitrary error that is retryable according to callOptions.
return status.Errorf(idempotentRetryCodes[0], "Synthetic error: partial failure of ApplyBulk")
}
return nil
}, retryOptions...)
if err != nil {
return nil, err
}
}
// Accumulate all of the errors into an array to return, interspersed with nils for successful
// entries. The absence of any errors means we should return nil.
var errs []error
var foundErr bool
for _, entry := range origEntries {
if entry.Err != nil {
foundErr = true
}
errs = append(errs, entry.Err)
}
if foundErr {
return errs, nil
}
return nil, nil
}
// getApplyBulkRetries returns the entries that need to be retried
func (t *Table) getApplyBulkRetries(entries []*entryErr) []*entryErr {
var retryEntries []*entryErr
for _, entry := range entries {
err := entry.Err
if err != nil && isIdempotentRetryCode[grpc.Code(err)] && mutationsAreRetryable(entry.Entry.Mutations) {
// There was an error and the entry is retryable.
retryEntries = append(retryEntries, entry)
}
}
return retryEntries
}
// doApplyBulk does the work of a single ApplyBulk invocation
func (t *Table) doApplyBulk(ctx context.Context, entryErrs []*entryErr, opts ...ApplyOption) error {
after := func(res proto.Message) {
for _, o := range opts {
o.after(res)
}
}
entries := make([]*btpb.MutateRowsRequest_Entry, len(entryErrs))
for i, entryErr := range entryErrs {
entries[i] = entryErr.Entry
}
req := &btpb.MutateRowsRequest{
TableName: t.c.fullTableName(t.table),
AppProfileId: t.c.appProfile,
Entries: entries,
}
stream, err := t.c.client.MutateRows(ctx, req)
if err != nil {
return err
}
for {
res, err := stream.Recv()
if err == io.EOF {
break
}
if err != nil {
return err
}
for i, entry := range res.Entries {
s := entry.Status
if s.Code == int32(codes.OK) {
entryErrs[i].Err = nil
} else {
entryErrs[i].Err = status.Errorf(codes.Code(s.Code), s.Message)
}
}
after(res)
}
return nil
}
// groupEntries groups entries into groups of a specified size without breaking up
// individual entries.
func groupEntries(entries []*entryErr, maxSize int) [][]*entryErr {
var (
res [][]*entryErr
start int
gmuts int
)
addGroup := func(end int) {
if end-start > 0 {
res = append(res, entries[start:end])
start = end
gmuts = 0
}
}
for i, e := range entries {
emuts := len(e.Entry.Mutations)
if gmuts+emuts > maxSize {
addGroup(i)
}
gmuts += emuts
}
addGroup(len(entries))
return res
}
// Timestamp is in units of microseconds since 1 January 1970.
type Timestamp int64
// ServerTime is a specific Timestamp that may be passed to (*Mutation).Set.
// It indicates that the server's timestamp should be used.
const ServerTime Timestamp = -1
// Time converts a time.Time into a Timestamp.
func Time(t time.Time) Timestamp { return Timestamp(t.UnixNano() / 1e3) }
// Now returns the Timestamp representation of the current time on the client.
func Now() Timestamp { return Time(time.Now()) }
// Time converts a Timestamp into a time.Time.
func (ts Timestamp) Time() time.Time { return time.Unix(0, int64(ts)*1e3) }
// TruncateToMilliseconds truncates a Timestamp to millisecond granularity,
// which is currently the only granularity supported.
func (ts Timestamp) TruncateToMilliseconds() Timestamp {
if ts == ServerTime {
return ts
}
return ts - ts%1000
}
// ApplyReadModifyWrite applies a ReadModifyWrite to a specific row.
// It returns the newly written cells.
func (t *Table) ApplyReadModifyWrite(ctx context.Context, row string, m *ReadModifyWrite) (Row, error) {
ctx = mergeOutgoingMetadata(ctx, t.md)
req := &btpb.ReadModifyWriteRowRequest{
TableName: t.c.fullTableName(t.table),
AppProfileId: t.c.appProfile,
RowKey: []byte(row),
Rules: m.ops,
}
res, err := t.c.client.ReadModifyWriteRow(ctx, req)
if err != nil {
return nil, err
}
if res.Row == nil {
return nil, errors.New("unable to apply ReadModifyWrite: res.Row=nil")
}
r := make(Row)
for _, fam := range res.Row.Families { // res is *btpb.Row, fam is *btpb.Family
decodeFamilyProto(r, row, fam)
}
return r, nil
}
// ReadModifyWrite represents a set of operations on a single row of a table.
// It is like Mutation but for non-idempotent changes.
// When applied, these operations operate on the latest values of the row's cells,
// and result in a new value being written to the relevant cell with a timestamp
// that is max(existing timestamp, current server time).
//
// The application of a ReadModifyWrite is atomic; concurrent ReadModifyWrites will
// be executed serially by the server.
type ReadModifyWrite struct {
ops []*btpb.ReadModifyWriteRule
}
// NewReadModifyWrite returns a new ReadModifyWrite.
func NewReadModifyWrite() *ReadModifyWrite { return new(ReadModifyWrite) }
// AppendValue appends a value to a specific cell's value.
// If the cell is unset, it will be treated as an empty value.
func (m *ReadModifyWrite) AppendValue(family, column string, v []byte) {
m.ops = append(m.ops, &btpb.ReadModifyWriteRule{
FamilyName: family,
ColumnQualifier: []byte(column),
Rule: &btpb.ReadModifyWriteRule_AppendValue{AppendValue: v},
})
}
// Increment interprets the value in a specific cell as a 64-bit big-endian signed integer,
// and adds a value to it. If the cell is unset, it will be treated as zero.
// If the cell is set and is not an 8-byte value, the entire ApplyReadModifyWrite
// operation will fail.
func (m *ReadModifyWrite) Increment(family, column string, delta int64) {
m.ops = append(m.ops, &btpb.ReadModifyWriteRule{
FamilyName: family,
ColumnQualifier: []byte(column),
Rule: &btpb.ReadModifyWriteRule_IncrementAmount{IncrementAmount: delta},
})
}
// mergeOutgoingMetadata returns a context populated by the existing outgoing metadata,
// if any, joined with internal metadata.
func mergeOutgoingMetadata(ctx context.Context, md metadata.MD) context.Context {
mdCopy, _ := metadata.FromOutgoingContext(ctx)
return metadata.NewOutgoingContext(ctx, metadata.Join(mdCopy, md))
}
// SampleRowKeys returns a sample of row keys in the table. The returned row keys will delimit contiguous sections of
// the table of approximately equal size, which can be used to break up the data for distributed tasks like mapreduces.
func (t *Table) SampleRowKeys(ctx context.Context) ([]string, error) {
ctx = mergeOutgoingMetadata(ctx, t.md)
var sampledRowKeys []string
err := gax.Invoke(ctx, func(ctx context.Context) error {
sampledRowKeys = nil
req := &btpb.SampleRowKeysRequest{
TableName: t.c.fullTableName(t.table),
AppProfileId: t.c.appProfile,
}
ctx, cancel := context.WithCancel(ctx) // for aborting the stream
defer cancel()
stream, err := t.c.client.SampleRowKeys(ctx, req)
if err != nil {
return err
}
for {
res, err := stream.Recv()
if err == io.EOF {
break
}
if err != nil {
return err
}
key := string(res.RowKey)
if key == "" {
continue
}
sampledRowKeys = append(sampledRowKeys, key)
}
return nil
}, retryOptions...)
return sampledRowKeys, err
}
/*
Copyright 2015 Google LLC
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
package bigtable
import (
"context"
"fmt"
"math/rand"
"strings"
"sync"
"testing"
"time"
"cloud.google.com/go/internal/testutil"
"github.com/google/go-cmp/cmp"
"google.golang.org/api/option"
btpb "google.golang.org/genproto/googleapis/bigtable/v2"
"google.golang.org/grpc"
)
func TestPrefix(t *testing.T) {
tests := []struct {
prefix, succ string
}{
{"", ""},
{"\xff", ""}, // when used, "" means Infinity
{"x\xff", "y"},
{"\xfe", "\xff"},
}
for _, tc := range tests {
got := prefixSuccessor(tc.prefix)
if got != tc.succ {
t.Errorf("prefixSuccessor(%q) = %q, want %s", tc.prefix, got, tc.succ)
continue
}
r := PrefixRange(tc.prefix)
if tc.succ == "" && r.limit != "" {
t.Errorf("PrefixRange(%q) got limit %q", tc.prefix, r.limit)
}
if tc.succ != "" && r.limit != tc.succ {
t.Errorf("PrefixRange(%q) got limit %q, want %q", tc.prefix, r.limit, tc.succ)
}
}
}
func TestApplyErrors(t *testing.T) {
ctx := context.Background()
table := &Table{
c: &Client{
project: "P",
instance: "I",
},
table: "t",
}
f := ColumnFilter("C")
m := NewMutation()
m.DeleteRow()
// Test nested conditional mutations.
cm := NewCondMutation(f, NewCondMutation(f, m, nil), nil)
if err := table.Apply(ctx, "x", cm); err == nil {
t.Error("got nil, want error")
}
cm = NewCondMutation(f, nil, NewCondMutation(f, m, nil))
if err := table.Apply(ctx, "x", cm); err == nil {
t.Error("got nil, want error")
}
}
func TestGroupEntries(t *testing.T) {
tests := []struct {
desc string
in []*entryErr
size int
want [][]*entryErr
}{
{
desc: "one entry less than max size is one group",
in: []*entryErr{buildEntry(5)},
size: 10,
want: [][]*entryErr{{buildEntry(5)}},
},
{
desc: "one entry equal to max size is one group",
in: []*entryErr{buildEntry(10)},
size: 10,
want: [][]*entryErr{{buildEntry(10)}},
},
{
desc: "one entry greater than max size is one group",
in: []*entryErr{buildEntry(15)},
size: 10,
want: [][]*entryErr{{buildEntry(15)}},
},
{
desc: "all entries fitting within max size are one group",
in: []*entryErr{buildEntry(10), buildEntry(10)},
size: 20,
want: [][]*entryErr{{buildEntry(10), buildEntry(10)}},
},
{
desc: "entries each under max size and together over max size are grouped separately",
in: []*entryErr{buildEntry(10), buildEntry(10)},
size: 15,
want: [][]*entryErr{{buildEntry(10)}, {buildEntry(10)}},
},
{
desc: "entries together over max size are grouped by max size",
in: []*entryErr{buildEntry(5), buildEntry(5), buildEntry(5)},
size: 10,
want: [][]*entryErr{{buildEntry(5), buildEntry(5)}, {buildEntry(5)}},
},
{
desc: "one entry over max size and one entry under max size are two groups",
in: []*entryErr{buildEntry(15), buildEntry(5)},
size: 10,
want: [][]*entryErr{{buildEntry(15)}, {buildEntry(5)}},
},
}
for _, test := range tests {
if got, want := groupEntries(test.in, test.size), test.want; !cmp.Equal(mutationCounts(got), mutationCounts(want)) {
t.Errorf("[%s] want = %v, got = %v", test.desc, mutationCounts(want), mutationCounts(got))
}
}
}
func buildEntry(numMutations int) *entryErr {
var muts []*btpb.Mutation
for i := 0; i < numMutations; i++ {
muts = append(muts, &btpb.Mutation{})
}
return &entryErr{Entry: &btpb.MutateRowsRequest_Entry{Mutations: muts}}
}
func mutationCounts(batched [][]*entryErr) []int {
var res []int
for _, entries := range batched {
var count int
for _, e := range entries {
count += len(e.Entry.Mutations)
}
res = append(res, count)
}
return res
}
func TestClientIntegration(t *testing.T) {
// TODO(jba): go1.9: Use subtests.
start := time.Now()
lastCheckpoint := start
checkpoint := func(s string) {
n := time.Now()
t.Logf("[%s] %v since start, %v since last checkpoint", s, n.Sub(start), n.Sub(lastCheckpoint))
lastCheckpoint = n
}
testEnv, err := NewIntegrationEnv()
if err != nil {
t.Fatalf("IntegrationEnv: %v", err)
}
var timeout time.Duration
if testEnv.Config().UseProd {
timeout = 10 * time.Minute
t.Logf("Running test against production")
} else {
timeout = 1 * time.Minute
t.Logf("bttest.Server running on %s", testEnv.Config().AdminEndpoint)
}
ctx, cancel := context.WithTimeout(context.Background(), timeout)
defer cancel()
client, err := testEnv.NewClient()
if err != nil {
t.Fatalf("Client: %v", err)
}
defer client.Close()
checkpoint("dialed Client")
adminClient, err := testEnv.NewAdminClient()
if err != nil {
t.Fatalf("AdminClient: %v", err)
}
defer adminClient.Close()
checkpoint("dialed AdminClient")
table := testEnv.Config().Table
// Delete the table at the end of the test.
// Do this even before creating the table so that if this is running
// against production and CreateTable fails there's a chance of cleaning it up.
defer adminClient.DeleteTable(ctx, table)
if err := adminClient.CreateTable(ctx, table); err != nil {
t.Fatalf("Creating table: %v", err)
}
checkpoint("created table")
if err := adminClient.CreateColumnFamily(ctx, table, "follows"); err != nil {
t.Fatalf("Creating column family: %v", err)
}
checkpoint(`created "follows" column family`)
tbl := client.Open(table)
// Insert some data.
initialData := map[string][]string{
"wmckinley": {"tjefferson"},
"gwashington": {"jadams"},
"tjefferson": {"gwashington", "jadams"}, // wmckinley set conditionally below
"jadams": {"gwashington", "tjefferson"},
}
for row, ss := range initialData {
mut := NewMutation()
for _, name := range ss {
mut.Set("follows", name, 1000, []byte("1"))
}
if err := tbl.Apply(ctx, row, mut); err != nil {
t.Errorf("Mutating row %q: %v", row, err)
}
}
checkpoint("inserted initial data")
// TODO(igorbernstein): re-enable this when ready
//if err := adminClient.WaitForReplication(ctx, table); err != nil {
// t.Errorf("Waiting for replication for table %q: %v", table, err)
//}
//checkpoint("waited for replication")
// Do a conditional mutation with a complex filter.
mutTrue := NewMutation()
mutTrue.Set("follows", "wmckinley", 1000, []byte("1"))
filter := ChainFilters(ColumnFilter("gwash[iz].*"), ValueFilter("."))
mut := NewCondMutation(filter, mutTrue, nil)
if err := tbl.Apply(ctx, "tjefferson", mut); err != nil {
t.Errorf("Conditionally mutating row: %v", err)
}
// Do a second condition mutation with a filter that does not match,
// and thus no changes should be made.
mutTrue = NewMutation()
mutTrue.DeleteRow()
filter = ColumnFilter("snoop.dogg")
mut = NewCondMutation(filter, mutTrue, nil)
if err := tbl.Apply(ctx, "tjefferson", mut); err != nil {
t.Errorf("Conditionally mutating row: %v", err)
}
checkpoint("did two conditional mutations")
// Fetch a row.
row, err := tbl.ReadRow(ctx, "jadams")
if err != nil {
t.Fatalf("Reading a row: %v", err)
}
wantRow := Row{
"follows": []ReadItem{
{Row: "jadams", Column: "follows:gwashington", Timestamp: 1000, Value: []byte("1")},
{Row: "jadams", Column: "follows:tjefferson", Timestamp: 1000, Value: []byte("1")},
},
}
if !testutil.Equal(row, wantRow) {
t.Errorf("Read row mismatch.\n got %#v\nwant %#v", row, wantRow)
}
checkpoint("tested ReadRow")
// Do a bunch of reads with filters.
readTests := []struct {
desc string
rr RowSet
filter Filter // may be nil
limit ReadOption // may be nil
// We do the read, grab all the cells, turn them into "<row>-<col>-<val>",
// and join with a comma.
want string
}{
{
desc: "read all, unfiltered",
rr: RowRange{},
want: "gwashington-jadams-1,jadams-gwashington-1,jadams-tjefferson-1,tjefferson-gwashington-1,tjefferson-jadams-1,tjefferson-wmckinley-1,wmckinley-tjefferson-1",
},
{
desc: "read with InfiniteRange, unfiltered",
rr: InfiniteRange("tjefferson"),
want: "tjefferson-gwashington-1,tjefferson-jadams-1,tjefferson-wmckinley-1,wmckinley-tjefferson-1",
},
{
desc: "read with NewRange, unfiltered",
rr: NewRange("gargamel", "hubbard"),
want: "gwashington-jadams-1",
},
{
desc: "read with PrefixRange, unfiltered",
rr: PrefixRange("jad"),
want: "jadams-gwashington-1,jadams-tjefferson-1",
},
{
desc: "read with SingleRow, unfiltered",
rr: SingleRow("wmckinley"),
want: "wmckinley-tjefferson-1",
},
{
desc: "read all, with ColumnFilter",
rr: RowRange{},
filter: ColumnFilter(".*j.*"), // matches "jadams" and "tjefferson"
want: "gwashington-jadams-1,jadams-tjefferson-1,tjefferson-jadams-1,wmckinley-tjefferson-1",
},
{
desc: "read all, with ColumnFilter, prefix",
rr: RowRange{},
filter: ColumnFilter("j"), // no matches
want: "",
},
{
desc: "read range, with ColumnRangeFilter",
rr: RowRange{},
filter: ColumnRangeFilter("follows", "h", "k"),
want: "gwashington-jadams-1,tjefferson-jadams-1",
},
{
desc: "read range from empty, with ColumnRangeFilter",
rr: RowRange{},
filter: ColumnRangeFilter("follows", "", "u"),
want: "gwashington-jadams-1,jadams-gwashington-1,jadams-tjefferson-1,tjefferson-gwashington-1,tjefferson-jadams-1,wmckinley-tjefferson-1",
},
{
desc: "read range from start to empty, with ColumnRangeFilter",
rr: RowRange{},
filter: ColumnRangeFilter("follows", "h", ""),
want: "gwashington-jadams-1,jadams-tjefferson-1,tjefferson-jadams-1,tjefferson-wmckinley-1,wmckinley-tjefferson-1",
},
{
desc: "read with RowKeyFilter",
rr: RowRange{},
filter: RowKeyFilter(".*wash.*"),
want: "gwashington-jadams-1",
},
{
desc: "read with RowKeyFilter, prefix",
rr: RowRange{},
filter: RowKeyFilter("gwash"),
want: "",
},
{
desc: "read with RowKeyFilter, no matches",
rr: RowRange{},
filter: RowKeyFilter(".*xxx.*"),
want: "",
},
{
desc: "read with FamilyFilter, no matches",
rr: RowRange{},
filter: FamilyFilter(".*xxx.*"),
want: "",
},
{
desc: "read with ColumnFilter + row limit",
rr: RowRange{},
filter: ColumnFilter(".*j.*"), // matches "jadams" and "tjefferson"
limit: LimitRows(2),
want: "gwashington-jadams-1,jadams-tjefferson-1",
},
{
desc: "read all, strip values",
rr: RowRange{},
filter: StripValueFilter(),
want: "gwashington-jadams-,jadams-gwashington-,jadams-tjefferson-,tjefferson-gwashington-,tjefferson-jadams-,tjefferson-wmckinley-,wmckinley-tjefferson-",
},
{
desc: "read with ColumnFilter + row limit + strip values",
rr: RowRange{},
filter: ChainFilters(ColumnFilter(".*j.*"), StripValueFilter()), // matches "jadams" and "tjefferson"
limit: LimitRows(2),
want: "gwashington-jadams-,jadams-tjefferson-",
},
{
desc: "read with condition, strip values on true",
rr: RowRange{},
filter: ConditionFilter(ColumnFilter(".*j.*"), StripValueFilter(), nil),
want: "gwashington-jadams-,jadams-gwashington-,jadams-tjefferson-,tjefferson-gwashington-,tjefferson-jadams-,tjefferson-wmckinley-,wmckinley-tjefferson-",
},
{
desc: "read with condition, strip values on false",
rr: RowRange{},
filter: ConditionFilter(ColumnFilter(".*xxx.*"), nil, StripValueFilter()),
want: "gwashington-jadams-,jadams-gwashington-,jadams-tjefferson-,tjefferson-gwashington-,tjefferson-jadams-,tjefferson-wmckinley-,wmckinley-tjefferson-",
},
{
desc: "read with ValueRangeFilter + row limit",
rr: RowRange{},
filter: ValueRangeFilter([]byte("1"), []byte("5")), // matches our value of "1"
limit: LimitRows(2),
want: "gwashington-jadams-1,jadams-gwashington-1,jadams-tjefferson-1",
},
{
desc: "read with ValueRangeFilter, no match on exclusive end",
rr: RowRange{},
filter: ValueRangeFilter([]byte("0"), []byte("1")), // no match
want: "",
},
{
desc: "read with ValueRangeFilter, no matches",
rr: RowRange{},
filter: ValueRangeFilter([]byte("3"), []byte("5")), // matches nothing
want: "",
},
{
desc: "read with InterleaveFilter, no matches on all filters",
rr: RowRange{},
filter: InterleaveFilters(ColumnFilter(".*x.*"), ColumnFilter(".*z.*")),
want: "",
},
{
desc: "read with InterleaveFilter, no duplicate cells",
rr: RowRange{},
filter: InterleaveFilters(ColumnFilter(".*g.*"), ColumnFilter(".*j.*")),
want: "gwashington-jadams-1,jadams-gwashington-1,jadams-tjefferson-1,tjefferson-gwashington-1,tjefferson-jadams-1,wmckinley-tjefferson-1",
},
{
desc: "read with InterleaveFilter, with duplicate cells",
rr: RowRange{},
filter: InterleaveFilters(ColumnFilter(".*g.*"), ColumnFilter(".*g.*")),
want: "jadams-gwashington-1,jadams-gwashington-1,tjefferson-gwashington-1,tjefferson-gwashington-1",
},
{
desc: "read with a RowRangeList and no filter",
rr: RowRangeList{NewRange("gargamel", "hubbard"), InfiniteRange("wmckinley")},
want: "gwashington-jadams-1,wmckinley-tjefferson-1",
},
{
desc: "chain that excludes rows and matches nothing, in a condition",
rr: RowRange{},
filter: ConditionFilter(ChainFilters(ColumnFilter(".*j.*"), ColumnFilter(".*mckinley.*")), StripValueFilter(), nil),
want: "",
},
{
desc: "chain that ends with an interleave that has no match. covers #804",
rr: RowRange{},
filter: ConditionFilter(ChainFilters(ColumnFilter(".*j.*"), InterleaveFilters(ColumnFilter(".*x.*"), ColumnFilter(".*z.*"))), StripValueFilter(), nil),
want: "",
},
}
for _, tc := range readTests {
var opts []ReadOption
if tc.filter != nil {
opts = append(opts, RowFilter(tc.filter))
}
if tc.limit != nil {
opts = append(opts, tc.limit)
}
var elt []string
err := tbl.ReadRows(ctx, tc.rr, func(r Row) bool {
for _, ris := range r {
for _, ri := range ris {
elt = append(elt, formatReadItem(ri))
}
}
return true
}, opts...)
if err != nil {
t.Errorf("%s: %v", tc.desc, err)
continue
}
if got := strings.Join(elt, ","); got != tc.want {
t.Errorf("%s: wrong reads.\n got %q\nwant %q", tc.desc, got, tc.want)
}
}
// Read a RowList
var elt []string
keys := RowList{"wmckinley", "gwashington", "jadams"}
want := "gwashington-jadams-1,jadams-gwashington-1,jadams-tjefferson-1,wmckinley-tjefferson-1"
err = tbl.ReadRows(ctx, keys, func(r Row) bool {
for _, ris := range r {
for _, ri := range ris {
elt = append(elt, formatReadItem(ri))
}
}
return true
})
if err != nil {
t.Errorf("read RowList: %v", err)
}
if got := strings.Join(elt, ","); got != want {
t.Errorf("bulk read: wrong reads.\n got %q\nwant %q", got, want)
}
checkpoint("tested ReadRows in a few ways")
// Do a scan and stop part way through.
// Verify that the ReadRows callback doesn't keep running.
stopped := false
err = tbl.ReadRows(ctx, InfiniteRange(""), func(r Row) bool {
if r.Key() < "h" {
return true
}
if !stopped {
stopped = true
return false
}
t.Errorf("ReadRows kept scanning to row %q after being told to stop", r.Key())
return false
})
if err != nil {
t.Errorf("Partial ReadRows: %v", err)
}
checkpoint("did partial ReadRows test")
// Delete a row and check it goes away.
mut = NewMutation()
mut.DeleteRow()
if err := tbl.Apply(ctx, "wmckinley", mut); err != nil {
t.Errorf("Apply DeleteRow: %v", err)
}
row, err = tbl.ReadRow(ctx, "wmckinley")
if err != nil {
t.Fatalf("Reading a row after DeleteRow: %v", err)
}
if len(row) != 0 {
t.Fatalf("Read non-zero row after DeleteRow: %v", row)
}
checkpoint("exercised DeleteRow")
// Check ReadModifyWrite.
if err := adminClient.CreateColumnFamily(ctx, table, "counter"); err != nil {
t.Fatalf("Creating column family: %v", err)
}
appendRMW := func(b []byte) *ReadModifyWrite {
rmw := NewReadModifyWrite()
rmw.AppendValue("counter", "likes", b)
return rmw
}
incRMW := func(n int64) *ReadModifyWrite {
rmw := NewReadModifyWrite()
rmw.Increment("counter", "likes", n)
return rmw
}
rmwSeq := []struct {
desc string
rmw *ReadModifyWrite
want []byte
}{
{
desc: "append #1",
rmw: appendRMW([]byte{0, 0, 0}),
want: []byte{0, 0, 0},
},
{
desc: "append #2",
rmw: appendRMW([]byte{0, 0, 0, 0, 17}), // the remaining 40 bits to make a big-endian 17
want: []byte{0, 0, 0, 0, 0, 0, 0, 17},
},
{
desc: "increment",
rmw: incRMW(8),
want: []byte{0, 0, 0, 0, 0, 0, 0, 25},
},
}
for _, step := range rmwSeq {
row, err := tbl.ApplyReadModifyWrite(ctx, "gwashington", step.rmw)
if err != nil {
t.Fatalf("ApplyReadModifyWrite %+v: %v", step.rmw, err)
}
// Make sure the modified cell returned by the RMW operation has a timestamp.
if row["counter"][0].Timestamp == 0 {
t.Errorf("RMW returned cell timestamp: got %v, want > 0", row["counter"][0].Timestamp)
}
clearTimestamps(row)
wantRow := Row{"counter": []ReadItem{{Row: "gwashington", Column: "counter:likes", Value: step.want}}}
if !testutil.Equal(row, wantRow) {
t.Fatalf("After %s,\n got %v\nwant %v", step.desc, row, wantRow)
}
}
// Check for google-cloud-go/issues/723. RMWs that insert new rows should keep row order sorted in the emulator.
_, err = tbl.ApplyReadModifyWrite(ctx, "issue-723-2", appendRMW([]byte{0}))
if err != nil {
t.Fatalf("ApplyReadModifyWrite null string: %v", err)
}
_, err = tbl.ApplyReadModifyWrite(ctx, "issue-723-1", appendRMW([]byte{0}))
if err != nil {
t.Fatalf("ApplyReadModifyWrite null string: %v", err)
}
// Get only the correct row back on read.
r, err := tbl.ReadRow(ctx, "issue-723-1")
if err != nil {
t.Fatalf("Reading row: %v", err)
}
if r.Key() != "issue-723-1" {
t.Errorf("ApplyReadModifyWrite: incorrect read after RMW,\n got %v\nwant %v", r.Key(), "issue-723-1")
}
checkpoint("tested ReadModifyWrite")
// Test arbitrary timestamps more thoroughly.
if err := adminClient.CreateColumnFamily(ctx, table, "ts"); err != nil {
t.Fatalf("Creating column family: %v", err)
}
const numVersions = 4
mut = NewMutation()
for i := 1; i < numVersions; i++ {
// Timestamps are used in thousands because the server
// only permits that granularity.
mut.Set("ts", "col", Timestamp(i*1000), []byte(fmt.Sprintf("val-%d", i)))
mut.Set("ts", "col2", Timestamp(i*1000), []byte(fmt.Sprintf("val-%d", i)))
}
if err := tbl.Apply(ctx, "testrow", mut); err != nil {
t.Fatalf("Mutating row: %v", err)
}
r, err = tbl.ReadRow(ctx, "testrow")
if err != nil {
t.Fatalf("Reading row: %v", err)
}
wantRow = Row{"ts": []ReadItem{
// These should be returned in descending timestamp order.
{Row: "testrow", Column: "ts:col", Timestamp: 3000, Value: []byte("val-3")},
{Row: "testrow", Column: "ts:col", Timestamp: 2000, Value: []byte("val-2")},
{Row: "testrow", Column: "ts:col", Timestamp: 1000, Value: []byte("val-1")},
{Row: "testrow", Column: "ts:col2", Timestamp: 3000, Value: []byte("val-3")},
{Row: "testrow", Column: "ts:col2", Timestamp: 2000, Value: []byte("val-2")},
{Row: "testrow", Column: "ts:col2", Timestamp: 1000, Value: []byte("val-1")},
}}
if !testutil.Equal(r, wantRow) {
t.Errorf("Cell with multiple versions,\n got %v\nwant %v", r, wantRow)
}
// Do the same read, but filter to the latest two versions.
r, err = tbl.ReadRow(ctx, "testrow", RowFilter(LatestNFilter(2)))
if err != nil {
t.Fatalf("Reading row: %v", err)
}
wantRow = Row{"ts": []ReadItem{
{Row: "testrow", Column: "ts:col", Timestamp: 3000, Value: []byte("val-3")},
{Row: "testrow", Column: "ts:col", Timestamp: 2000, Value: []byte("val-2")},
{Row: "testrow", Column: "ts:col2", Timestamp: 3000, Value: []byte("val-3")},
{Row: "testrow", Column: "ts:col2", Timestamp: 2000, Value: []byte("val-2")},
}}
if !testutil.Equal(r, wantRow) {
t.Errorf("Cell with multiple versions and LatestNFilter(2),\n got %v\nwant %v", r, wantRow)
}
// Check cell offset / limit
r, err = tbl.ReadRow(ctx, "testrow", RowFilter(CellsPerRowLimitFilter(3)))
if err != nil {
t.Fatalf("Reading row: %v", err)
}
wantRow = Row{"ts": []ReadItem{
{Row: "testrow", Column: "ts:col", Timestamp: 3000, Value: []byte("val-3")},
{Row: "testrow", Column: "ts:col", Timestamp: 2000, Value: []byte("val-2")},
{Row: "testrow", Column: "ts:col", Timestamp: 1000, Value: []byte("val-1")},
}}
if !testutil.Equal(r, wantRow) {
t.Errorf("Cell with multiple versions and CellsPerRowLimitFilter(3),\n got %v\nwant %v", r, wantRow)
}
r, err = tbl.ReadRow(ctx, "testrow", RowFilter(CellsPerRowOffsetFilter(3)))
if err != nil {
t.Fatalf("Reading row: %v", err)
}
wantRow = Row{"ts": []ReadItem{
{Row: "testrow", Column: "ts:col2", Timestamp: 3000, Value: []byte("val-3")},
{Row: "testrow", Column: "ts:col2", Timestamp: 2000, Value: []byte("val-2")},
{Row: "testrow", Column: "ts:col2", Timestamp: 1000, Value: []byte("val-1")},
}}
if !testutil.Equal(r, wantRow) {
t.Errorf("Cell with multiple versions and CellsPerRowOffsetFilter(3),\n got %v\nwant %v", r, wantRow)
}
// Check timestamp range filtering (with truncation)
r, err = tbl.ReadRow(ctx, "testrow", RowFilter(TimestampRangeFilterMicros(1001, 3000)))
if err != nil {
t.Fatalf("Reading row: %v", err)
}
wantRow = Row{"ts": []ReadItem{
{Row: "testrow", Column: "ts:col", Timestamp: 2000, Value: []byte("val-2")},
{Row: "testrow", Column: "ts:col", Timestamp: 1000, Value: []byte("val-1")},
{Row: "testrow", Column: "ts:col2", Timestamp: 2000, Value: []byte("val-2")},
{Row: "testrow", Column: "ts:col2", Timestamp: 1000, Value: []byte("val-1")},
}}
if !testutil.Equal(r, wantRow) {
t.Errorf("Cell with multiple versions and TimestampRangeFilter(1000, 3000),\n got %v\nwant %v", r, wantRow)
}
r, err = tbl.ReadRow(ctx, "testrow", RowFilter(TimestampRangeFilterMicros(1000, 0)))
if err != nil {
t.Fatalf("Reading row: %v", err)
}
wantRow = Row{"ts": []ReadItem{
{Row: "testrow", Column: "ts:col", Timestamp: 3000, Value: []byte("val-3")},
{Row: "testrow", Column: "ts:col", Timestamp: 2000, Value: []byte("val-2")},
{Row: "testrow", Column: "ts:col", Timestamp: 1000, Value: []byte("val-1")},
{Row: "testrow", Column: "ts:col2", Timestamp: 3000, Value: []byte("val-3")},
{Row: "testrow", Column: "ts:col2", Timestamp: 2000, Value: []byte("val-2")},
{Row: "testrow", Column: "ts:col2", Timestamp: 1000, Value: []byte("val-1")},
}}
if !testutil.Equal(r, wantRow) {
t.Errorf("Cell with multiple versions and TimestampRangeFilter(1000, 0),\n got %v\nwant %v", r, wantRow)
}
// Delete non-existing cells, no such column family in this row
// Should not delete anything
if err := adminClient.CreateColumnFamily(ctx, table, "non-existing"); err != nil {
t.Fatalf("Creating column family: %v", err)
}
mut = NewMutation()
mut.DeleteTimestampRange("non-existing", "col", 2000, 3000) // half-open interval
if err := tbl.Apply(ctx, "testrow", mut); err != nil {
t.Fatalf("Mutating row: %v", err)
}
r, err = tbl.ReadRow(ctx, "testrow", RowFilter(LatestNFilter(3)))
if err != nil {
t.Fatalf("Reading row: %v", err)
}
if !testutil.Equal(r, wantRow) {
t.Errorf("Cell was deleted unexpectly,\n got %v\nwant %v", r, wantRow)
}
// Delete non-existing cells, no such column in this column family
// Should not delete anything
mut = NewMutation()
mut.DeleteTimestampRange("ts", "non-existing", 2000, 3000) // half-open interval
if err := tbl.Apply(ctx, "testrow", mut); err != nil {
t.Fatalf("Mutating row: %v", err)
}
r, err = tbl.ReadRow(ctx, "testrow", RowFilter(LatestNFilter(3)))
if err != nil {
t.Fatalf("Reading row: %v", err)
}
if !testutil.Equal(r, wantRow) {
t.Errorf("Cell was deleted unexpectly,\n got %v\nwant %v", r, wantRow)
}
// Delete the cell with timestamp 2000 and repeat the last read,
// checking that we get ts 3000 and ts 1000.
mut = NewMutation()
mut.DeleteTimestampRange("ts", "col", 2001, 3000) // half-open interval
if err := tbl.Apply(ctx, "testrow", mut); err != nil {
t.Fatalf("Mutating row: %v", err)
}
r, err = tbl.ReadRow(ctx, "testrow", RowFilter(LatestNFilter(2)))
if err != nil {
t.Fatalf("Reading row: %v", err)
}
wantRow = Row{"ts": []ReadItem{
{Row: "testrow", Column: "ts:col", Timestamp: 3000, Value: []byte("val-3")},
{Row: "testrow", Column: "ts:col", Timestamp: 1000, Value: []byte("val-1")},
{Row: "testrow", Column: "ts:col2", Timestamp: 3000, Value: []byte("val-3")},
{Row: "testrow", Column: "ts:col2", Timestamp: 2000, Value: []byte("val-2")},
}}
if !testutil.Equal(r, wantRow) {
t.Errorf("Cell with multiple versions and LatestNFilter(2), after deleting timestamp 2000,\n got %v\nwant %v", r, wantRow)
}
checkpoint("tested multiple versions in a cell")
// Check DeleteCellsInFamily
if err := adminClient.CreateColumnFamily(ctx, table, "status"); err != nil {
t.Fatalf("Creating column family: %v", err)
}
mut = NewMutation()
mut.Set("status", "start", 2000, []byte("2"))
mut.Set("status", "end", 3000, []byte("3"))
mut.Set("ts", "col", 1000, []byte("1"))
if err := tbl.Apply(ctx, "row1", mut); err != nil {
t.Errorf("Mutating row: %v", err)
}
if err := tbl.Apply(ctx, "row2", mut); err != nil {
t.Errorf("Mutating row: %v", err)
}
mut = NewMutation()
mut.DeleteCellsInFamily("status")
if err := tbl.Apply(ctx, "row1", mut); err != nil {
t.Errorf("Delete cf: %v", err)
}
// ColumnFamily removed
r, err = tbl.ReadRow(ctx, "row1")
if err != nil {
t.Fatalf("Reading row: %v", err)
}
wantRow = Row{"ts": []ReadItem{
{Row: "row1", Column: "ts:col", Timestamp: 1000, Value: []byte("1")},
}}
if !testutil.Equal(r, wantRow) {
t.Errorf("column family was not deleted.\n got %v\n want %v", r, wantRow)
}
// ColumnFamily not removed
r, err = tbl.ReadRow(ctx, "row2")
if err != nil {
t.Fatalf("Reading row: %v", err)
}
wantRow = Row{
"ts": []ReadItem{
{Row: "row2", Column: "ts:col", Timestamp: 1000, Value: []byte("1")},
},
"status": []ReadItem{
{Row: "row2", Column: "status:end", Timestamp: 3000, Value: []byte("3")},
{Row: "row2", Column: "status:start", Timestamp: 2000, Value: []byte("2")},
},
}
if !testutil.Equal(r, wantRow) {
t.Errorf("Column family was deleted unexpectly.\n got %v\n want %v", r, wantRow)
}
checkpoint("tested family delete")
// Check DeleteCellsInColumn
mut = NewMutation()
mut.Set("status", "start", 2000, []byte("2"))
mut.Set("status", "middle", 3000, []byte("3"))
mut.Set("status", "end", 1000, []byte("1"))
if err := tbl.Apply(ctx, "row3", mut); err != nil {
t.Errorf("Mutating row: %v", err)
}
mut = NewMutation()
mut.DeleteCellsInColumn("status", "middle")
if err := tbl.Apply(ctx, "row3", mut); err != nil {
t.Errorf("Delete column: %v", err)
}
r, err = tbl.ReadRow(ctx, "row3")
if err != nil {
t.Fatalf("Reading row: %v", err)
}
wantRow = Row{
"status": []ReadItem{
{Row: "row3", Column: "status:end", Timestamp: 1000, Value: []byte("1")},
{Row: "row3", Column: "status:start", Timestamp: 2000, Value: []byte("2")},
},
}
if !testutil.Equal(r, wantRow) {
t.Errorf("Column was not deleted.\n got %v\n want %v", r, wantRow)
}
mut = NewMutation()
mut.DeleteCellsInColumn("status", "start")
if err := tbl.Apply(ctx, "row3", mut); err != nil {
t.Errorf("Delete column: %v", err)
}
r, err = tbl.ReadRow(ctx, "row3")
if err != nil {
t.Fatalf("Reading row: %v", err)
}
wantRow = Row{
"status": []ReadItem{
{Row: "row3", Column: "status:end", Timestamp: 1000, Value: []byte("1")},
},
}
if !testutil.Equal(r, wantRow) {
t.Errorf("Column was not deleted.\n got %v\n want %v", r, wantRow)
}
mut = NewMutation()
mut.DeleteCellsInColumn("status", "end")
if err := tbl.Apply(ctx, "row3", mut); err != nil {
t.Errorf("Delete column: %v", err)
}
r, err = tbl.ReadRow(ctx, "row3")
if err != nil {
t.Fatalf("Reading row: %v", err)
}
if len(r) != 0 {
t.Errorf("Delete column: got %v, want empty row", r)
}
// Add same cell after delete
mut = NewMutation()
mut.Set("status", "end", 1000, []byte("1"))
if err := tbl.Apply(ctx, "row3", mut); err != nil {
t.Errorf("Mutating row: %v", err)
}
r, err = tbl.ReadRow(ctx, "row3")
if err != nil {
t.Fatalf("Reading row: %v", err)
}
if !testutil.Equal(r, wantRow) {
t.Errorf("Column was not deleted correctly.\n got %v\n want %v", r, wantRow)
}
checkpoint("tested column delete")
// Do highly concurrent reads/writes.
// TODO(dsymonds): Raise this to 1000 when https://github.com/grpc/grpc-go/issues/205 is resolved.
const maxConcurrency = 100
var wg sync.WaitGroup
for i := 0; i < maxConcurrency; i++ {
wg.Add(1)
go func() {
defer wg.Done()
switch r := rand.Intn(100); { // r ∈ [0,100)
case 0 <= r && r < 30:
// Do a read.
_, err := tbl.ReadRow(ctx, "testrow", RowFilter(LatestNFilter(1)))
if err != nil {
t.Errorf("Concurrent read: %v", err)
}
case 30 <= r && r < 100:
// Do a write.
mut := NewMutation()
mut.Set("ts", "col", 1000, []byte("data"))
if err := tbl.Apply(ctx, "testrow", mut); err != nil {
t.Errorf("Concurrent write: %v", err)
}
}
}()
}
wg.Wait()
checkpoint("tested high concurrency")
// Large reads, writes and scans.
bigBytes := make([]byte, 5<<20) // 5 MB is larger than current default gRPC max of 4 MB, but less than the max we set.
nonsense := []byte("lorem ipsum dolor sit amet, ")
fill(bigBytes, nonsense)
mut = NewMutation()
mut.Set("ts", "col", 1000, bigBytes)
if err := tbl.Apply(ctx, "bigrow", mut); err != nil {
t.Errorf("Big write: %v", err)
}
r, err = tbl.ReadRow(ctx, "bigrow")
if err != nil {
t.Errorf("Big read: %v", err)
}
wantRow = Row{"ts": []ReadItem{
{Row: "bigrow", Column: "ts:col", Timestamp: 1000, Value: bigBytes},
}}
if !testutil.Equal(r, wantRow) {
t.Errorf("Big read returned incorrect bytes: %v", r)
}
// Now write 1000 rows, each with 82 KB values, then scan them all.
medBytes := make([]byte, 82<<10)
fill(medBytes, nonsense)
sem := make(chan int, 50) // do up to 50 mutations at a time.
for i := 0; i < 1000; i++ {
mut := NewMutation()
mut.Set("ts", "big-scan", 1000, medBytes)
row := fmt.Sprintf("row-%d", i)
wg.Add(1)
go func() {
defer wg.Done()
defer func() { <-sem }()
sem <- 1
if err := tbl.Apply(ctx, row, mut); err != nil {
t.Errorf("Preparing large scan: %v", err)
}
}()
}
wg.Wait()
n := 0
err = tbl.ReadRows(ctx, PrefixRange("row-"), func(r Row) bool {
for _, ris := range r {
for _, ri := range ris {
n += len(ri.Value)
}
}
return true
}, RowFilter(ColumnFilter("big-scan")))
if err != nil {
t.Errorf("Doing large scan: %v", err)
}
if want := 1000 * len(medBytes); n != want {
t.Errorf("Large scan returned %d bytes, want %d", n, want)
}
// Scan a subset of the 1000 rows that we just created, using a LimitRows ReadOption.
rc := 0
wantRc := 3
err = tbl.ReadRows(ctx, PrefixRange("row-"), func(r Row) bool {
rc++
return true
}, LimitRows(int64(wantRc)))
if err != nil {
t.Fatal(err)
}
if rc != wantRc {
t.Errorf("Scan with row limit returned %d rows, want %d", rc, wantRc)
}
checkpoint("tested big read/write/scan")
// Test bulk mutations
if err := adminClient.CreateColumnFamily(ctx, table, "bulk"); err != nil {
t.Fatalf("Creating column family: %v", err)
}
bulkData := map[string][]string{
"red sox": {"2004", "2007", "2013"},
"patriots": {"2001", "2003", "2004", "2014"},
"celtics": {"1981", "1984", "1986", "2008"},
}
var rowKeys []string
var muts []*Mutation
for row, ss := range bulkData {
mut := NewMutation()
for _, name := range ss {
mut.Set("bulk", name, 1000, []byte("1"))
}
rowKeys = append(rowKeys, row)
muts = append(muts, mut)
}
status, err := tbl.ApplyBulk(ctx, rowKeys, muts)
if err != nil {
t.Fatalf("Bulk mutating rows %q: %v", rowKeys, err)
}
if status != nil {
t.Errorf("non-nil errors: %v", err)
}
checkpoint("inserted bulk data")
// Read each row back
for rowKey, ss := range bulkData {
row, err := tbl.ReadRow(ctx, rowKey)
if err != nil {
t.Fatalf("Reading a bulk row: %v", err)
}
var wantItems []ReadItem
for _, val := range ss {
wantItems = append(wantItems, ReadItem{Row: rowKey, Column: "bulk:" + val, Timestamp: 1000, Value: []byte("1")})
}
wantRow := Row{"bulk": wantItems}
if !testutil.Equal(row, wantRow) {
t.Errorf("Read row mismatch.\n got %#v\nwant %#v", row, wantRow)
}
}
checkpoint("tested reading from bulk insert")
// Test bulk write errors.
// Note: Setting timestamps as ServerTime makes sure the mutations are not retried on error.
badMut := NewMutation()
badMut.Set("badfamily", "col", ServerTime, nil)
badMut2 := NewMutation()
badMut2.Set("badfamily2", "goodcol", ServerTime, []byte("1"))
status, err = tbl.ApplyBulk(ctx, []string{"badrow", "badrow2"}, []*Mutation{badMut, badMut2})
if err != nil {
t.Fatalf("Bulk mutating rows %q: %v", rowKeys, err)
}
if status == nil {
t.Errorf("No errors for bad bulk mutation")
} else if status[0] == nil || status[1] == nil {
t.Errorf("No error for bad bulk mutation")
}
}
type requestCountingInterceptor struct {
grpc.ClientStream
requestCallback func()
}
func (i *requestCountingInterceptor) SendMsg(m interface{}) error {
i.requestCallback()
return i.ClientStream.SendMsg(m)
}
func (i *requestCountingInterceptor) RecvMsg(m interface{}) error {
return i.ClientStream.RecvMsg(m)
}
func requestCallback(callback func()) func(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, streamer grpc.Streamer, opts ...grpc.CallOption) (grpc.ClientStream, error) {
return func(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, streamer grpc.Streamer, opts ...grpc.CallOption) (grpc.ClientStream, error) {
clientStream, err := streamer(ctx, desc, cc, method, opts...)
return &requestCountingInterceptor{
ClientStream: clientStream,
requestCallback: callback,
}, err
}
}
// TestReadRowsInvalidRowSet verifies that the client doesn't send ReadRows() requests with invalid RowSets.
func TestReadRowsInvalidRowSet(t *testing.T) {
testEnv, err := NewEmulatedEnv(IntegrationTestConfig{})
if err != nil {
t.Fatalf("NewEmulatedEnv failed: %v", err)
}
var requestCount int
incrementRequestCount := func() { requestCount++ }
conn, err := grpc.Dial(testEnv.server.Addr, grpc.WithInsecure(), grpc.WithBlock(),
grpc.WithDefaultCallOptions(grpc.MaxCallSendMsgSize(100<<20), grpc.MaxCallRecvMsgSize(100<<20)),
grpc.WithStreamInterceptor(requestCallback(incrementRequestCount)),
)
if err != nil {
t.Fatalf("grpc.Dial failed: %v", err)
}
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
defer cancel()
adminClient, err := NewAdminClient(ctx, testEnv.config.Project, testEnv.config.Instance, option.WithGRPCConn(conn))
if err != nil {
t.Fatalf("NewClient failed: %v", err)
}
defer adminClient.Close()
if err := adminClient.CreateTable(ctx, testEnv.config.Table); err != nil {
t.Fatalf("CreateTable(%v) failed: %v", testEnv.config.Table, err)
}
client, err := NewClient(ctx, testEnv.config.Project, testEnv.config.Instance, option.WithGRPCConn(conn))
if err != nil {
t.Fatalf("NewClient failed: %v", err)
}
defer client.Close()
table := client.Open(testEnv.config.Table)
tests := []struct {
rr RowSet
valid bool
}{
{
rr: RowRange{},
valid: true,
},
{
rr: RowRange{start: "b"},
valid: true,
},
{
rr: RowRange{start: "b", limit: "c"},
valid: true,
},
{
rr: RowRange{start: "b", limit: "a"},
valid: false,
},
{
rr: RowList{"a"},
valid: true,
},
{
rr: RowList{},
valid: false,
},
}
for _, test := range tests {
requestCount = 0
err = table.ReadRows(ctx, test.rr, func(r Row) bool { return true })
if err != nil {
t.Fatalf("ReadRows(%v) failed: %v", test.rr, err)
}
requestValid := requestCount != 0
if requestValid != test.valid {
t.Errorf("%s: got %v, want %v", test.rr, requestValid, test.valid)
}
}
}
func formatReadItem(ri ReadItem) string {
// Use the column qualifier only to make the test data briefer.
col := ri.Column[strings.Index(ri.Column, ":")+1:]
return fmt.Sprintf("%s-%s-%s", ri.Row, col, ri.Value)
}
func fill(b, sub []byte) {
for len(b) > len(sub) {
n := copy(b, sub)
b = b[n:]
}
}
func clearTimestamps(r Row) {
for _, ris := range r {
for i := range ris {
ris[i].Timestamp = 0
}
}
}
func TestSampleRowKeys(t *testing.T) {
start := time.Now()
lastCheckpoint := start
checkpoint := func(s string) {
n := time.Now()
t.Logf("[%s] %v since start, %v since last checkpoint", s, n.Sub(start), n.Sub(lastCheckpoint))
lastCheckpoint = n
}
ctx := context.Background()
client, adminClient, table, err := doSetup(ctx)
if err != nil {
t.Fatalf("%v", err)
}
defer client.Close()
defer adminClient.Close()
tbl := client.Open(table)
// Delete the table at the end of the test.
// Do this even before creating the table so that if this is running
// against production and CreateTable fails there's a chance of cleaning it up.
defer adminClient.DeleteTable(ctx, table)
// Insert some data.
initialData := map[string][]string{
"wmckinley11": {"tjefferson11"},
"gwashington77": {"jadams77"},
"tjefferson0": {"gwashington0", "jadams0"},
}
for row, ss := range initialData {
mut := NewMutation()
for _, name := range ss {
mut.Set("follows", name, 1000, []byte("1"))
}
if err := tbl.Apply(ctx, row, mut); err != nil {
t.Errorf("Mutating row %q: %v", row, err)
}
}
checkpoint("inserted initial data")
sampleKeys, err := tbl.SampleRowKeys(context.Background())
if err != nil {
t.Errorf("%s: %v", "SampleRowKeys:", err)
}
if len(sampleKeys) == 0 {
t.Error("SampleRowKeys length 0")
}
checkpoint("tested SampleRowKeys.")
}
func doSetup(ctx context.Context) (*Client, *AdminClient, string, error) {
start := time.Now()
lastCheckpoint := start
checkpoint := func(s string) {
n := time.Now()
fmt.Printf("[%s] %v since start, %v since last checkpoint", s, n.Sub(start), n.Sub(lastCheckpoint))
lastCheckpoint = n
}
testEnv, err := NewIntegrationEnv()
if err != nil {
return nil, nil, "", fmt.Errorf("IntegrationEnv: %v", err)
}
var timeout time.Duration
if testEnv.Config().UseProd {
timeout = 10 * time.Minute
fmt.Printf("Running test against production")
} else {
timeout = 1 * time.Minute
fmt.Printf("bttest.Server running on %s", testEnv.Config().AdminEndpoint)
}
ctx, cancel := context.WithTimeout(ctx, timeout)
defer cancel()
client, err := testEnv.NewClient()
if err != nil {
return nil, nil, "", fmt.Errorf("Client: %v", err)
}
checkpoint("dialed Client")
adminClient, err := testEnv.NewAdminClient()
if err != nil {
return nil, nil, "", fmt.Errorf("AdminClient: %v", err)
}
checkpoint("dialed AdminClient")
table := testEnv.Config().Table
if err := adminClient.CreateTable(ctx, table); err != nil {
return nil, nil, "", fmt.Errorf("Creating table: %v", err)
}
checkpoint("created table")
if err := adminClient.CreateColumnFamily(ctx, table, "follows"); err != nil {
return nil, nil, "", fmt.Errorf("Creating column family: %v", err)
}
checkpoint(`created "follows" column family`)
return client, adminClient, table, nil
}
/*
Copyright 2016 Google LLC
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
package bttest_test
import (
"context"
"fmt"
"log"
"cloud.google.com/go/bigtable"
"cloud.google.com/go/bigtable/bttest"
"google.golang.org/api/option"
"google.golang.org/grpc"
)
func ExampleNewServer() {
srv, err := bttest.NewServer("localhost:0")
if err != nil {
log.Fatalln(err)
}
ctx := context.Background()
conn, err := grpc.Dial(srv.Addr, grpc.WithInsecure())
if err != nil {
log.Fatalln(err)
}
proj, instance := "proj", "instance"
adminClient, err := bigtable.NewAdminClient(ctx, proj, instance, option.WithGRPCConn(conn))
if err != nil {
log.Fatalln(err)
}
if err = adminClient.CreateTable(ctx, "example"); err != nil {
log.Fatalln(err)
}
if err = adminClient.CreateColumnFamily(ctx, "example", "links"); err != nil {
log.Fatalln(err)
}
client, err := bigtable.NewClient(ctx, proj, instance, option.WithGRPCConn(conn))
if err != nil {
log.Fatalln(err)
}
tbl := client.Open("example")
mut := bigtable.NewMutation()
mut.Set("links", "golang.org", bigtable.Now(), []byte("Gophers!"))
if err = tbl.Apply(ctx, "com.google.cloud", mut); err != nil {
log.Fatalln(err)
}
if row, err := tbl.ReadRow(ctx, "com.google.cloud"); err != nil {
log.Fatalln(err)
} else {
for _, column := range row["links"] {
fmt.Println(column.Column)
fmt.Println(string(column.Value))
}
}
// Output:
// links:golang.org
// Gophers!
}
/*
Copyright 2015 Google LLC
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
/*
Package bttest contains test helpers for working with the bigtable package.
To use a Server, create it, and then connect to it with no security:
(The project/instance values are ignored.)
srv, err := bttest.NewServer("localhost:0")
...
conn, err := grpc.Dial(srv.Addr, grpc.WithInsecure())
...
client, err := bigtable.NewClient(ctx, proj, instance,
option.WithGRPCConn(conn))
...
*/
package bttest // import "cloud.google.com/go/bigtable/bttest"
import (
"bytes"
"context"
"encoding/binary"
"fmt"
"log"
"math/rand"
"net"
"regexp"
"sort"
"strings"
"sync"
"time"
emptypb "github.com/golang/protobuf/ptypes/empty"
"github.com/golang/protobuf/ptypes/wrappers"
"github.com/google/btree"
btapb "google.golang.org/genproto/googleapis/bigtable/admin/v2"
btpb "google.golang.org/genproto/googleapis/bigtable/v2"
"google.golang.org/genproto/googleapis/longrunning"
statpb "google.golang.org/genproto/googleapis/rpc/status"
"google.golang.org/grpc"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/status"
)
const (
// MilliSeconds field of the minimum valid Timestamp.
minValidMilliSeconds = 0
// MilliSeconds field of the max valid Timestamp.
maxValidMilliSeconds = int64(time.Millisecond) * 253402300800
)
var (
validLabelTransformer = regexp.MustCompile(`[a-z0-9\-]{1,15}`)
)
// Server is an in-memory Cloud Bigtable fake.
// It is unauthenticated, and only a rough approximation.
type Server struct {
Addr string
l net.Listener
srv *grpc.Server
s *server
}
// server is the real implementation of the fake.
// It is a separate and unexported type so the API won't be cluttered with
// methods that are only relevant to the fake's implementation.
type server struct {
mu sync.Mutex
tables map[string]*table // keyed by fully qualified name
gcc chan int // set when gcloop starts, closed when server shuts down
// Any unimplemented methods will cause a panic.
btapb.BigtableTableAdminServer
btpb.BigtableServer
}
// NewServer creates a new Server.
// The Server will be listening for gRPC connections, without TLS,
// on the provided address. The resolved address is named by the Addr field.
func NewServer(laddr string, opt ...grpc.ServerOption) (*Server, error) {
l, err := net.Listen("tcp", laddr)
if err != nil {
return nil, err
}
s := &Server{
Addr: l.Addr().String(),
l: l,
srv: grpc.NewServer(opt...),
s: &server{
tables: make(map[string]*table),
},
}
btapb.RegisterBigtableTableAdminServer(s.srv, s.s)
btpb.RegisterBigtableServer(s.srv, s.s)
go s.srv.Serve(s.l)
return s, nil
}
// Close shuts down the server.
func (s *Server) Close() {
s.s.mu.Lock()
if s.s.gcc != nil {
close(s.s.gcc)
}
s.s.mu.Unlock()
s.srv.Stop()
s.l.Close()
}
func (s *server) CreateTable(ctx context.Context, req *btapb.CreateTableRequest) (*btapb.Table, error) {
tbl := req.Parent + "/tables/" + req.TableId
s.mu.Lock()
if _, ok := s.tables[tbl]; ok {
s.mu.Unlock()
return nil, status.Errorf(codes.AlreadyExists, "table %q already exists", tbl)
}
s.tables[tbl] = newTable(req)
s.mu.Unlock()
return &btapb.Table{Name: tbl}, nil
}
func (s *server) CreateTableFromSnapshot(context.Context, *btapb.CreateTableFromSnapshotRequest) (*longrunning.Operation, error) {
return nil, status.Errorf(codes.Unimplemented, "the emulator does not currently support snapshots")
}
func (s *server) ListTables(ctx context.Context, req *btapb.ListTablesRequest) (*btapb.ListTablesResponse, error) {
res := &btapb.ListTablesResponse{}
prefix := req.Parent + "/tables/"
s.mu.Lock()
for tbl := range s.tables {
if strings.HasPrefix(tbl, prefix) {
res.Tables = append(res.Tables, &btapb.Table{Name: tbl})
}
}
s.mu.Unlock()
return res, nil
}
func (s *server) GetTable(ctx context.Context, req *btapb.GetTableRequest) (*btapb.Table, error) {
tbl := req.Name
s.mu.Lock()
tblIns, ok := s.tables[tbl]
s.mu.Unlock()
if !ok {
return nil, status.Errorf(codes.NotFound, "table %q not found", tbl)
}
return &btapb.Table{
Name: tbl,
ColumnFamilies: toColumnFamilies(tblIns.columnFamilies()),
}, nil
}
func (s *server) DeleteTable(ctx context.Context, req *btapb.DeleteTableRequest) (*emptypb.Empty, error) {
s.mu.Lock()
defer s.mu.Unlock()
if _, ok := s.tables[req.Name]; !ok {
return nil, status.Errorf(codes.NotFound, "table %q not found", req.Name)
}
delete(s.tables, req.Name)
return &emptypb.Empty{}, nil
}
func (s *server) ModifyColumnFamilies(ctx context.Context, req *btapb.ModifyColumnFamiliesRequest) (*btapb.Table, error) {
s.mu.Lock()
tbl, ok := s.tables[req.Name]
s.mu.Unlock()
if !ok {
return nil, status.Errorf(codes.NotFound, "table %q not found", req.Name)
}
tbl.mu.Lock()
defer tbl.mu.Unlock()
for _, mod := range req.Modifications {
if create := mod.GetCreate(); create != nil {
if _, ok := tbl.families[mod.Id]; ok {
return nil, status.Errorf(codes.AlreadyExists, "family %q already exists", mod.Id)
}
newcf := &columnFamily{
name: req.Name + "/columnFamilies/" + mod.Id,
order: tbl.counter,
gcRule: create.GcRule,
}
tbl.counter++
tbl.families[mod.Id] = newcf
} else if mod.GetDrop() {
if _, ok := tbl.families[mod.Id]; !ok {
return nil, fmt.Errorf("can't delete unknown family %q", mod.Id)
}
delete(tbl.families, mod.Id)
} else if modify := mod.GetUpdate(); modify != nil {
if _, ok := tbl.families[mod.Id]; !ok {
return nil, fmt.Errorf("no such family %q", mod.Id)
}
newcf := &columnFamily{
name: req.Name + "/columnFamilies/" + mod.Id,
gcRule: modify.GcRule,
}
// assume that we ALWAYS want to replace by the new setting
// we may need partial update through
tbl.families[mod.Id] = newcf
}
}
s.needGC()
return &btapb.Table{
Name: req.Name,
ColumnFamilies: toColumnFamilies(tbl.families),
Granularity: btapb.Table_TimestampGranularity(btapb.Table_MILLIS),
}, nil
}
func (s *server) DropRowRange(ctx context.Context, req *btapb.DropRowRangeRequest) (*emptypb.Empty, error) {
s.mu.Lock()
defer s.mu.Unlock()
tbl, ok := s.tables[req.Name]
if !ok {
return nil, status.Errorf(codes.NotFound, "table %q not found", req.Name)
}
if req.GetDeleteAllDataFromTable() {
tbl.rows = btree.New(btreeDegree)
} else {
// Delete rows by prefix.
prefixBytes := req.GetRowKeyPrefix()
if prefixBytes == nil {
return nil, fmt.Errorf("missing row key prefix")
}
prefix := string(prefixBytes)
// The BTree does not specify what happens if rows are deleted during
// iteration, and it provides no "delete range" method.
// So we collect the rows first, then delete them one by one.
var rowsToDelete []*row
tbl.rows.AscendGreaterOrEqual(btreeKey(prefix), func(i btree.Item) bool {
r := i.(*row)
if strings.HasPrefix(r.key, prefix) {
rowsToDelete = append(rowsToDelete, r)
return true
}
return false // stop iteration
})
for _, r := range rowsToDelete {
tbl.rows.Delete(r)
}
}
return &emptypb.Empty{}, nil
}
func (s *server) GenerateConsistencyToken(ctx context.Context, req *btapb.GenerateConsistencyTokenRequest) (*btapb.GenerateConsistencyTokenResponse, error) {
// Check that the table exists.
_, ok := s.tables[req.Name]
if !ok {
return nil, status.Errorf(codes.NotFound, "table %q not found", req.Name)
}
return &btapb.GenerateConsistencyTokenResponse{
ConsistencyToken: "TokenFor-" + req.Name,
}, nil
}
func (s *server) CheckConsistency(ctx context.Context, req *btapb.CheckConsistencyRequest) (*btapb.CheckConsistencyResponse, error) {
// Check that the table exists.
_, ok := s.tables[req.Name]
if !ok {
return nil, status.Errorf(codes.NotFound, "table %q not found", req.Name)
}
// Check this is the right token.
if req.ConsistencyToken != "TokenFor-"+req.Name {
return nil, status.Errorf(codes.InvalidArgument, "token %q not valid", req.ConsistencyToken)
}
// Single cluster instances are always consistent.
return &btapb.CheckConsistencyResponse{
Consistent: true,
}, nil
}
func (s *server) SnapshotTable(context.Context, *btapb.SnapshotTableRequest) (*longrunning.Operation, error) {
return nil, status.Errorf(codes.Unimplemented, "the emulator does not currently support snapshots")
}
func (s *server) GetSnapshot(context.Context, *btapb.GetSnapshotRequest) (*btapb.Snapshot, error) {
return nil, status.Errorf(codes.Unimplemented, "the emulator does not currently support snapshots")
}
func (s *server) ListSnapshots(context.Context, *btapb.ListSnapshotsRequest) (*btapb.ListSnapshotsResponse, error) {
return nil, status.Errorf(codes.Unimplemented, "the emulator does not currently support snapshots")
}
func (s *server) DeleteSnapshot(context.Context, *btapb.DeleteSnapshotRequest) (*emptypb.Empty, error) {
return nil, status.Errorf(codes.Unimplemented, "the emulator does not currently support snapshots")
}
func (s *server) ReadRows(req *btpb.ReadRowsRequest, stream btpb.Bigtable_ReadRowsServer) error {
s.mu.Lock()
tbl, ok := s.tables[req.TableName]
s.mu.Unlock()
if !ok {
return status.Errorf(codes.NotFound, "table %q not found", req.TableName)
}
// Rows to read can be specified by a set of row keys and/or a set of row ranges.
// Output is a stream of sorted, de-duped rows.
tbl.mu.RLock()
rowSet := make(map[string]*row)
addRow := func(i btree.Item) bool {
r := i.(*row)
rowSet[r.key] = r
return true
}
if req.Rows != nil &&
len(req.Rows.RowKeys)+len(req.Rows.RowRanges) > 0 {
// Add the explicitly given keys
for _, key := range req.Rows.RowKeys {
k := string(key)
if i := tbl.rows.Get(btreeKey(k)); i != nil {
addRow(i)
}
}
// Add keys from row ranges
for _, rr := range req.Rows.RowRanges {
var start, end string
switch sk := rr.StartKey.(type) {
case *btpb.RowRange_StartKeyClosed:
start = string(sk.StartKeyClosed)
case *btpb.RowRange_StartKeyOpen:
start = string(sk.StartKeyOpen) + "\x00"
}
switch ek := rr.EndKey.(type) {
case *btpb.RowRange_EndKeyClosed:
end = string(ek.EndKeyClosed) + "\x00"
case *btpb.RowRange_EndKeyOpen:
end = string(ek.EndKeyOpen)
}
switch {
case start == "" && end == "":
tbl.rows.Ascend(addRow) // all rows
case start == "":
tbl.rows.AscendLessThan(btreeKey(end), addRow)
case end == "":
tbl.rows.AscendGreaterOrEqual(btreeKey(start), addRow)
default:
tbl.rows.AscendRange(btreeKey(start), btreeKey(end), addRow)
}
}
} else {
// Read all rows
tbl.rows.Ascend(addRow)
}
tbl.mu.RUnlock()
rows := make([]*row, 0, len(rowSet))
for _, r := range rowSet {
rows = append(rows, r)
}
sort.Sort(byRowKey(rows))
limit := int(req.RowsLimit)
count := 0
for _, r := range rows {
if limit > 0 && count >= limit {
return nil
}
streamed, err := streamRow(stream, r, req.Filter)
if err != nil {
return err
}
if streamed {
count++
}
}
return nil
}
// streamRow filters the given row and sends it via the given stream.
// Returns true if at least one cell matched the filter and was streamed, false otherwise.
func streamRow(stream btpb.Bigtable_ReadRowsServer, r *row, f *btpb.RowFilter) (bool, error) {
r.mu.Lock()
nr := r.copy()
r.mu.Unlock()
r = nr
match, err := filterRow(f, r)
if err != nil {
return false, err
}
if !match {
return false, nil
}
rrr := &btpb.ReadRowsResponse{}
families := r.sortedFamilies()
for _, fam := range families {
for _, colName := range fam.colNames {
cells := fam.cells[colName]
if len(cells) == 0 {
continue
}
for _, cell := range cells {
rrr.Chunks = append(rrr.Chunks, &btpb.ReadRowsResponse_CellChunk{
RowKey: []byte(r.key),
FamilyName: &wrappers.StringValue{Value: fam.name},
Qualifier: &wrappers.BytesValue{Value: []byte(colName)},
TimestampMicros: cell.ts,
Value: cell.value,
Labels: cell.labels,
})
}
}
}
// We can't have a cell with just COMMIT set, which would imply a new empty cell.
// So modify the last cell to have the COMMIT flag set.
if len(rrr.Chunks) > 0 {
rrr.Chunks[len(rrr.Chunks)-1].RowStatus = &btpb.ReadRowsResponse_CellChunk_CommitRow{CommitRow: true}
}
return true, stream.Send(rrr)
}
// filterRow modifies a row with the given filter. Returns true if at least one cell from the row matches,
// false otherwise. If a filter is invalid, filterRow returns false and an error.
func filterRow(f *btpb.RowFilter, r *row) (bool, error) {
if f == nil {
return true, nil
}
// Handle filters that apply beyond just including/excluding cells.
switch f := f.Filter.(type) {
case *btpb.RowFilter_BlockAllFilter:
return !f.BlockAllFilter, nil
case *btpb.RowFilter_PassAllFilter:
return f.PassAllFilter, nil
case *btpb.RowFilter_Chain_:
for _, sub := range f.Chain.Filters {
match, err := filterRow(sub, r)
if err != nil {
return false, err
}
if !match {
return false, nil
}
}
return true, nil
case *btpb.RowFilter_Interleave_:
srs := make([]*row, 0, len(f.Interleave.Filters))
for _, sub := range f.Interleave.Filters {
sr := r.copy()
filterRow(sub, sr)
srs = append(srs, sr)
}
// merge
// TODO(dsymonds): is this correct?
r.families = make(map[string]*family)
for _, sr := range srs {
for _, fam := range sr.families {
f := r.getOrCreateFamily(fam.name, fam.order)
for colName, cs := range fam.cells {
f.cells[colName] = append(f.cellsByColumn(colName), cs...)
}
}
}
var count int
for _, fam := range r.families {
for _, cs := range fam.cells {
sort.Sort(byDescTS(cs))
count += len(cs)
}
}
return count > 0, nil
case *btpb.RowFilter_CellsPerColumnLimitFilter:
lim := int(f.CellsPerColumnLimitFilter)
for _, fam := range r.families {
for col, cs := range fam.cells {
if len(cs) > lim {
fam.cells[col] = cs[:lim]
}
}
}
return true, nil
case *btpb.RowFilter_Condition_:
match, err := filterRow(f.Condition.PredicateFilter, r.copy())
if err != nil {
return false, err
}
if match {
if f.Condition.TrueFilter == nil {
return false, nil
}
return filterRow(f.Condition.TrueFilter, r)
}
if f.Condition.FalseFilter == nil {
return false, nil
}
return filterRow(f.Condition.FalseFilter, r)
case *btpb.RowFilter_RowKeyRegexFilter:
rx, err := newRegexp(f.RowKeyRegexFilter)
if err != nil {
return false, status.Errorf(codes.InvalidArgument, "Error in field 'rowkey_regex_filter' : %v", err)
}
if !rx.MatchString(r.key) {
return false, nil
}
case *btpb.RowFilter_CellsPerRowLimitFilter:
// Grab the first n cells in the row.
lim := int(f.CellsPerRowLimitFilter)
for _, fam := range r.families {
for _, col := range fam.colNames {
cs := fam.cells[col]
if len(cs) > lim {
fam.cells[col] = cs[:lim]
lim = 0
} else {
lim -= len(cs)
}
}
}
return true, nil
case *btpb.RowFilter_CellsPerRowOffsetFilter:
// Skip the first n cells in the row.
offset := int(f.CellsPerRowOffsetFilter)
for _, fam := range r.families {
for _, col := range fam.colNames {
cs := fam.cells[col]
if len(cs) > offset {
fam.cells[col] = cs[offset:]
offset = 0
return true, nil
}
fam.cells[col] = cs[:0]
offset -= len(cs)
}
}
return true, nil
case *btpb.RowFilter_RowSampleFilter:
// The row sample filter "matches all cells from a row with probability
// p, and matches no cells from the row with probability 1-p."
// See https://github.com/googleapis/googleapis/blob/master/google/bigtable/v2/data.proto
if f.RowSampleFilter <= 0.0 || f.RowSampleFilter >= 1.0 {
return false, status.Error(codes.InvalidArgument, "row_sample_filter argument must be between 0.0 and 1.0")
}
return randFloat() < f.RowSampleFilter, nil
}
// Any other case, operate on a per-cell basis.
cellCount := 0
for _, fam := range r.families {
for colName, cs := range fam.cells {
filtered, err := filterCells(f, fam.name, colName, cs)
if err != nil {
return false, err
}
fam.cells[colName] = filtered
cellCount += len(fam.cells[colName])
}
}
return cellCount > 0, nil
}
var randFloat = rand.Float64
func filterCells(f *btpb.RowFilter, fam, col string, cs []cell) ([]cell, error) {
var ret []cell
for _, cell := range cs {
include, err := includeCell(f, fam, col, cell)
if err != nil {
return nil, err
}
if include {
cell, err = modifyCell(f, cell)
if err != nil {
return nil, err
}
ret = append(ret, cell)
}
}
return ret, nil
}
func modifyCell(f *btpb.RowFilter, c cell) (cell, error) {
if f == nil {
return c, nil
}
// Consider filters that may modify the cell contents
switch filter := f.Filter.(type) {
case *btpb.RowFilter_StripValueTransformer:
return cell{ts: c.ts}, nil
case *btpb.RowFilter_ApplyLabelTransformer:
if !validLabelTransformer.MatchString(filter.ApplyLabelTransformer) {
return cell{}, status.Errorf(
codes.InvalidArgument,
`apply_label_transformer must match RE2([a-z0-9\-]+), but found %v`,
filter.ApplyLabelTransformer,
)
}
return cell{ts: c.ts, value: c.value, labels: []string{filter.ApplyLabelTransformer}}, nil
default:
return c, nil
}
}
func includeCell(f *btpb.RowFilter, fam, col string, cell cell) (bool, error) {
if f == nil {
return true, nil
}
// TODO(dsymonds): Implement many more filters.
switch f := f.Filter.(type) {
case *btpb.RowFilter_CellsPerColumnLimitFilter:
// Don't log, row-level filter
return true, nil
case *btpb.RowFilter_RowKeyRegexFilter:
// Don't log, row-level filter
return true, nil
case *btpb.RowFilter_StripValueTransformer:
// Don't log, cell-modifying filter
return true, nil
case *btpb.RowFilter_ApplyLabelTransformer:
// Don't log, cell-modifying filter
return true, nil
default:
log.Printf("WARNING: don't know how to handle filter of type %T (ignoring it)", f)
return true, nil
case *btpb.RowFilter_FamilyNameRegexFilter:
rx, err := newRegexp([]byte(f.FamilyNameRegexFilter))
if err != nil {
return false, status.Errorf(codes.InvalidArgument, "Error in field 'family_name_regex_filter' : %v", err)
}
return rx.MatchString(fam), nil
case *btpb.RowFilter_ColumnQualifierRegexFilter:
rx, err := newRegexp(f.ColumnQualifierRegexFilter)
if err != nil {
return false, status.Errorf(codes.InvalidArgument, "Error in field 'column_qualifier_regex_filter' : %v", err)
}
return rx.MatchString(toUTF8([]byte(col))), nil
case *btpb.RowFilter_ValueRegexFilter:
rx, err := newRegexp(f.ValueRegexFilter)
if err != nil {
return false, status.Errorf(codes.InvalidArgument, "Error in field 'value_regex_filter' : %v", err)
}
return rx.Match(cell.value), nil
case *btpb.RowFilter_ColumnRangeFilter:
if fam != f.ColumnRangeFilter.FamilyName {
return false, nil
}
// Start qualifier defaults to empty string closed
inRangeStart := func() bool { return col >= "" }
switch sq := f.ColumnRangeFilter.StartQualifier.(type) {
case *btpb.ColumnRange_StartQualifierOpen:
inRangeStart = func() bool { return col > string(sq.StartQualifierOpen) }
case *btpb.ColumnRange_StartQualifierClosed:
inRangeStart = func() bool { return col >= string(sq.StartQualifierClosed) }
}
// End qualifier defaults to no upper boundary
inRangeEnd := func() bool { return true }
switch eq := f.ColumnRangeFilter.EndQualifier.(type) {
case *btpb.ColumnRange_EndQualifierClosed:
inRangeEnd = func() bool { return col <= string(eq.EndQualifierClosed) }
case *btpb.ColumnRange_EndQualifierOpen:
inRangeEnd = func() bool { return col < string(eq.EndQualifierOpen) }
}
return inRangeStart() && inRangeEnd(), nil
case *btpb.RowFilter_TimestampRangeFilter:
// Lower bound is inclusive and defaults to 0, upper bound is exclusive and defaults to infinity.
return cell.ts >= f.TimestampRangeFilter.StartTimestampMicros &&
(f.TimestampRangeFilter.EndTimestampMicros == 0 || cell.ts < f.TimestampRangeFilter.EndTimestampMicros), nil
case *btpb.RowFilter_ValueRangeFilter:
v := cell.value
// Start value defaults to empty string closed
inRangeStart := func() bool { return bytes.Compare(v, []byte{}) >= 0 }
switch sv := f.ValueRangeFilter.StartValue.(type) {
case *btpb.ValueRange_StartValueOpen:
inRangeStart = func() bool { return bytes.Compare(v, sv.StartValueOpen) > 0 }
case *btpb.ValueRange_StartValueClosed:
inRangeStart = func() bool { return bytes.Compare(v, sv.StartValueClosed) >= 0 }
}
// End value defaults to no upper boundary
inRangeEnd := func() bool { return true }
switch ev := f.ValueRangeFilter.EndValue.(type) {
case *btpb.ValueRange_EndValueClosed:
inRangeEnd = func() bool { return bytes.Compare(v, ev.EndValueClosed) <= 0 }
case *btpb.ValueRange_EndValueOpen:
inRangeEnd = func() bool { return bytes.Compare(v, ev.EndValueOpen) < 0 }
}
return inRangeStart() && inRangeEnd(), nil
}
}
func toUTF8(bs []byte) string {
var rs []rune
for _, b := range bs {
rs = append(rs, rune(b))
}
return string(rs)
}
func newRegexp(patBytes []byte) (*regexp.Regexp, error) {
pat := toUTF8(patBytes)
re, err := regexp.Compile("^" + pat + "$") // match entire target
if err != nil {
log.Printf("Bad pattern %q: %v", pat, err)
}
return re, err
}
func (s *server) MutateRow(ctx context.Context, req *btpb.MutateRowRequest) (*btpb.MutateRowResponse, error) {
s.mu.Lock()
tbl, ok := s.tables[req.TableName]
s.mu.Unlock()
if !ok {
return nil, status.Errorf(codes.NotFound, "table %q not found", req.TableName)
}
fs := tbl.columnFamilies()
r := tbl.mutableRow(string(req.RowKey))
r.mu.Lock()
defer r.mu.Unlock()
if err := applyMutations(tbl, r, req.Mutations, fs); err != nil {
return nil, err
}
return &btpb.MutateRowResponse{}, nil
}
func (s *server) MutateRows(req *btpb.MutateRowsRequest, stream btpb.Bigtable_MutateRowsServer) error {
s.mu.Lock()
tbl, ok := s.tables[req.TableName]
s.mu.Unlock()
if !ok {
return status.Errorf(codes.NotFound, "table %q not found", req.TableName)
}
res := &btpb.MutateRowsResponse{Entries: make([]*btpb.MutateRowsResponse_Entry, len(req.Entries))}
fs := tbl.columnFamilies()
for i, entry := range req.Entries {
r := tbl.mutableRow(string(entry.RowKey))
r.mu.Lock()
code, msg := int32(codes.OK), ""
if err := applyMutations(tbl, r, entry.Mutations, fs); err != nil {
code = int32(codes.Internal)
msg = err.Error()
}
res.Entries[i] = &btpb.MutateRowsResponse_Entry{
Index: int64(i),
Status: &statpb.Status{Code: code, Message: msg},
}
r.mu.Unlock()
}
return stream.Send(res)
}
func (s *server) CheckAndMutateRow(ctx context.Context, req *btpb.CheckAndMutateRowRequest) (*btpb.CheckAndMutateRowResponse, error) {
s.mu.Lock()
tbl, ok := s.tables[req.TableName]
s.mu.Unlock()
if !ok {
return nil, status.Errorf(codes.NotFound, "table %q not found", req.TableName)
}
res := &btpb.CheckAndMutateRowResponse{}
fs := tbl.columnFamilies()
r := tbl.mutableRow(string(req.RowKey))
r.mu.Lock()
defer r.mu.Unlock()
// Figure out which mutation to apply.
whichMut := false
if req.PredicateFilter == nil {
// Use true_mutations iff row contains any cells.
whichMut = !r.isEmpty()
} else {
// Use true_mutations iff any cells in the row match the filter.
// TODO(dsymonds): This could be cheaper.
nr := r.copy()
filterRow(req.PredicateFilter, nr)
whichMut = !nr.isEmpty()
}
res.PredicateMatched = whichMut
muts := req.FalseMutations
if whichMut {
muts = req.TrueMutations
}
if err := applyMutations(tbl, r, muts, fs); err != nil {
return nil, err
}
return res, nil
}
// applyMutations applies a sequence of mutations to a row.
// fam should be a snapshot of the keys of tbl.families.
// It assumes r.mu is locked.
func applyMutations(tbl *table, r *row, muts []*btpb.Mutation, fs map[string]*columnFamily) error {
for _, mut := range muts {
switch mut := mut.Mutation.(type) {
default:
return fmt.Errorf("can't handle mutation type %T", mut)
case *btpb.Mutation_SetCell_:
set := mut.SetCell
if _, ok := fs[set.FamilyName]; !ok {
return fmt.Errorf("unknown family %q", set.FamilyName)
}
ts := set.TimestampMicros
if ts == -1 { // bigtable.ServerTime
ts = newTimestamp()
}
if !tbl.validTimestamp(ts) {
return fmt.Errorf("invalid timestamp %d", ts)
}
fam := set.FamilyName
col := string(set.ColumnQualifier)
newCell := cell{ts: ts, value: set.Value}
f := r.getOrCreateFamily(fam, fs[fam].order)
f.cells[col] = appendOrReplaceCell(f.cellsByColumn(col), newCell)
case *btpb.Mutation_DeleteFromColumn_:
del := mut.DeleteFromColumn
if _, ok := fs[del.FamilyName]; !ok {
return fmt.Errorf("unknown family %q", del.FamilyName)
}
fam := del.FamilyName
col := string(del.ColumnQualifier)
if _, ok := r.families[fam]; ok {
cs := r.families[fam].cells[col]
if del.TimeRange != nil {
tsr := del.TimeRange
if !tbl.validTimestamp(tsr.StartTimestampMicros) {
return fmt.Errorf("invalid timestamp %d", tsr.StartTimestampMicros)
}
if !tbl.validTimestamp(tsr.EndTimestampMicros) && tsr.EndTimestampMicros != 0 {
return fmt.Errorf("invalid timestamp %d", tsr.EndTimestampMicros)
}
if tsr.StartTimestampMicros >= tsr.EndTimestampMicros && tsr.EndTimestampMicros != 0 {
return fmt.Errorf("inverted or invalid timestamp range [%d, %d]", tsr.StartTimestampMicros, tsr.EndTimestampMicros)
}
// Find half-open interval to remove.
// Cells are in descending timestamp order,
// so the predicates to sort.Search are inverted.
si, ei := 0, len(cs)
if tsr.StartTimestampMicros > 0 {
ei = sort.Search(len(cs), func(i int) bool { return cs[i].ts < tsr.StartTimestampMicros })
}
if tsr.EndTimestampMicros > 0 {
si = sort.Search(len(cs), func(i int) bool { return cs[i].ts < tsr.EndTimestampMicros })
}
if si < ei {
copy(cs[si:], cs[ei:])
cs = cs[:len(cs)-(ei-si)]
}
} else {
cs = nil
}
if len(cs) == 0 {
delete(r.families[fam].cells, col)
colNames := r.families[fam].colNames
i := sort.Search(len(colNames), func(i int) bool { return colNames[i] >= col })
if i < len(colNames) && colNames[i] == col {
r.families[fam].colNames = append(colNames[:i], colNames[i+1:]...)
}
if len(r.families[fam].cells) == 0 {
delete(r.families, fam)
}
} else {
r.families[fam].cells[col] = cs
}
}
case *btpb.Mutation_DeleteFromRow_:
r.families = make(map[string]*family)
case *btpb.Mutation_DeleteFromFamily_:
fampre := mut.DeleteFromFamily.FamilyName
delete(r.families, fampre)
}
}
return nil
}
func maxTimestamp(x, y int64) int64 {
if x > y {
return x
}
return y
}
func newTimestamp() int64 {
ts := time.Now().UnixNano() / 1e3
ts -= ts % 1000 // round to millisecond granularity
return ts
}
func appendOrReplaceCell(cs []cell, newCell cell) []cell {
replaced := false
for i, cell := range cs {
if cell.ts == newCell.ts {
cs[i] = newCell
replaced = true
break
}
}
if !replaced {
cs = append(cs, newCell)
}
sort.Sort(byDescTS(cs))
return cs
}
func (s *server) ReadModifyWriteRow(ctx context.Context, req *btpb.ReadModifyWriteRowRequest) (*btpb.ReadModifyWriteRowResponse, error) {
s.mu.Lock()
tbl, ok := s.tables[req.TableName]
s.mu.Unlock()
if !ok {
return nil, status.Errorf(codes.NotFound, "table %q not found", req.TableName)
}
fs := tbl.columnFamilies()
rowKey := string(req.RowKey)
r := tbl.mutableRow(rowKey)
resultRow := newRow(rowKey) // copy of updated cells
// This must be done before the row lock, acquired below, is released.
r.mu.Lock()
defer r.mu.Unlock()
// Assume all mutations apply to the most recent version of the cell.
// TODO(dsymonds): Verify this assumption and document it in the proto.
for _, rule := range req.Rules {
if _, ok := fs[rule.FamilyName]; !ok {
return nil, fmt.Errorf("unknown family %q", rule.FamilyName)
}
fam := rule.FamilyName
col := string(rule.ColumnQualifier)
isEmpty := false
f := r.getOrCreateFamily(fam, fs[fam].order)
cs := f.cells[col]
isEmpty = len(cs) == 0
ts := newTimestamp()
var newCell, prevCell cell
if !isEmpty {
cells := r.families[fam].cells[col]
prevCell = cells[0]
// ts is the max of now or the prev cell's timestamp in case the
// prev cell is in the future
ts = maxTimestamp(ts, prevCell.ts)
}
switch rule := rule.Rule.(type) {
default:
return nil, fmt.Errorf("unknown RMW rule oneof %T", rule)
case *btpb.ReadModifyWriteRule_AppendValue:
newCell = cell{ts: ts, value: append(prevCell.value, rule.AppendValue...)}
case *btpb.ReadModifyWriteRule_IncrementAmount:
var v int64
if !isEmpty {
prevVal := prevCell.value
if len(prevVal) != 8 {
return nil, fmt.Errorf("increment on non-64-bit value")
}
v = int64(binary.BigEndian.Uint64(prevVal))
}
v += rule.IncrementAmount
var val [8]byte
binary.BigEndian.PutUint64(val[:], uint64(v))
newCell = cell{ts: ts, value: val[:]}
}
// Store the new cell
f.cells[col] = appendOrReplaceCell(f.cellsByColumn(col), newCell)
// Store a copy for the result row
resultFamily := resultRow.getOrCreateFamily(fam, fs[fam].order)
resultFamily.cellsByColumn(col) // create the column
resultFamily.cells[col] = []cell{newCell} // overwrite the cells
}
// Build the response using the result row
res := &btpb.Row{
Key: req.RowKey,
Families: make([]*btpb.Family, len(resultRow.families)),
}
for i, family := range resultRow.sortedFamilies() {
res.Families[i] = &btpb.Family{
Name: family.name,
Columns: make([]*btpb.Column, len(family.colNames)),
}
for j, colName := range family.colNames {
res.Families[i].Columns[j] = &btpb.Column{
Qualifier: []byte(colName),
Cells: []*btpb.Cell{{
TimestampMicros: family.cells[colName][0].ts,
Value: family.cells[colName][0].value,
}},
}
}
}
return &btpb.ReadModifyWriteRowResponse{Row: res}, nil
}
func (s *server) SampleRowKeys(req *btpb.SampleRowKeysRequest, stream btpb.Bigtable_SampleRowKeysServer) error {
s.mu.Lock()
tbl, ok := s.tables[req.TableName]
s.mu.Unlock()
if !ok {
return status.Errorf(codes.NotFound, "table %q not found", req.TableName)
}
tbl.mu.RLock()
defer tbl.mu.RUnlock()
// The return value of SampleRowKeys is very loosely defined. Return at least the
// final row key in the table and choose other row keys randomly.
var offset int64
var err error
i := 0
tbl.rows.Ascend(func(it btree.Item) bool {
row := it.(*row)
if i == tbl.rows.Len()-1 || rand.Int31n(100) == 0 {
resp := &btpb.SampleRowKeysResponse{
RowKey: []byte(row.key),
OffsetBytes: offset,
}
err = stream.Send(resp)
if err != nil {
return false
}
}
offset += int64(row.size())
i++
return true
})
return err
}
// needGC is invoked whenever the server needs gcloop running.
func (s *server) needGC() {
s.mu.Lock()
if s.gcc == nil {
s.gcc = make(chan int)
go s.gcloop(s.gcc)
}
s.mu.Unlock()
}
func (s *server) gcloop(done <-chan int) {
const (
minWait = 500 // ms
maxWait = 1500 // ms
)
for {
// Wait for a random time interval.
d := time.Duration(minWait+rand.Intn(maxWait-minWait)) * time.Millisecond
select {
case <-time.After(d):
case <-done:
return // server has been closed
}
// Do a GC pass over all tables.
var tables []*table
s.mu.Lock()
for _, tbl := range s.tables {
tables = append(tables, tbl)
}
s.mu.Unlock()
for _, tbl := range tables {
tbl.gc()
}
}
}
type table struct {
mu sync.RWMutex
counter uint64 // increment by 1 when a new family is created
families map[string]*columnFamily // keyed by plain family name
rows *btree.BTree // indexed by row key
}
const btreeDegree = 16
func newTable(ctr *btapb.CreateTableRequest) *table {
fams := make(map[string]*columnFamily)
c := uint64(0)
if ctr.Table != nil {
for id, cf := range ctr.Table.ColumnFamilies {
fams[id] = &columnFamily{
name: ctr.Parent + "/columnFamilies/" + id,
order: c,
gcRule: cf.GcRule,
}
c++
}
}
return &table{
families: fams,
counter: c,
rows: btree.New(btreeDegree),
}
}
func (t *table) validTimestamp(ts int64) bool {
if ts < minValidMilliSeconds || ts > maxValidMilliSeconds {
return false
}
// Assume millisecond granularity is required.
return ts%1000 == 0
}
func (t *table) columnFamilies() map[string]*columnFamily {
cp := make(map[string]*columnFamily)
t.mu.RLock()
for fam, cf := range t.families {
cp[fam] = cf
}
t.mu.RUnlock()
return cp
}
func (t *table) mutableRow(key string) *row {
bkey := btreeKey(key)
// Try fast path first.
t.mu.RLock()
i := t.rows.Get(bkey)
t.mu.RUnlock()
if i != nil {
return i.(*row)
}
// We probably need to create the row.
t.mu.Lock()
defer t.mu.Unlock()
i = t.rows.Get(bkey)
if i != nil {
return i.(*row)
}
r := newRow(key)
t.rows.ReplaceOrInsert(r)
return r
}
func (t *table) gc() {
// This method doesn't add or remove rows, so we only need a read lock for the table.
t.mu.RLock()
defer t.mu.RUnlock()
// Gather GC rules we'll apply.
rules := make(map[string]*btapb.GcRule) // keyed by "fam"
for fam, cf := range t.families {
if cf.gcRule != nil {
rules[fam] = cf.gcRule
}
}
if len(rules) == 0 {
return
}
t.rows.Ascend(func(i btree.Item) bool {
r := i.(*row)
r.mu.Lock()
r.gc(rules)
r.mu.Unlock()
return true
})
}
type byRowKey []*row
func (b byRowKey) Len() int { return len(b) }
func (b byRowKey) Swap(i, j int) { b[i], b[j] = b[j], b[i] }
func (b byRowKey) Less(i, j int) bool { return b[i].key < b[j].key }
type row struct {
key string
mu sync.Mutex
families map[string]*family // keyed by family name
}
func newRow(key string) *row {
return &row{
key: key,
families: make(map[string]*family),
}
}
// copy returns a copy of the row.
// Cell values are aliased.
// r.mu should be held.
func (r *row) copy() *row {
nr := newRow(r.key)
for _, fam := range r.families {
nr.families[fam.name] = &family{
name: fam.name,
order: fam.order,
colNames: fam.colNames,
cells: make(map[string][]cell),
}
for col, cs := range fam.cells {
// Copy the []cell slice, but not the []byte inside each cell.
nr.families[fam.name].cells[col] = append([]cell(nil), cs...)
}
}
return nr
}
// isEmpty returns true if a row doesn't contain any cell
func (r *row) isEmpty() bool {
for _, fam := range r.families {
for _, cs := range fam.cells {
if len(cs) > 0 {
return false
}
}
}
return true
}
// sortedFamilies returns a column family set
// sorted in ascending creation order in a row.
func (r *row) sortedFamilies() []*family {
var families []*family
for _, fam := range r.families {
families = append(families, fam)
}
sort.Sort(byCreationOrder(families))
return families
}
func (r *row) getOrCreateFamily(name string, order uint64) *family {
if _, ok := r.families[name]; !ok {
r.families[name] = &family{
name: name,
order: order,
cells: make(map[string][]cell),
}
}
return r.families[name]
}
// gc applies the given GC rules to the row.
// r.mu should be held.
func (r *row) gc(rules map[string]*btapb.GcRule) {
for _, fam := range r.families {
rule, ok := rules[fam.name]
if !ok {
continue
}
for col, cs := range fam.cells {
r.families[fam.name].cells[col] = applyGC(cs, rule)
}
}
}
// size returns the total size of all cell values in the row.
func (r *row) size() int {
size := 0
for _, fam := range r.families {
for _, cells := range fam.cells {
for _, cell := range cells {
size += len(cell.value)
}
}
}
return size
}
// Less implements btree.Less.
func (r *row) Less(i btree.Item) bool {
return r.key < i.(*row).key
}
// btreeKey returns a row for use as a key into the BTree.
func btreeKey(s string) *row { return &row{key: s} }
func (r *row) String() string {
return r.key
}
var gcTypeWarn sync.Once
// applyGC applies the given GC rule to the cells.
func applyGC(cells []cell, rule *btapb.GcRule) []cell {
switch rule := rule.Rule.(type) {
default:
// TODO(dsymonds): Support GcRule_Intersection_
gcTypeWarn.Do(func() {
log.Printf("Unsupported GC rule type %T", rule)
})
case *btapb.GcRule_Union_:
for _, sub := range rule.Union.Rules {
cells = applyGC(cells, sub)
}
return cells
case *btapb.GcRule_MaxAge:
// Timestamps are in microseconds.
cutoff := time.Now().UnixNano() / 1e3
cutoff -= rule.MaxAge.Seconds * 1e6
cutoff -= int64(rule.MaxAge.Nanos) / 1e3
// The slice of cells in in descending timestamp order.
// This sort.Search will return the index of the first cell whose timestamp is chronologically before the cutoff.
si := sort.Search(len(cells), func(i int) bool { return cells[i].ts < cutoff })
if si < len(cells) {
log.Printf("bttest: GC MaxAge(%v) deleted %d cells.", rule.MaxAge, len(cells)-si)
}
return cells[:si]
case *btapb.GcRule_MaxNumVersions:
n := int(rule.MaxNumVersions)
if len(cells) > n {
cells = cells[:n]
}
return cells
}
return cells
}
type family struct {
name string // Column family name
order uint64 // Creation order of column family
colNames []string // Column names are sorted in lexicographical ascending order
cells map[string][]cell // Keyed by column name; cells are in descending timestamp order
}
type byCreationOrder []*family
func (b byCreationOrder) Len() int { return len(b) }
func (b byCreationOrder) Swap(i, j int) { b[i], b[j] = b[j], b[i] }
func (b byCreationOrder) Less(i, j int) bool { return b[i].order < b[j].order }
// cellsByColumn adds the column name to colNames set if it does not exist
// and returns all cells within a column
func (f *family) cellsByColumn(name string) []cell {
if _, ok := f.cells[name]; !ok {
f.colNames = append(f.colNames, name)
sort.Strings(f.colNames)
}
return f.cells[name]
}
type cell struct {
ts int64
value []byte
labels []string
}
type byDescTS []cell
func (b byDescTS) Len() int { return len(b) }
func (b byDescTS) Swap(i, j int) { b[i], b[j] = b[j], b[i] }
func (b byDescTS) Less(i, j int) bool { return b[i].ts > b[j].ts }
type columnFamily struct {
name string
order uint64 // Creation order of column family
gcRule *btapb.GcRule
}
func (c *columnFamily) proto() *btapb.ColumnFamily {
return &btapb.ColumnFamily{
GcRule: c.gcRule,
}
}
func toColumnFamilies(families map[string]*columnFamily) map[string]*btapb.ColumnFamily {
fs := make(map[string]*btapb.ColumnFamily)
for k, v := range families {
fs[k] = v.proto()
}
return fs
}
// Copyright 2016 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package bttest
import (
"context"
"fmt"
"math/rand"
"strconv"
"sync"
"sync/atomic"
"testing"
"time"
"github.com/golang/protobuf/proto"
"github.com/google/go-cmp/cmp"
"github.com/google/go-cmp/cmp/cmpopts"
btapb "google.golang.org/genproto/googleapis/bigtable/admin/v2"
btpb "google.golang.org/genproto/googleapis/bigtable/v2"
"google.golang.org/grpc"
)
func TestConcurrentMutationsReadModifyAndGC(t *testing.T) {
s := &server{
tables: make(map[string]*table),
}
ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
defer cancel()
if _, err := s.CreateTable(
ctx,
&btapb.CreateTableRequest{Parent: "cluster", TableId: "t"}); err != nil {
t.Fatal(err)
}
const name = `cluster/tables/t`
tbl := s.tables[name]
req := &btapb.ModifyColumnFamiliesRequest{
Name: name,
Modifications: []*btapb.ModifyColumnFamiliesRequest_Modification{{
Id: "cf",
Mod: &btapb.ModifyColumnFamiliesRequest_Modification_Create{Create: &btapb.ColumnFamily{}},
}},
}
_, err := s.ModifyColumnFamilies(ctx, req)
if err != nil {
t.Fatal(err)
}
req = &btapb.ModifyColumnFamiliesRequest{
Name: name,
Modifications: []*btapb.ModifyColumnFamiliesRequest_Modification{{
Id: "cf",
Mod: &btapb.ModifyColumnFamiliesRequest_Modification_Update{Update: &btapb.ColumnFamily{
GcRule: &btapb.GcRule{Rule: &btapb.GcRule_MaxNumVersions{MaxNumVersions: 1}},
}},
}},
}
if _, err := s.ModifyColumnFamilies(ctx, req); err != nil {
t.Fatal(err)
}
var wg sync.WaitGroup
var ts int64
ms := func() []*btpb.Mutation {
return []*btpb.Mutation{{
Mutation: &btpb.Mutation_SetCell_{SetCell: &btpb.Mutation_SetCell{
FamilyName: "cf",
ColumnQualifier: []byte(`col`),
TimestampMicros: atomic.AddInt64(&ts, 1000),
}},
}}
}
rmw := func() *btpb.ReadModifyWriteRowRequest {
return &btpb.ReadModifyWriteRowRequest{
TableName: name,
RowKey: []byte(fmt.Sprint(rand.Intn(100))),
Rules: []*btpb.ReadModifyWriteRule{{
FamilyName: "cf",
ColumnQualifier: []byte("col"),
Rule: &btpb.ReadModifyWriteRule_IncrementAmount{IncrementAmount: 1},
}},
}
}
for i := 0; i < 100; i++ {
wg.Add(1)
go func() {
defer wg.Done()
for ctx.Err() == nil {
req := &btpb.MutateRowRequest{
TableName: name,
RowKey: []byte(fmt.Sprint(rand.Intn(100))),
Mutations: ms(),
}
if _, err := s.MutateRow(ctx, req); err != nil {
panic(err) // can't use t.Fatal in goroutine
}
}
}()
wg.Add(1)
go func() {
defer wg.Done()
for ctx.Err() == nil {
_, _ = s.ReadModifyWriteRow(ctx, rmw())
}
}()
wg.Add(1)
go func() {
defer wg.Done()
tbl.gc()
}()
}
done := make(chan struct{})
go func() {
wg.Wait()
close(done)
}()
select {
case <-done:
case <-time.After(1 * time.Second):
t.Error("Concurrent mutations and GCs haven't completed after 1s")
}
}
func TestCreateTableWithFamily(t *testing.T) {
// The Go client currently doesn't support creating a table with column families
// in one operation but it is allowed by the API. This must still be supported by the
// fake server so this test lives here instead of in the main bigtable
// integration test.
s := &server{
tables: make(map[string]*table),
}
ctx := context.Background()
newTbl := btapb.Table{
ColumnFamilies: map[string]*btapb.ColumnFamily{
"cf1": {GcRule: &btapb.GcRule{Rule: &btapb.GcRule_MaxNumVersions{MaxNumVersions: 123}}},
"cf2": {GcRule: &btapb.GcRule{Rule: &btapb.GcRule_MaxNumVersions{MaxNumVersions: 456}}},
},
}
cTbl, err := s.CreateTable(ctx, &btapb.CreateTableRequest{Parent: "cluster", TableId: "t", Table: &newTbl})
if err != nil {
t.Fatalf("Creating table: %v", err)
}
tbl, err := s.GetTable(ctx, &btapb.GetTableRequest{Name: cTbl.Name})
if err != nil {
t.Fatalf("Getting table: %v", err)
}
cf := tbl.ColumnFamilies["cf1"]
if cf == nil {
t.Fatalf("Missing col family cf1")
}
if got, want := cf.GcRule.GetMaxNumVersions(), int32(123); got != want {
t.Errorf("Invalid MaxNumVersions: wanted:%d, got:%d", want, got)
}
cf = tbl.ColumnFamilies["cf2"]
if cf == nil {
t.Fatalf("Missing col family cf2")
}
if got, want := cf.GcRule.GetMaxNumVersions(), int32(456); got != want {
t.Errorf("Invalid MaxNumVersions: wanted:%d, got:%d", want, got)
}
}
type MockSampleRowKeysServer struct {
responses []*btpb.SampleRowKeysResponse
grpc.ServerStream
}
func (s *MockSampleRowKeysServer) Send(resp *btpb.SampleRowKeysResponse) error {
s.responses = append(s.responses, resp)
return nil
}
func TestSampleRowKeys(t *testing.T) {
s := &server{
tables: make(map[string]*table),
}
ctx := context.Background()
newTbl := btapb.Table{
ColumnFamilies: map[string]*btapb.ColumnFamily{
"cf": {GcRule: &btapb.GcRule{Rule: &btapb.GcRule_MaxNumVersions{MaxNumVersions: 1}}},
},
}
tbl, err := s.CreateTable(ctx, &btapb.CreateTableRequest{Parent: "cluster", TableId: "t", Table: &newTbl})
if err != nil {
t.Fatalf("Creating table: %v", err)
}
// Populate the table
val := []byte("value")
rowCount := 1000
for i := 0; i < rowCount; i++ {
req := &btpb.MutateRowRequest{
TableName: tbl.Name,
RowKey: []byte("row-" + strconv.Itoa(i)),
Mutations: []*btpb.Mutation{{
Mutation: &btpb.Mutation_SetCell_{SetCell: &btpb.Mutation_SetCell{
FamilyName: "cf",
ColumnQualifier: []byte("col"),
TimestampMicros: 1000,
Value: val,
}},
}},
}
if _, err := s.MutateRow(ctx, req); err != nil {
t.Fatalf("Populating table: %v", err)
}
}
mock := &MockSampleRowKeysServer{}
if err := s.SampleRowKeys(&btpb.SampleRowKeysRequest{TableName: tbl.Name}, mock); err != nil {
t.Errorf("SampleRowKeys error: %v", err)
}
if len(mock.responses) == 0 {
t.Fatal("Response count: got 0, want > 0")
}
// Make sure the offset of the final response is the offset of the final row
got := mock.responses[len(mock.responses)-1].OffsetBytes
want := int64((rowCount - 1) * len(val))
if got != want {
t.Errorf("Invalid offset: got %d, want %d", got, want)
}
}
func TestDropRowRange(t *testing.T) {
s := &server{
tables: make(map[string]*table),
}
ctx := context.Background()
newTbl := btapb.Table{
ColumnFamilies: map[string]*btapb.ColumnFamily{
"cf": {GcRule: &btapb.GcRule{Rule: &btapb.GcRule_MaxNumVersions{MaxNumVersions: 1}}},
},
}
tblInfo, err := s.CreateTable(ctx, &btapb.CreateTableRequest{Parent: "cluster", TableId: "t", Table: &newTbl})
if err != nil {
t.Fatalf("Creating table: %v", err)
}
tbl := s.tables[tblInfo.Name]
// Populate the table
prefixes := []string{"AAA", "BBB", "CCC", "DDD"}
count := 3
doWrite := func() {
for _, prefix := range prefixes {
for i := 0; i < count; i++ {
req := &btpb.MutateRowRequest{
TableName: tblInfo.Name,
RowKey: []byte(prefix + strconv.Itoa(i)),
Mutations: []*btpb.Mutation{{
Mutation: &btpb.Mutation_SetCell_{SetCell: &btpb.Mutation_SetCell{
FamilyName: "cf",
ColumnQualifier: []byte("col"),
TimestampMicros: 1000,
Value: []byte{},
}},
}},
}
if _, err := s.MutateRow(ctx, req); err != nil {
t.Fatalf("Populating table: %v", err)
}
}
}
}
doWrite()
tblSize := tbl.rows.Len()
req := &btapb.DropRowRangeRequest{
Name: tblInfo.Name,
Target: &btapb.DropRowRangeRequest_RowKeyPrefix{RowKeyPrefix: []byte("AAA")},
}
if _, err = s.DropRowRange(ctx, req); err != nil {
t.Fatalf("Dropping first range: %v", err)
}
got, want := tbl.rows.Len(), tblSize-count
if got != want {
t.Errorf("Row count after first drop: got %d (%v), want %d", got, tbl.rows, want)
}
req = &btapb.DropRowRangeRequest{
Name: tblInfo.Name,
Target: &btapb.DropRowRangeRequest_RowKeyPrefix{RowKeyPrefix: []byte("DDD")},
}
if _, err = s.DropRowRange(ctx, req); err != nil {
t.Fatalf("Dropping second range: %v", err)
}
got, want = tbl.rows.Len(), tblSize-(2*count)
if got != want {
t.Errorf("Row count after second drop: got %d (%v), want %d", got, tbl.rows, want)
}
req = &btapb.DropRowRangeRequest{
Name: tblInfo.Name,
Target: &btapb.DropRowRangeRequest_RowKeyPrefix{RowKeyPrefix: []byte("XXX")},
}
if _, err = s.DropRowRange(ctx, req); err != nil {
t.Fatalf("Dropping invalid range: %v", err)
}
got, want = tbl.rows.Len(), tblSize-(2*count)
if got != want {
t.Errorf("Row count after invalid drop: got %d (%v), want %d", got, tbl.rows, want)
}
req = &btapb.DropRowRangeRequest{
Name: tblInfo.Name,
Target: &btapb.DropRowRangeRequest_DeleteAllDataFromTable{DeleteAllDataFromTable: true},
}
if _, err = s.DropRowRange(ctx, req); err != nil {
t.Fatalf("Dropping all data: %v", err)
}
got, want = tbl.rows.Len(), 0
if got != want {
t.Errorf("Row count after drop all: got %d, want %d", got, want)
}
// Test that we can write rows, delete some and then write them again.
count = 1
doWrite()
req = &btapb.DropRowRangeRequest{
Name: tblInfo.Name,
Target: &btapb.DropRowRangeRequest_DeleteAllDataFromTable{DeleteAllDataFromTable: true},
}
if _, err = s.DropRowRange(ctx, req); err != nil {
t.Fatalf("Dropping all data: %v", err)
}
got, want = tbl.rows.Len(), 0
if got != want {
t.Errorf("Row count after drop all: got %d, want %d", got, want)
}
doWrite()
got, want = tbl.rows.Len(), len(prefixes)
if got != want {
t.Errorf("Row count after rewrite: got %d, want %d", got, want)
}
req = &btapb.DropRowRangeRequest{
Name: tblInfo.Name,
Target: &btapb.DropRowRangeRequest_RowKeyPrefix{RowKeyPrefix: []byte("BBB")},
}
if _, err = s.DropRowRange(ctx, req); err != nil {
t.Fatalf("Dropping range: %v", err)
}
doWrite()
got, want = tbl.rows.Len(), len(prefixes)
if got != want {
t.Errorf("Row count after drop range: got %d, want %d", got, want)
}
}
type MockReadRowsServer struct {
responses []*btpb.ReadRowsResponse
grpc.ServerStream
}
func (s *MockReadRowsServer) Send(resp *btpb.ReadRowsResponse) error {
s.responses = append(s.responses, resp)
return nil
}
func TestReadRows(t *testing.T) {
ctx := context.Background()
s := &server{
tables: make(map[string]*table),
}
newTbl := btapb.Table{
ColumnFamilies: map[string]*btapb.ColumnFamily{
"cf0": {GcRule: &btapb.GcRule{Rule: &btapb.GcRule_MaxNumVersions{MaxNumVersions: 1}}},
},
}
tblInfo, err := s.CreateTable(ctx, &btapb.CreateTableRequest{Parent: "cluster", TableId: "t", Table: &newTbl})
if err != nil {
t.Fatalf("Creating table: %v", err)
}
mreq := &btpb.MutateRowRequest{
TableName: tblInfo.Name,
RowKey: []byte("row"),
Mutations: []*btpb.Mutation{{
Mutation: &btpb.Mutation_SetCell_{SetCell: &btpb.Mutation_SetCell{
FamilyName: "cf0",
ColumnQualifier: []byte("col"),
TimestampMicros: 1000,
Value: []byte{},
}},
}},
}
if _, err := s.MutateRow(ctx, mreq); err != nil {
t.Fatalf("Populating table: %v", err)
}
for _, rowset := range []*btpb.RowSet{
{RowKeys: [][]byte{[]byte("row")}},
{RowRanges: []*btpb.RowRange{{StartKey: &btpb.RowRange_StartKeyClosed{StartKeyClosed: []byte("")}}}},
{RowRanges: []*btpb.RowRange{{StartKey: &btpb.RowRange_StartKeyClosed{StartKeyClosed: []byte("r")}}}},
{RowRanges: []*btpb.RowRange{{
StartKey: &btpb.RowRange_StartKeyClosed{StartKeyClosed: []byte("")},
EndKey: &btpb.RowRange_EndKeyOpen{EndKeyOpen: []byte("s")},
}}},
} {
mock := &MockReadRowsServer{}
req := &btpb.ReadRowsRequest{TableName: tblInfo.Name, Rows: rowset}
if err = s.ReadRows(req, mock); err != nil {
t.Fatalf("ReadRows error: %v", err)
}
if got, want := len(mock.responses), 1; got != want {
t.Errorf("%+v: response count: got %d, want %d", rowset, got, want)
}
}
}
func TestReadRowsError(t *testing.T) {
ctx := context.Background()
s := &server{
tables: make(map[string]*table),
}
newTbl := btapb.Table{
ColumnFamilies: map[string]*btapb.ColumnFamily{
"cf0": {GcRule: &btapb.GcRule{Rule: &btapb.GcRule_MaxNumVersions{MaxNumVersions: 1}}},
},
}
tblInfo, err := s.CreateTable(ctx, &btapb.CreateTableRequest{Parent: "cluster", TableId: "t", Table: &newTbl})
if err != nil {
t.Fatalf("Creating table: %v", err)
}
mreq := &btpb.MutateRowRequest{
TableName: tblInfo.Name,
RowKey: []byte("row"),
Mutations: []*btpb.Mutation{{
Mutation: &btpb.Mutation_SetCell_{SetCell: &btpb.Mutation_SetCell{
FamilyName: "cf0",
ColumnQualifier: []byte("col"),
TimestampMicros: 1000,
Value: []byte{},
}},
}},
}
if _, err := s.MutateRow(ctx, mreq); err != nil {
t.Fatalf("Populating table: %v", err)
}
mock := &MockReadRowsServer{}
req := &btpb.ReadRowsRequest{TableName: tblInfo.Name, Filter: &btpb.RowFilter{
Filter: &btpb.RowFilter_RowKeyRegexFilter{RowKeyRegexFilter: []byte("[")}}, // Invalid regex.
}
if err = s.ReadRows(req, mock); err == nil {
t.Fatal("ReadRows got no error, want error")
}
}
func TestReadRowsOrder(t *testing.T) {
s := &server{
tables: make(map[string]*table),
}
ctx := context.Background()
newTbl := btapb.Table{
ColumnFamilies: map[string]*btapb.ColumnFamily{
"cf0": {GcRule: &btapb.GcRule{Rule: &btapb.GcRule_MaxNumVersions{MaxNumVersions: 1}}},
},
}
tblInfo, err := s.CreateTable(ctx, &btapb.CreateTableRequest{Parent: "cluster", TableId: "t", Table: &newTbl})
if err != nil {
t.Fatalf("Creating table: %v", err)
}
count := 3
mcf := func(i int) *btapb.ModifyColumnFamiliesRequest {
return &btapb.ModifyColumnFamiliesRequest{
Name: tblInfo.Name,
Modifications: []*btapb.ModifyColumnFamiliesRequest_Modification{{
Id: "cf" + strconv.Itoa(i),
Mod: &btapb.ModifyColumnFamiliesRequest_Modification_Create{Create: &btapb.ColumnFamily{}},
}},
}
}
for i := 1; i <= count; i++ {
_, err = s.ModifyColumnFamilies(ctx, mcf(i))
if err != nil {
t.Fatal(err)
}
}
// Populate the table
for fc := 0; fc < count; fc++ {
for cc := count; cc > 0; cc-- {
for tc := 0; tc < count; tc++ {
req := &btpb.MutateRowRequest{
TableName: tblInfo.Name,
RowKey: []byte("row"),
Mutations: []*btpb.Mutation{{
Mutation: &btpb.Mutation_SetCell_{SetCell: &btpb.Mutation_SetCell{
FamilyName: "cf" + strconv.Itoa(fc),
ColumnQualifier: []byte("col" + strconv.Itoa(cc)),
TimestampMicros: int64((tc + 1) * 1000),
Value: []byte{},
}},
}},
}
if _, err := s.MutateRow(ctx, req); err != nil {
t.Fatalf("Populating table: %v", err)
}
}
}
}
req := &btpb.ReadRowsRequest{
TableName: tblInfo.Name,
Rows: &btpb.RowSet{RowKeys: [][]byte{[]byte("row")}},
}
mock := &MockReadRowsServer{}
if err = s.ReadRows(req, mock); err != nil {
t.Errorf("ReadRows error: %v", err)
}
if len(mock.responses) == 0 {
t.Fatal("Response count: got 0, want > 0")
}
if len(mock.responses[0].Chunks) != 27 {
t.Fatalf("Chunk count: got %d, want 27", len(mock.responses[0].Chunks))
}
testOrder := func(ms *MockReadRowsServer) {
var prevFam, prevCol string
var prevTime int64
for _, cc := range ms.responses[0].Chunks {
if prevFam == "" {
prevFam = cc.FamilyName.Value
prevCol = string(cc.Qualifier.Value)
prevTime = cc.TimestampMicros
continue
}
if cc.FamilyName.Value < prevFam {
t.Errorf("Family order is not correct: got %s < %s", cc.FamilyName.Value, prevFam)
} else if cc.FamilyName.Value == prevFam {
if string(cc.Qualifier.Value) < prevCol {
t.Errorf("Column order is not correct: got %s < %s", string(cc.Qualifier.Value), prevCol)
} else if string(cc.Qualifier.Value) == prevCol {
if cc.TimestampMicros > prevTime {
t.Errorf("cell order is not correct: got %d > %d", cc.TimestampMicros, prevTime)
}
}
}
prevFam = cc.FamilyName.Value
prevCol = string(cc.Qualifier.Value)
prevTime = cc.TimestampMicros
}
}
testOrder(mock)
// Read with interleave filter
inter := &btpb.RowFilter_Interleave{}
fnr := &btpb.RowFilter{Filter: &btpb.RowFilter_FamilyNameRegexFilter{FamilyNameRegexFilter: "cf1"}}
cqr := &btpb.RowFilter{Filter: &btpb.RowFilter_ColumnQualifierRegexFilter{ColumnQualifierRegexFilter: []byte("col2")}}
inter.Filters = append(inter.Filters, fnr, cqr)
req = &btpb.ReadRowsRequest{
TableName: tblInfo.Name,
Rows: &btpb.RowSet{RowKeys: [][]byte{[]byte("row")}},
Filter: &btpb.RowFilter{
Filter: &btpb.RowFilter_Interleave_{Interleave: inter},
},
}
mock = &MockReadRowsServer{}
if err = s.ReadRows(req, mock); err != nil {
t.Errorf("ReadRows error: %v", err)
}
if len(mock.responses) == 0 {
t.Fatal("Response count: got 0, want > 0")
}
if len(mock.responses[0].Chunks) != 18 {
t.Fatalf("Chunk count: got %d, want 18", len(mock.responses[0].Chunks))
}
testOrder(mock)
// Check order after ReadModifyWriteRow
rmw := func(i int) *btpb.ReadModifyWriteRowRequest {
return &btpb.ReadModifyWriteRowRequest{
TableName: tblInfo.Name,
RowKey: []byte("row"),
Rules: []*btpb.ReadModifyWriteRule{{
FamilyName: "cf3",
ColumnQualifier: []byte("col" + strconv.Itoa(i)),
Rule: &btpb.ReadModifyWriteRule_IncrementAmount{IncrementAmount: 1},
}},
}
}
for i := count; i > 0; i-- {
if _, err := s.ReadModifyWriteRow(ctx, rmw(i)); err != nil {
t.Fatal(err)
}
}
req = &btpb.ReadRowsRequest{
TableName: tblInfo.Name,
Rows: &btpb.RowSet{RowKeys: [][]byte{[]byte("row")}},
}
mock = &MockReadRowsServer{}
if err = s.ReadRows(req, mock); err != nil {
t.Errorf("ReadRows error: %v", err)
}
if len(mock.responses) == 0 {
t.Fatal("Response count: got 0, want > 0")
}
if len(mock.responses[0].Chunks) != 30 {
t.Fatalf("Chunk count: got %d, want 30", len(mock.responses[0].Chunks))
}
testOrder(mock)
}
func TestReadRowsWithlabelTransformer(t *testing.T) {
ctx := context.Background()
s := &server{
tables: make(map[string]*table),
}
newTbl := btapb.Table{
ColumnFamilies: map[string]*btapb.ColumnFamily{
"cf0": {GcRule: &btapb.GcRule{Rule: &btapb.GcRule_MaxNumVersions{MaxNumVersions: 1}}},
},
}
tblInfo, err := s.CreateTable(ctx, &btapb.CreateTableRequest{Parent: "cluster", TableId: "t", Table: &newTbl})
if err != nil {
t.Fatalf("Creating table: %v", err)
}
mreq := &btpb.MutateRowRequest{
TableName: tblInfo.Name,
RowKey: []byte("row"),
Mutations: []*btpb.Mutation{{
Mutation: &btpb.Mutation_SetCell_{SetCell: &btpb.Mutation_SetCell{
FamilyName: "cf0",
ColumnQualifier: []byte("col"),
TimestampMicros: 1000,
Value: []byte{},
}},
}},
}
if _, err := s.MutateRow(ctx, mreq); err != nil {
t.Fatalf("Populating table: %v", err)
}
mock := &MockReadRowsServer{}
req := &btpb.ReadRowsRequest{
TableName: tblInfo.Name,
Filter: &btpb.RowFilter{
Filter: &btpb.RowFilter_ApplyLabelTransformer{
ApplyLabelTransformer: "label",
},
},
}
if err = s.ReadRows(req, mock); err != nil {
t.Fatalf("ReadRows error: %v", err)
}
if got, want := len(mock.responses), 1; got != want {
t.Fatalf("response count: got %d, want %d", got, want)
}
resp := mock.responses[0]
if got, want := len(resp.Chunks), 1; got != want {
t.Fatalf("chunks count: got %d, want %d", got, want)
}
chunk := resp.Chunks[0]
if got, want := len(chunk.Labels), 1; got != want {
t.Fatalf("labels count: got %d, want %d", got, want)
}
if got, want := chunk.Labels[0], "label"; got != want {
t.Fatalf("label: got %s, want %s", got, want)
}
mock = &MockReadRowsServer{}
req = &btpb.ReadRowsRequest{
TableName: tblInfo.Name,
Filter: &btpb.RowFilter{
Filter: &btpb.RowFilter_ApplyLabelTransformer{
ApplyLabelTransformer: "", // invalid label
},
},
}
if err = s.ReadRows(req, mock); err == nil {
t.Fatal("ReadRows want invalid label error, got none")
}
}
func TestCheckAndMutateRowWithoutPredicate(t *testing.T) {
s := &server{
tables: make(map[string]*table),
}
ctx := context.Background()
newTbl := btapb.Table{
ColumnFamilies: map[string]*btapb.ColumnFamily{
"cf": {GcRule: &btapb.GcRule{Rule: &btapb.GcRule_MaxNumVersions{MaxNumVersions: 1}}},
},
}
tbl, err := s.CreateTable(ctx, &btapb.CreateTableRequest{Parent: "cluster", TableId: "t", Table: &newTbl})
if err != nil {
t.Fatalf("Creating table: %v", err)
}
// Populate the table
val := []byte("value")
mrreq := &btpb.MutateRowRequest{
TableName: tbl.Name,
RowKey: []byte("row-present"),
Mutations: []*btpb.Mutation{{
Mutation: &btpb.Mutation_SetCell_{SetCell: &btpb.Mutation_SetCell{
FamilyName: "cf",
ColumnQualifier: []byte("col"),
TimestampMicros: 1000,
Value: val,
}},
}},
}
if _, err := s.MutateRow(ctx, mrreq); err != nil {
t.Fatalf("Populating table: %v", err)
}
req := &btpb.CheckAndMutateRowRequest{
TableName: tbl.Name,
RowKey: []byte("row-not-present"),
}
if res, err := s.CheckAndMutateRow(ctx, req); err != nil {
t.Errorf("CheckAndMutateRow error: %v", err)
} else if got, want := res.PredicateMatched, false; got != want {
t.Errorf("Invalid PredicateMatched value: got %t, want %t", got, want)
}
req = &btpb.CheckAndMutateRowRequest{
TableName: tbl.Name,
RowKey: []byte("row-present"),
}
if res, err := s.CheckAndMutateRow(ctx, req); err != nil {
t.Errorf("CheckAndMutateRow error: %v", err)
} else if got, want := res.PredicateMatched, true; got != want {
t.Errorf("Invalid PredicateMatched value: got %t, want %t", got, want)
}
}
func TestServer_ReadModifyWriteRow(t *testing.T) {
s := &server{
tables: make(map[string]*table),
}
ctx := context.Background()
newTbl := btapb.Table{
ColumnFamilies: map[string]*btapb.ColumnFamily{
"cf": {GcRule: &btapb.GcRule{Rule: &btapb.GcRule_MaxNumVersions{MaxNumVersions: 1}}},
},
}
tbl, err := s.CreateTable(ctx, &btapb.CreateTableRequest{Parent: "cluster", TableId: "t", Table: &newTbl})
if err != nil {
t.Fatalf("Creating table: %v", err)
}
req := &btpb.ReadModifyWriteRowRequest{
TableName: tbl.Name,
RowKey: []byte("row-key"),
Rules: []*btpb.ReadModifyWriteRule{
{
FamilyName: "cf",
ColumnQualifier: []byte("q1"),
Rule: &btpb.ReadModifyWriteRule_AppendValue{
AppendValue: []byte("a"),
},
},
// multiple ops for same cell
{
FamilyName: "cf",
ColumnQualifier: []byte("q1"),
Rule: &btpb.ReadModifyWriteRule_AppendValue{
AppendValue: []byte("b"),
},
},
// different cell whose qualifier should sort before the prior rules
{
FamilyName: "cf",
ColumnQualifier: []byte("q0"),
Rule: &btpb.ReadModifyWriteRule_IncrementAmount{
IncrementAmount: 1,
},
},
},
}
got, err := s.ReadModifyWriteRow(ctx, req)
if err != nil {
t.Fatalf("ReadModifyWriteRow error: %v", err)
}
want := &btpb.ReadModifyWriteRowResponse{
Row: &btpb.Row{
Key: []byte("row-key"),
Families: []*btpb.Family{{
Name: "cf",
Columns: []*btpb.Column{
{
Qualifier: []byte("q0"),
Cells: []*btpb.Cell{{
Value: []byte{0, 0, 0, 0, 0, 0, 0, 1},
}},
},
{
Qualifier: []byte("q1"),
Cells: []*btpb.Cell{{
Value: []byte("ab"),
}},
},
},
}},
},
}
diff := cmp.Diff(got, want, cmpopts.IgnoreFields(btpb.Cell{}, "TimestampMicros"))
if diff != "" {
t.Errorf("unexpected response: %s", diff)
}
}
// helper function to populate table data
func populateTable(ctx context.Context, s *server) (*btapb.Table, error) {
newTbl := btapb.Table{
ColumnFamilies: map[string]*btapb.ColumnFamily{
"cf0": {GcRule: &btapb.GcRule{Rule: &btapb.GcRule_MaxNumVersions{1}}},
},
}
tblInfo, err := s.CreateTable(ctx, &btapb.CreateTableRequest{Parent: "cluster", TableId: "t", Table: &newTbl})
if err != nil {
return nil, err
}
count := 3
mcf := func(i int) *btapb.ModifyColumnFamiliesRequest {
return &btapb.ModifyColumnFamiliesRequest{
Name: tblInfo.Name,
Modifications: []*btapb.ModifyColumnFamiliesRequest_Modification{{
Id: "cf" + strconv.Itoa(i),
Mod: &btapb.ModifyColumnFamiliesRequest_Modification_Create{&btapb.ColumnFamily{}},
}},
}
}
for i := 1; i <= count; i++ {
_, err = s.ModifyColumnFamilies(ctx, mcf(i))
if err != nil {
return nil, err
}
}
// Populate the table
for fc := 0; fc < count; fc++ {
for cc := count; cc > 0; cc-- {
for tc := 0; tc < count; tc++ {
req := &btpb.MutateRowRequest{
TableName: tblInfo.Name,
RowKey: []byte("row"),
Mutations: []*btpb.Mutation{{
Mutation: &btpb.Mutation_SetCell_{&btpb.Mutation_SetCell{
FamilyName: "cf" + strconv.Itoa(fc),
ColumnQualifier: []byte("col" + strconv.Itoa(cc)),
TimestampMicros: int64((tc + 1) * 1000),
Value: []byte{},
}},
}},
}
if _, err := s.MutateRow(ctx, req); err != nil {
return nil, err
}
}
}
}
return tblInfo, nil
}
func TestFilters(t *testing.T) {
tests := []struct {
in *btpb.RowFilter
out int
}{
{in: &btpb.RowFilter{Filter: &btpb.RowFilter_BlockAllFilter{true}}, out: 0},
{in: &btpb.RowFilter{Filter: &btpb.RowFilter_BlockAllFilter{false}}, out: 1},
{in: &btpb.RowFilter{Filter: &btpb.RowFilter_PassAllFilter{true}}, out: 1},
{in: &btpb.RowFilter{Filter: &btpb.RowFilter_PassAllFilter{false}}, out: 0},
}
ctx := context.Background()
s := &server{
tables: make(map[string]*table),
}
tblInfo, err := populateTable(ctx, s)
if err != nil {
t.Fatal(err)
}
req := &btpb.ReadRowsRequest{
TableName: tblInfo.Name,
Rows: &btpb.RowSet{RowKeys: [][]byte{[]byte("row")}},
}
for _, tc := range tests {
req.Filter = tc.in
mock := &MockReadRowsServer{}
if err = s.ReadRows(req, mock); err != nil {
t.Errorf("ReadRows error: %v", err)
continue
}
if len(mock.responses) != tc.out {
t.Errorf("Response count: got %d, want %d", len(mock.responses), tc.out)
continue
}
}
}
func Test_Mutation_DeleteFromColumn(t *testing.T) {
ctx := context.Background()
s := &server{
tables: make(map[string]*table),
}
tblInfo, err := populateTable(ctx, s)
if err != nil {
t.Fatal(err)
}
tests := []struct {
in *btpb.MutateRowRequest
fail bool
}{
{in: &btpb.MutateRowRequest{
TableName: tblInfo.Name,
RowKey: []byte("row"),
Mutations: []*btpb.Mutation{{
Mutation: &btpb.Mutation_DeleteFromColumn_{DeleteFromColumn: &btpb.Mutation_DeleteFromColumn{
FamilyName: "cf1",
ColumnQualifier: []byte("col1"),
TimeRange: &btpb.TimestampRange{
StartTimestampMicros: 2000,
EndTimestampMicros: 1000,
},
}},
}},
},
fail: true,
},
{in: &btpb.MutateRowRequest{
TableName: tblInfo.Name,
RowKey: []byte("row"),
Mutations: []*btpb.Mutation{{
Mutation: &btpb.Mutation_DeleteFromColumn_{DeleteFromColumn: &btpb.Mutation_DeleteFromColumn{
FamilyName: "cf2",
ColumnQualifier: []byte("col2"),
TimeRange: &btpb.TimestampRange{
StartTimestampMicros: 1000,
EndTimestampMicros: 2000,
},
}},
}},
},
fail: false,
},
{in: &btpb.MutateRowRequest{
TableName: tblInfo.Name,
RowKey: []byte("row"),
Mutations: []*btpb.Mutation{{
Mutation: &btpb.Mutation_DeleteFromColumn_{DeleteFromColumn: &btpb.Mutation_DeleteFromColumn{
FamilyName: "cf3",
ColumnQualifier: []byte("col3"),
TimeRange: &btpb.TimestampRange{
StartTimestampMicros: 1000,
EndTimestampMicros: 0,
},
}},
}},
},
fail: false,
},
{in: &btpb.MutateRowRequest{
TableName: tblInfo.Name,
RowKey: []byte("row"),
Mutations: []*btpb.Mutation{{
Mutation: &btpb.Mutation_DeleteFromColumn_{DeleteFromColumn: &btpb.Mutation_DeleteFromColumn{
FamilyName: "cf4",
ColumnQualifier: []byte("col4"),
TimeRange: &btpb.TimestampRange{
StartTimestampMicros: 0,
EndTimestampMicros: 1000,
},
}},
}},
},
fail: true,
},
}
for _, tst := range tests {
_, err = s.MutateRow(ctx, tst.in)
if err != nil && !tst.fail {
t.Errorf("expected passed got failure for : %v \n with err: %v", tst.in, err)
}
if err == nil && tst.fail {
t.Errorf("expected failure got passed for : %v", tst)
}
}
}
func TestFilterRow(t *testing.T) {
row := &row{
key: "row",
families: map[string]*family{
"fam": {
name: "fam",
cells: map[string][]cell{
"col": {{ts: 100, value: []byte("val")}},
},
},
},
}
for _, test := range []struct {
filter *btpb.RowFilter
want bool
}{
// The regexp-based filters perform whole-string, case-sensitive matches.
{&btpb.RowFilter{Filter: &btpb.RowFilter_RowKeyRegexFilter{[]byte("row")}}, true},
{&btpb.RowFilter{Filter: &btpb.RowFilter_RowKeyRegexFilter{[]byte("ro")}}, false},
{&btpb.RowFilter{Filter: &btpb.RowFilter_RowKeyRegexFilter{[]byte("ROW")}}, false},
{&btpb.RowFilter{Filter: &btpb.RowFilter_RowKeyRegexFilter{[]byte("moo")}}, false},
{&btpb.RowFilter{Filter: &btpb.RowFilter_FamilyNameRegexFilter{"fam"}}, true},
{&btpb.RowFilter{Filter: &btpb.RowFilter_FamilyNameRegexFilter{"f.*"}}, true},
{&btpb.RowFilter{Filter: &btpb.RowFilter_FamilyNameRegexFilter{"[fam]+"}}, true},
{&btpb.RowFilter{Filter: &btpb.RowFilter_FamilyNameRegexFilter{"fa"}}, false},
{&btpb.RowFilter{Filter: &btpb.RowFilter_FamilyNameRegexFilter{"FAM"}}, false},
{&btpb.RowFilter{Filter: &btpb.RowFilter_FamilyNameRegexFilter{"moo"}}, false},
{&btpb.RowFilter{Filter: &btpb.RowFilter_ColumnQualifierRegexFilter{[]byte("col")}}, true},
{&btpb.RowFilter{Filter: &btpb.RowFilter_ColumnQualifierRegexFilter{[]byte("co")}}, false},
{&btpb.RowFilter{Filter: &btpb.RowFilter_ColumnQualifierRegexFilter{[]byte("COL")}}, false},
{&btpb.RowFilter{Filter: &btpb.RowFilter_ColumnQualifierRegexFilter{[]byte("moo")}}, false},
{&btpb.RowFilter{Filter: &btpb.RowFilter_ValueRegexFilter{[]byte("val")}}, true},
{&btpb.RowFilter{Filter: &btpb.RowFilter_ValueRegexFilter{[]byte("va")}}, false},
{&btpb.RowFilter{Filter: &btpb.RowFilter_ValueRegexFilter{[]byte("VAL")}}, false},
{&btpb.RowFilter{Filter: &btpb.RowFilter_ValueRegexFilter{[]byte("moo")}}, false},
} {
got, _ := filterRow(test.filter, row.copy())
if got != test.want {
t.Errorf("%s: got %t, want %t", proto.CompactTextString(test.filter), got, test.want)
}
}
}
func TestFilterRowWithErrors(t *testing.T) {
row := &row{
key: "row",
families: map[string]*family{
"fam": {
name: "fam",
cells: map[string][]cell{
"col": {{ts: 100, value: []byte("val")}},
},
},
},
}
for _, test := range []struct {
badRegex *btpb.RowFilter
}{
{&btpb.RowFilter{Filter: &btpb.RowFilter_RowKeyRegexFilter{[]byte("[")}}},
{&btpb.RowFilter{Filter: &btpb.RowFilter_FamilyNameRegexFilter{"["}}},
{&btpb.RowFilter{Filter: &btpb.RowFilter_ColumnQualifierRegexFilter{[]byte("[")}}},
{&btpb.RowFilter{Filter: &btpb.RowFilter_ValueRegexFilter{[]byte("[")}}},
{&btpb.RowFilter{Filter: &btpb.RowFilter_Chain_{
Chain: &btpb.RowFilter_Chain{Filters: []*btpb.RowFilter{
{Filter: &btpb.RowFilter_ValueRegexFilter{[]byte("[")}}},
},
}}},
{&btpb.RowFilter{Filter: &btpb.RowFilter_Condition_{
Condition: &btpb.RowFilter_Condition{
PredicateFilter: &btpb.RowFilter{Filter: &btpb.RowFilter_ValueRegexFilter{[]byte("[")}},
},
}}},
{&btpb.RowFilter{Filter: &btpb.RowFilter_RowSampleFilter{0.0}}}, // 0.0 is invalid.
{&btpb.RowFilter{Filter: &btpb.RowFilter_RowSampleFilter{1.0}}}, // 1.0 is invalid.
} {
got, err := filterRow(test.badRegex, row.copy())
if got != false {
t.Errorf("%s: got true, want false", proto.CompactTextString(test.badRegex))
}
if err == nil {
t.Errorf("%s: got no error, want error", proto.CompactTextString(test.badRegex))
}
}
}
func TestFilterRowWithRowSampleFilter(t *testing.T) {
prev := randFloat
randFloat = func() float64 { return 0.5 }
defer func() { randFloat = prev }()
for _, test := range []struct {
p float64
want bool
}{
{0.1, false}, // Less than random float. Return no rows.
{0.5, false}, // Equal to random float. Return no rows.
{0.9, true}, // Greater than random float. Return all rows.
} {
got, err := filterRow(&btpb.RowFilter{Filter: &btpb.RowFilter_RowSampleFilter{test.p}}, &row{})
if err != nil {
t.Fatalf("%f: %v", test.p, err)
}
if got != test.want {
t.Errorf("%v: got %t, want %t", test.p, got, test.want)
}
}
}
func TestFilterRowWithBinaryColumnQualifier(t *testing.T) {
rs := []byte{128, 128}
row := &row{
key: string(rs),
families: map[string]*family{
"fam": {
name: "fam",
cells: map[string][]cell{
string(rs): {{ts: 100, value: []byte("val")}},
},
},
},
}
for _, test := range []struct {
filter []byte
want bool
}{
{[]byte{128, 128}, true}, // succeeds, exact match
{[]byte{128, 129}, false}, // fails
{[]byte{128}, false}, // fails, because the regexp must match the entire input
{[]byte{128, '*'}, true}, // succeeds: 0 or more 128s
{[]byte{'[', 127, 128, ']', '{', '2', '}'}, true}, // succeeds: exactly two of either 127 or 128
} {
got, _ := filterRow(&btpb.RowFilter{Filter: &btpb.RowFilter_ColumnQualifierRegexFilter{test.filter}}, row.copy())
if got != test.want {
t.Errorf("%v: got %t, want %t", test.filter, got, test.want)
}
}
}
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment