package pilosa import ( "bytes" "encoding/json" "fmt" "log" "strings" "time" "github.com/featurebasedb/featurebase/v3/dax" "github.com/featurebasedb/featurebase/v3/pql" "github.com/pkg/errors" ) // WireQueryResponse is the standard featurebase response type which can be // serialized and sent over the wire. type WireQueryResponse struct { Schema WireQuerySchema `json:"schema"` Data [][]interface{} `json:"data"` Error string `json:"error"` Warnings []string `json:"warnings"` QueryPlan map[string]interface{} `json:"query-plan"` ExecutionTime int64 `json:"execution-time"` } // WireQuerySchema is a list of Fields which map to the data columns in the // Response. type WireQuerySchema struct { Fields []*WireQueryField `json:"fields"` } // WireQueryField is a field name along with a supported BaseType and type // information. type WireQueryField struct { Name dax.FieldName `json:"name"` Type string `json:"type"` // human readable display (e.g. "decimal(2)") BaseType dax.BaseType `json:"base-type"` // for programmatic switching on type (e.g. "decimal") TypeInfo map[string]interface{} `json:"type-info"` // type modifiers (like scale), but not constraints (like min/max) } // UnmarshalJSON is a custom unmarshaller for the SQLResponse that converts the // value types in `Data` based on the types in `Schema`. func (s *WireQueryResponse) UnmarshalJSON(in []byte) error { return s.UnmarshalJSONTyped(in, false) } // UnmarshalJSONTyped is a temporary until we send typed values back in sql // responses. At that point, we can get rid of the typed=false path. In order to // do that, we need sql3 to return typed values, and we need the sql3/test/defs // to define results as typed values (like `IDSet`) instead of (for example) // `[]int64`. func (s *WireQueryResponse) UnmarshalJSONTyped(in []byte, typed bool) error { type Alias WireQueryResponse var aux Alias dec := json.NewDecoder(bytes.NewReader(in)) dec.UseNumber() err := dec.Decode(&aux) if err != nil { return err } *s = WireQueryResponse(aux) // If the SQLResponse contains an error, don't bother doing any conversions // on the data. if s.Error != "" { return nil } // Try to convert the types in the TypeInfo map for each field in the // schema. for _, fld := range s.Schema.Fields { // TODO(tlt): we can remove these two "ToLower" calls once sql3 is // returning dax.FieldType (i.e. lowercase). fld.Type = strings.ToLower(fld.Type) fld.BaseType = dax.BaseType(strings.ToLower(string(fld.BaseType))) for k, v := range fld.TypeInfo { switch k { case "scale": switch n := v.(type) { case float64: fld.TypeInfo[k] = int64(n) case json.Number: fld.TypeInfo[k], _ = n.Int64() } } } } // Try to convert the data types based on the headers. for i := range s.Data { for j, hdr := range s.Schema.Fields { switch hdr.BaseType { case dax.BaseTypeID, dax.BaseTypeInt: jn := s.Data[i][j] if v, ok := jn.(json.Number); ok { if x, err := v.Int64(); err == nil { s.Data[i][j] = x } else { return errors.Wrap(err, "can't be decoded as int64") } } case dax.BaseTypeIDSet: if src, ok := s.Data[i][j].([]interface{}); ok { if typed { val := make(IDSet, len(src)) for k := range src { v := src[k].(json.Number) if x, err := v.Int64(); err == nil { val[k] = x } else { return errors.Wrap(err, "can't be decoded as int64") } } s.Data[i][j] = val } else { val := make([]int64, len(src)) for k := range src { v := src[k].(json.Number) if x, err := v.Int64(); err == nil { val[k] = x } else { return errors.Wrap(err, "can't be decoded as int64") } } s.Data[i][j] = val } } case dax.BaseTypeDecimal: if jn, ok := s.Data[i][j].(json.Number); ok { var scale int64 if scaleVal, ok := hdr.TypeInfo["scale"]; !ok { return errors.New("decimal does not have a scale") } else if scaleInt64, ok := scaleVal.(int64); !ok { return errors.New("scale can't be cast to int64") } else { scale = scaleInt64 } format := fmt.Sprintf("%%.%df", scale) f, err := jn.Float64() if err != nil { return errors.Wrap(err, "parsing decimal") } dec, err := pql.ParseDecimal(fmt.Sprintf(format, f)) if err != nil { return errors.Wrap(err, "parsing decimal") } if dec.Scale != scale { dec = pql.NewDecimal(dec.ToInt64(scale), scale) } s.Data[i][j] = dec } case dax.BaseTypeStringSet: if src, ok := s.Data[i][j].([]interface{}); ok { if typed { val := make(StringSet, len(src)) for k := range src { val[k] = src[k].(string) } s.Data[i][j] = val } else { val := make([]string, len(src)) for k := range src { val[k] = src[k].(string) } s.Data[i][j] = val } } case dax.BaseTypeTimestamp: if src, ok := s.Data[i][j].(string); ok && src != "" { val, err := time.ParseInLocation(time.RFC3339Nano, src, time.UTC) if err != nil { return errors.Wrap(err, "parsing timestamp") } s.Data[i][j] = val } case dax.BaseTypeBool, dax.BaseTypeString: // no need to convert default: log.Printf("WARNING: unimplemented: %s", hdr.BaseType) } } } return nil } // IDSet is a return type specific to SQLResponse types. type IDSet []int64 func (ii IDSet) String() string { var sb strings.Builder sb.WriteString("[") for i := range ii { if i > 0 { sb.WriteString(", ") } sb.WriteString(fmt.Sprintf("%d", ii[i])) } sb.WriteString("]") return sb.String() } // StringSet is a return type specific to SQLResponse types. type StringSet []string func (ss StringSet) String() string { var sb strings.Builder sb.WriteString("[") for i := range ss { if i > 0 { sb.WriteString(", ") } sb.WriteString("'" + ss[i] + "'") } sb.WriteString("]") return sb.String() } // ShowColumnsResponse returns a structure which is specific to a `SHOW COLUMNS` // statement, derived from the results in the WireQueryResponse. This is kind of // a crude way to unmarshal a WireQueryResponse into a type which is specific to // the sql operation. // TODO(tlt): see if we can standardize on this logic, because it would be useful to // have the same thing for SHOW DATABASES and SHOW TABLES. // TODO(tlt): the fields unmarshalled in this method are a subset of the actual // columns available; at the moment, we only handled the ones we need in the // CLI. func (s *WireQueryResponse) ShowColumnsResponse() (*ShowColumnsResponse, error) { // Make a map of header names to index position. We do this to avoid // breaking things if for some reason the format of the SHOW COLUMNS // response defined in the sql3 package changes. m := make(map[string]int) for i, fld := range s.Schema.Fields { m[string(fld.Name)] = i } flds := make([]*dax.Field, 0, len(s.Data)) for _, row := range s.Data { // name var name dax.FieldName if val, ok := row[m["name"]].(string); ok { name = dax.FieldName(val) } // type var typ dax.BaseType if val, ok := row[m["type"]].(string); ok { typ = dax.BaseType(val) } // scale var scale int64 if val, ok := row[m["scale"]].(int64); ok { scale = val } flds = append(flds, &dax.Field{ Name: name, Type: typ, Options: dax.FieldOptions{ Scale: scale, }, }) } return &ShowColumnsResponse{ Fields: flds, }, nil } // ShowColumnsResponse is a type used to marshal the results of a `SHOW COLUMNS` // statement. type ShowColumnsResponse struct { Fields []*dax.Field } // Field returns the field by name. If the field is not found, nil is returned. func (s *ShowColumnsResponse) Field(name dax.FieldName) *dax.Field { for i := range s.Fields { if s.Fields[i].Name == name { return s.Fields[i] } } return nil }