From f52b88a96288d51346c1281e963f3ddf34b81cf0 Mon Sep 17 00:00:00 2001 From: Todd Gruben Date: Mon, 22 Mar 2021 08:49:54 -0500 Subject: [PATCH 01/17] etcd tls configuration support --- ctl/server.go | 7 ++++++- etcd/embed.go | 18 ++++++++++++++++++ server/config.go | 6 ++++++ 3 files changed, 30 insertions(+), 1 deletion(-) diff --git a/ctl/server.go b/ctl/server.go index e6469ffbd..65f854078 100644 --- a/ctl/server.go +++ b/ctl/server.go @@ -63,8 +63,13 @@ func BuildServerFlags(cmd *cobra.Command, srv *server.Command) { flags.StringVar(&srv.Config.Etcd.LPeerURL, "etcd.listen-peer-address", srv.Config.Etcd.LPeerURL, "Listen peer address.") flags.StringVar(&srv.Config.Etcd.APeerURL, "etcd.advertise-peer-address", srv.Config.Etcd.APeerURL, "Advertise peer address. If not provided, uses the listen peer address.") flags.StringVar(&srv.Config.Etcd.ClusterURL, "etcd.cluster-url", srv.Config.Etcd.ClusterURL, "Cluster URL to join.") - // Etcd.ClusterName uses Cluster.Name for its value. flags.StringVar(&srv.Config.Etcd.InitCluster, "etcd.initial-cluster", srv.Config.Etcd.InitCluster, "Initial cluster name1=apurl1,name2=apurl2") + // Etcd specified tls configuration options + flags.StringVar(&srv.Config.Etcd.TrustedCAFile, "etcd.tls-trusted-cafile", srv.Config.Etcd.TrustedCAFile, "Certificate Authority certificate") + flags.StringVar(&srv.Config.Etcd.ClientCertFile, "etcd.tls-cert-file", srv.Config.Etcd.ClientCertFile, "Client certificate (required for tls)") + flags.StringVar(&srv.Config.Etcd.ClientKeyFile, "etcd.tls-key-file", srv.Config.Etcd.ClientKeyFile, "Client key file (required for tls)") + flags.StringVar(&srv.Config.Etcd.PeerCertFile, "etcd.tls-peer-cert-file", srv.Config.Etcd.PeerCertFile, "Peer certificate (required for tls)") + flags.StringVar(&srv.Config.Etcd.PeerKeyFile, "etcd.tls-peer-key-file", srv.Config.Etcd.PeerKeyFile, "Peer key file (required for tls)") // AntiEntropy flags.DurationVar((*time.Duration)(&srv.Config.AntiEntropy.Interval), "anti-entropy.interval", (time.Duration)(srv.Config.AntiEntropy.Interval), "Interval at which to run anti-entropy routine.") diff --git a/etcd/embed.go b/etcd/embed.go index 0092ef42b..09afb7b22 100644 --- a/etcd/embed.go +++ b/etcd/embed.go @@ -34,6 +34,7 @@ import ( "go.etcd.io/etcd/embed" "go.etcd.io/etcd/etcdserver/api/v3client" "go.etcd.io/etcd/mvcc/mvccpb" + "go.etcd.io/etcd/pkg/transport" "go.etcd.io/etcd/pkg/types" ) @@ -48,6 +49,12 @@ type Options struct { InitCluster string `toml:"initial-cluster"` ClusterName string `toml:"cluster-name"` HeartbeatTTL int64 `toml:"heartbeat-ttl"` + // TLS provided tls files + TrustedCAFile string `toml:"tls-trusted-cafile"` + ClientCertFile string `toml:"tls-cert-file"` + ClientKeyFile string `toml:"tls-key-file"` + PeerCertFile string `toml:"tls-peer-cert-file"` + PeerKeyFile string `toml:"tls-peer-key-file"` LPeerSocket []*net.TCPListener LClientSocket []*net.TCPListener @@ -171,6 +178,17 @@ func parseOptions(opt Options) *embed.Config { id, name := memberAdd(cli, opt.APeerURL) log.Printf("\tid: %d, name: %s\n", id, name) } + // can only use tls if not using pre-configured listeners + cfg.ClientTLSInfo = transport.TLSInfo{ + TrustedCAFile: opt.TrustedCAFile, + CertFile: opt.ClientCertFile, + KeyFile: opt.ClientKeyFile, + } + cfg.PeerTLSInfo = transport.TLSInfo{ + TrustedCAFile: opt.TrustedCAFile, + CertFile: opt.PeerCertFile, + KeyFile: opt.PeerKeyFile, + } return cfg } diff --git a/server/config.go b/server/config.go index 1c004b7e8..edb69d5ec 100644 --- a/server/config.go +++ b/server/config.go @@ -351,6 +351,12 @@ func NewConfig() *Config { c.Etcd.InitCluster = c.Name + "=" + c.Etcd.LPeerURL c.Etcd.HeartbeatTTL = 5 + c.Etcd.TrustedCAFile = "" + c.Etcd.ClientCertFile = "" + c.Etcd.ClientKeyFile = "" + c.Etcd.PeerCertFile = "" + c.Etcd.PeerKeyFile = "" + return c } From 8131e807bfff1c1872be541b97360a4ea34581b5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Kuba=20Podg=C3=B3rski?= Date: Mon, 22 Mar 2021 15:42:02 +0100 Subject: [PATCH 02/17] Add client package with go-pilosa implementation --- Makefile | 4 +- client.go | 3 +- client/batch.go | 1404 ++++++++++++++++++++ client/batch_test.go | 1311 ++++++++++++++++++ client/client.go | 1811 +++++++++++++++++++++++++ client/client_internal_it_test.go | 85 ++ client/client_it_test.go | 2058 +++++++++++++++++++++++++++++ client/client_test.go | 280 ++++ client/cluster.go | 127 ++ client/cluster_test.go | 98 ++ client/csv/csv.go | 194 +++ client/csv/csv_it_test.go | 78 ++ client/csv/csv_test.go | 267 ++++ client/doc.go | 83 ++ client/egpool/egpool.go | 123 ++ client/egpool/egpool_test.go | 50 + client/error.go | 56 + client/logimport.go | 45 + client/logimport_test.go | 155 +++ client/metrics.go | 29 + client/orm.go | 1628 +++++++++++++++++++++++ client/orm_test.go | 1257 ++++++++++++++++++ client/record.go | 90 ++ client/record_test.go | 98 ++ client/response.go | 610 +++++++++ client/response_test.go | 363 +++++ client/shardnodes.go | 49 + client/tracer.go | 104 ++ client/validate.go | 69 + client/validate_test.go | 88 ++ client/version.go | 4 + go.mod | 33 +- go.sum | 127 +- handler.go | 2 +- net/uri.go | 50 +- net/uri_internal_test.go | 14 +- pilosa.go | 3 +- pilosa_test.go | 2 +- test/pilosa.go | 2 +- 39 files changed, 12751 insertions(+), 103 deletions(-) create mode 100644 client/batch.go create mode 100644 client/batch_test.go create mode 100644 client/client.go create mode 100644 client/client_internal_it_test.go create mode 100644 client/client_it_test.go create mode 100644 client/client_test.go create mode 100644 client/cluster.go create mode 100644 client/cluster_test.go create mode 100644 client/csv/csv.go create mode 100644 client/csv/csv_it_test.go create mode 100644 client/csv/csv_test.go create mode 100644 client/doc.go create mode 100644 client/egpool/egpool.go create mode 100644 client/egpool/egpool_test.go create mode 100644 client/error.go create mode 100644 client/logimport.go create mode 100644 client/logimport_test.go create mode 100644 client/metrics.go create mode 100644 client/orm.go create mode 100644 client/orm_test.go create mode 100644 client/record.go create mode 100644 client/record_test.go create mode 100644 client/response.go create mode 100644 client/response_test.go create mode 100644 client/shardnodes.go create mode 100644 client/tracer.go create mode 100644 client/validate.go create mode 100644 client/validate_test.go create mode 100644 client/version.go diff --git a/Makefile b/Makefile index f93ea4f52..966ffe755 100644 --- a/Makefile +++ b/Makefile @@ -63,7 +63,7 @@ testv-race: topt-race testvsub-race # find which test is hung/deadlocked. # testvsub: - set -e; for i in boltdb ctl http pg pql rbf roaring server sql txkey; do \ + set -e; for i in boltdb client ctl http pg pql rbf roaring server sql txkey; do \ echo; echo "___ testing subpkg $$i"; \ cd $$i; pwd; \ $(GO) test -tags='$(BUILD_TAGS) $(TEST_TAGS)' $(TESTFLAGS) -v -timeout 60m || break; \ @@ -72,7 +72,7 @@ testvsub: done testvsub-race: - set -e; for i in boltdb ctl http pg pql rbf roaring server sql txkey; do \ + set -e; for i in boltdb client ctl http pg pql rbf roaring server sql txkey; do \ echo; echo "___ testing subpkg $$i -race"; \ cd $$i; pwd; \ CGO_ENABLED=1 $(GO) test -tags='$(BUILD_TAGS) $(TEST_TAGS)' $(TESTFLAGS) -v -race -timeout 60m || break; \ diff --git a/client.go b/client.go index fd2bf45e0..7da57a9f7 100644 --- a/client.go +++ b/client.go @@ -47,7 +47,8 @@ type FieldValue struct { // something hasn't been architected correctly. // While I understand that putting the entire Client behind an interface might require this many methods, // I don't want to let it go unquestioned. -// Another note from Travis: I think we eventually want to unify `InternalClient` with the `go-pilosa` client. +// Another note from Travis: I think we eventually want to unify `InternalClient` with +// the `github.com/pilosa/pilosa/v2/client` client. // Doing that may obviate the need to refactor this. type InternalClient interface { InternalQueryClient diff --git a/client/batch.go b/client/batch.go new file mode 100644 index 000000000..753ee8b4d --- /dev/null +++ b/client/batch.go @@ -0,0 +1,1404 @@ +// Copyright 2017 Pilosa Corp. +// +// 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 client + +import ( + "time" + + "github.com/pilosa/pilosa/v2/client/egpool" + "github.com/pilosa/pilosa/v2/logger" + "github.com/pilosa/pilosa/v2/roaring" + "github.com/pkg/errors" +) + +// Batch defaults. +const ( + DefaultKeyTranslateBatchSize = 100000 +) + +// TODO if using column translation, column ids might get way out of +// order. Could be worth sorting everything after translation (as an +// option?). Instead of sorting all simultaneously, it might be faster +// (more cache friendly) to sort ids and save the swap ops to apply to +// everything else that needs to be sorted. + +// TODO support clearing values? nil values in records are ignored, +// but perhaps we could have a special type indicating that a bit or +// value should explicitly be cleared? + +// RecordBatch is a Pilosa ingest interface designed to allow for +// maximum throughput on common workloads. Users should call Add() +// with a Row object until it returns ErrBatchNowFull, at which time +// they should call Import(), and then repeat. +// +// Add will not modify or otherwise retain the Row once it returns, so +// it is recommended that callers reuse the same Row with repeated +// calls to Add, just modifying its values appropriately in between +// calls. This avoids allocating a new slice of Values for each +// inserted Row. +// +// The supported types of the values in Row.Values are implementation +// defined. Similarly, the supported types for Row.ID are +// implementation defined. +type RecordBatch interface { + Add(Row) error + + // Import does translation, creates the fragment files, and then, + // if we're not using split batch mode, imports everything to + // Pilosa. It then resets internal data structures for the next + // batch. If we are using split batch mode, it saves the fragment + // data to the batch, resets all other internal structures, and + // continues. + // Split batch mode DOES NOT CURRENTLY SUPPORT MUTEX OR INT FIELDS! + Import() error + + // Len reports the number of records which have been added to the + // batch since the last call to Import (or since it was created). + Len() int + + // Flush is only applicable in split batch mode where it actually + // imports the stored data to Pilosa. Otherwise it simply returns + // nil. + Flush() error +} + +// agedTranslation combines a translation with a recording of when it was last used. +type agedTranslation struct { + id uint64 + lastUsed uint64 +} + +// Batch implements RecordBatch. +// +// It supports Values of type string, uint64, int64, or nil. The +// following table describes what Pilosa field each type of value must +// map to. Fields are set up when calling "NewBatch". +// +// | type | pilosa field type | options | +// |--------+-------------------+-----------| +// | string | set | keys=true | +// | uint64 | set | any | +// | int64 | int | any | +// | float64| decimal | scale | +// | nil | any | | +// +// nil values are ignored. +type Batch struct { + client *Client + index *Index + header []*Field + headerMap map[string]*Field + + // prevDuration records the time that each doImport() takes. This + // is used to set the timeout for transactions to a reasonable + // value based on the last import. It starts with a conservative + // default set in NewBatch. + prevDuration time.Duration + + // ids is a slice of length batchSize of record IDs + ids []uint64 + + // rowIDs is a map of field index (in the header) to slices of + // length batchSize which contain row IDs. + rowIDs map[int][]uint64 + // clearRowIDs is a map[fieldIndex][idsIndex]rowID we don't expect + // clears to happen very often, so we store the idIndex/value + // mapping in a map rather than a slice as we do for rowIDs. This + // is a potentially temporary workaround to allow packed boolean + // fields to clear "false" values. Packed fields may be more + // completely supported by Pilosa in future. + clearRowIDs map[int]map[int]uint64 + + // rowIDSets is a map from field name to a batchSize slice of + // slices of row IDs. When a given record can have more than one + // value for a field, rowIDSets stores that information. + rowIDSets map[string][][]uint64 + + // values holds the values for each record of an int field + values map[string][]int64 + + // times holds a time for each record. (if any of the fields are time fields) + times []QuantizedTime + + // nullIndices holds a slice of indices into b.ids for each + // integer field which has nil values. + nullIndices map[string][]uint64 + + // TODO support bool fields. + + // for each field, keep a map of key to which record indexes that key mapped to + toTranslate map[int]map[string][]int + toTranslateClear map[int]map[string][]int + + // toTranslateSets is a map from field name to a map of string + // keys that need to be translated to sets of record indexes which + // those keys map to. + toTranslateSets map[string]map[string][]int + + // toTranslateID maps each string key to a record index - this + // will get translated into Batch.rowIDs + toTranslateID []string + + colTranslations map[string]agedTranslation + rowTranslations map[string]map[string]agedTranslation + cycle uint64 + maxAge uint64 + + // staleTime tracks the time the first record of the batch was inserted + // plus the maxStaleness, in order to raise ErrBatchNowStale if the + // maxStaleness has elapsed + staleTime time.Time + maxStaleness time.Duration + + // Maximum number of keys to translate at one time. + keyTranslateBatchSize int + + log logger.Logger + + // experimental — only used by FlushToFragments which is an + // alternative to Import which just builds the bitmap data for a + // batch without actually importing it. + splitBatchMode bool + frags fragments + clearFrags fragments +} + +func (b *Batch) Len() int { return len(b.ids) } + +// BatchOption is a functional option for Batch objects. +type BatchOption func(b *Batch) error + +func OptLogger(l logger.Logger) BatchOption { + return func(b *Batch) error { + b.log = l + return nil + } +} + +func OptSplitBatchMode(on bool) BatchOption { + return func(b *Batch) error { + b.splitBatchMode = on + return nil + } +} + +func OptCacheMaxAge(age uint64) BatchOption { + return func(b *Batch) error { + b.maxAge = age + return nil + } +} + +func OptMaxStaleness(t time.Duration) BatchOption { + return func(b *Batch) error { + b.maxStaleness = t + return nil + } +} + +func OptKeyTranslateBatchSize(v int) BatchOption { + return func(b *Batch) error { + b.keyTranslateBatchSize = v + return nil + } +} + +// NewBatch initializes a new Batch object which will use the given +// Pilosa client, index, set of fields, and will take "size" records +// before returning ErrBatchNowFull. The positions of the Fields in +// 'fields' correspond to the positions of values in the Row's Values +// passed to Batch.Add(). +func NewBatch(client *Client, size int, index *Index, fields []*Field, opts ...BatchOption) (*Batch, error) { + if len(fields) == 0 || size == 0 { + return nil, errors.New("can't batch with no fields or batch size") + } + headerMap := make(map[string]*Field, len(fields)) + rowIDs := make(map[int][]uint64, len(fields)) + values := make(map[string][]int64) + tt := make(map[int]map[string][]int, len(fields)) + ttSets := make(map[string]map[string][]int) + hasTime := false + for i, field := range fields { + headerMap[field.Name()] = field + opts := field.Opts() + switch typ := opts.Type(); typ { + case FieldTypeDefault, FieldTypeSet, FieldTypeTime: + if opts.Keys() { + tt[i] = make(map[string][]int) + ttSets[field.Name()] = make(map[string][]int) + } + hasTime = typ == FieldTypeTime || hasTime + case FieldTypeInt, FieldTypeDecimal: + // tt line only needed if int field is string foreign key + tt[i] = make(map[string][]int) + values[field.Name()] = make([]int64, 0, size) + case FieldTypeMutex: + // similar to set/time fields, but no need to support sets + // of values (hence no ttSets) + if opts.Keys() { + tt[i] = make(map[string][]int) + } + rowIDs[i] = make([]uint64, 0, size) + default: + return nil, errors.Errorf("field type '%s' is not currently supported through Batch", typ) + } + } + + b := &Batch{ + client: client, + header: fields, + headerMap: headerMap, + prevDuration: time.Minute * 11, + index: index, + ids: make([]uint64, 0, size), + rowIDs: rowIDs, + clearRowIDs: make(map[int]map[int]uint64), + rowIDSets: make(map[string][][]uint64), + values: values, + nullIndices: make(map[string][]uint64), + toTranslate: tt, + toTranslateClear: make(map[int]map[string][]int), + toTranslateSets: ttSets, + colTranslations: make(map[string]agedTranslation), + rowTranslations: make(map[string]map[string]agedTranslation), + maxAge: 64, + maxStaleness: time.Duration(0), + keyTranslateBatchSize: DefaultKeyTranslateBatchSize, + + log: logger.NopLogger, + + frags: make(fragments), + clearFrags: make(fragments), + } + if hasTime { + b.times = make([]QuantizedTime, 0, size) + } + for _, opt := range opts { + err := opt(b) + if err != nil { + return nil, errors.Wrap(err, "applying options") + } + } + return b, nil +} + +// Row represents a single record which can be added to a Batch. +type Row struct { + ID interface{} + // Values map to the slice of fields in Batch.header + Values []interface{} + // Clears' int key is an index into Batch.header + Clears map[int]interface{} + // Time applies to all time fields + Time QuantizedTime +} + +// QuantizedTime represents a moment in time down to some granularity +// (year, month, day, or hour). +type QuantizedTime struct { + ymdh [10]byte +} + +// Set sets the Quantized time to the given timestamp (down to hour +// granularity). +func (qt *QuantizedTime) Set(t time.Time) { + copy(qt.ymdh[:], t.Format("2006010215")) +} + +// SetYear sets the quantized time's year, but leaves month, day, and +// hour untouched. +func (qt *QuantizedTime) SetYear(year string) { + copy(qt.ymdh[:4], year) +} + +// SetMonth sets the QuantizedTime's month, but leaves year, day, and +// hour untouched. +func (qt *QuantizedTime) SetMonth(month string) { + copy(qt.ymdh[4:6], month) +} + +// SetDay sets the QuantizedTime's day, but leaves year, month, and +// hour untouched. +func (qt *QuantizedTime) SetDay(day string) { + copy(qt.ymdh[6:8], day) +} + +// SetHour sets the QuantizedTime's hour, but leaves year, month, and +// day untouched. +func (qt *QuantizedTime) SetHour(hour string) { + copy(qt.ymdh[8:10], hour) +} + +// Reset sets the time to the zero value which generates no time views. +func (qt *QuantizedTime) Reset() { + for i := range qt.ymdh { + qt.ymdh[i] = 0 + } +} + +// views builds the list of Pilosa views for this particular time, +// given a quantum. +func (qt *QuantizedTime) views(q TimeQuantum) ([]string, error) { + zero := QuantizedTime{} + if *qt == zero { + return nil, nil + } + views := make([]string, 0, len(q)) + for _, unit := range q { + switch unit { + case 'Y': + if qt.ymdh[0] == 0 { + return nil, errors.New("no data set for year") + } + views = append(views, string(qt.ymdh[:4])) + case 'M': + if qt.ymdh[4] == 0 { + return nil, errors.New("no data set for month") + } + views = append(views, string(qt.ymdh[:6])) + case 'D': + if qt.ymdh[6] == 0 { + return nil, errors.New("no data set for day") + } + views = append(views, string(qt.ymdh[:8])) + case 'H': + if qt.ymdh[8] == 0 { + return nil, errors.New("no data set for hour") + } + views = append(views, string(qt.ymdh[:10])) + } + } + return views, nil +} + +func (b *Batch) getColTranslation(key string) (uint64, bool) { + trans, ok := b.colTranslations[key] + if ok { + trans.lastUsed = b.cycle + b.colTranslations[key] = trans + } + return trans.id, ok +} + +func (b *Batch) getRowTranslation(field, key string) (uint64, bool) { + trans, ok := b.rowTranslations[field][key] + if ok { + trans.lastUsed = b.cycle + b.rowTranslations[field][key] = trans + } + return trans.id, ok +} + +func (b *Batch) addRowTranslations(fieldName string, keys []string, ids []uint64) { + rowCache := b.rowTranslations[fieldName] + if rowCache == nil { + rowCache = make(map[string]agedTranslation) + b.rowTranslations[fieldName] = rowCache + } + for i, k := range keys { + rowCache[k] = agedTranslation{ + id: ids[i], + lastUsed: b.cycle, + } + } +} + +func (b *Batch) addColTranslations(keys []string, ids []uint64) { + for i, k := range keys { + b.colTranslations[k] = agedTranslation{ + id: ids[i], + lastUsed: b.cycle, + } + } +} + +// Add adds a record to the batch. Performance will be best if record +// IDs are shard-sorted. That is, all records which belong to the same +// Pilosa shard are added adjacent to each other. If the records are +// also in-order within a shard this will likely help as well. Add +// clears rec.Clears when it returns normally (either a nil error or +// BatchNowFull). +func (b *Batch) Add(rec Row) error { + // Clear recValues and rec.Clears upon return. + defer func() { + for i := range rec.Values { + rec.Values[i] = nil + } + for k := range rec.Clears { + delete(rec.Clears, k) + } + }() + + if len(b.ids) == cap(b.ids) { + return ErrBatchAlreadyFull + } + if len(rec.Values) != len(b.header) { + return errors.Errorf("record needs to match up with batch fields, got %d fields and %d record", len(b.header), len(rec.Values)) + } + + handleStringID := func(rid string) error { + if rid == "" { + return errors.Errorf("record identifier cannot be an empty string") + } + if colID, ok := b.getColTranslation(rid); ok { + b.ids = append(b.ids, colID) + } else { + if b.toTranslateID == nil { + b.toTranslateID = make([]string, cap(b.ids)) + } + b.toTranslateID[len(b.ids)] = rid + b.ids = append(b.ids, 0) + } + return nil + } + var err error + + switch rid := rec.ID.(type) { + case uint64: + b.ids = append(b.ids, rid) + case string: + err := handleStringID(rid) + if err != nil { + return err + } + case []byte: + err = handleStringID(string(rid)) + if err != nil { + return err + } + default: // TODO support nil ID as being auto-allocated. + return errors.Errorf("unsupported id type %T value %v", rid, rid) + } + + // curPos is the current position in b.ids, rowIDs[*], etc. + curPos := len(b.ids) - 1 + + if b.times != nil { + b.times = append(b.times, rec.Time) + } + + for i := 0; i < len(rec.Values); i++ { + field := b.header[i] + switch val := rec.Values[i].(type) { + case string: + if field.Opts().Type() != FieldTypeInt { + // nil-extend + for len(b.rowIDs[i]) < curPos { + b.rowIDs[i] = append(b.rowIDs[i], nilSentinel) + } + rowIDs := b.rowIDs[i] + // empty string is not a valid value at this point (Pilosa refuses to translate it) + if val == "" { // + b.rowIDs[i] = append(rowIDs, nilSentinel) + + } else if rowID, ok := b.getRowTranslation(field.Name(), val); ok { + b.rowIDs[i] = append(rowIDs, rowID) + } else { + ints, ok := b.toTranslate[i][val] + if !ok { + ints = make([]int, 0) + } + ints = append(ints, curPos) + b.toTranslate[i][val] = ints + b.rowIDs[i] = append(rowIDs, 0) + } + } else if field.Opts().Type() == FieldTypeInt { + if val == "" { + // copied from the `case nil:` section for ints and decimals + b.values[field.Name()] = append(b.values[field.Name()], 0) + nullIndices, ok := b.nullIndices[field.Name()] + if !ok { + nullIndices = make([]uint64, 0) + } + nullIndices = append(nullIndices, uint64(curPos)) + b.nullIndices[field.Name()] = nullIndices + } else if intVal, ok := b.getRowTranslation(field.Name(), val); ok { + b.values[field.Name()] = append(b.values[field.Name()], int64(intVal)) + } else { + ints, ok := b.toTranslate[i][val] + if !ok { + ints = make([]int, 0) + } + ints = append(ints, curPos) + b.toTranslate[i][val] = ints + b.values[field.Name()] = append(b.values[field.Name()], 0) + } + } + case uint64: + // nil-extend + for len(b.rowIDs[i]) < curPos { + b.rowIDs[i] = append(b.rowIDs[i], nilSentinel) + } + b.rowIDs[i] = append(b.rowIDs[i], val) + case int64: + b.values[field.Name()] = append(b.values[field.Name()], val) + case []string: + if len(val) == 0 { + continue + } + rowIDSets, ok := b.rowIDSets[field.Name()] + if !ok { + rowIDSets = make([][]uint64, len(b.ids)-1, cap(b.ids)) + b.rowIDSets[field.Name()] = rowIDSets + } + for len(rowIDSets) < len(b.ids)-1 { + rowIDSets = append(rowIDSets, nil) // nil extend + } + + rowIDs := make([]uint64, 0, len(val)) + for _, k := range val { + if k == "" { + continue + } + if rowID, ok := b.getRowTranslation(field.Name(), k); ok { + rowIDs = append(rowIDs, rowID) + } else { + ttsets, ok := b.toTranslateSets[field.Name()] + if !ok { + ttsets = make(map[string][]int) + b.toTranslateSets[field.Name()] = make(map[string][]int) + } + ints, ok := ttsets[k] + if !ok { + ints = make([]int, 0, 1) + } + ints = append(ints, curPos) + b.toTranslateSets[field.Name()][k] = ints + } + } + b.rowIDSets[field.Name()] = append(rowIDSets, rowIDs) + case []uint64: + if len(val) == 0 { + continue + } + rowIDSets, ok := b.rowIDSets[field.Name()] + if !ok { + rowIDSets = make([][]uint64, len(b.ids)-1, cap(b.ids)) + } + for len(rowIDSets) < len(b.ids)-1 { + rowIDSets = append(rowIDSets, nil) // nil extend + } + b.rowIDSets[field.Name()] = append(rowIDSets, val) + case nil: + if field.Opts().Type() == FieldTypeInt || field.Opts().Type() == FieldTypeDecimal { + b.values[field.Name()] = append(b.values[field.Name()], 0) + nullIndices, ok := b.nullIndices[field.Name()] + if !ok { + nullIndices = make([]uint64, 0) + } + nullIndices = append(nullIndices, uint64(curPos)) + b.nullIndices[field.Name()] = nullIndices + + } else { + // only append nil to rowIDs if this field already has + // rowIDs. Otherwise, this could be a []string or + // []uint64 field where we've only seen nil values so + // far. when we see a uint64 or string value, we'll + // "nil-extend" rowIDs to make sure it's the right + // length. + if rowIDs, ok := b.rowIDs[i]; ok { + b.rowIDs[i] = append(rowIDs, nilSentinel) + } + } + default: + return errors.Errorf("Val %v Type %[1]T is not currently supported. Use string, uint64 (row id), or int64 (integer value)", val) + } + } + + for i, uval := range rec.Clears { + field := b.header[i] + if _, ok := b.clearRowIDs[i]; !ok { + b.clearRowIDs[i] = make(map[int]uint64) + } + switch val := uval.(type) { + case string: + clearRows := b.clearRowIDs[i] + // translate val and add to clearRows + if rowID, ok := b.getRowTranslation(field.Name(), val); ok { + clearRows[curPos] = rowID + } else { + _, ok := b.toTranslateClear[i] + if !ok { + b.toTranslateClear[i] = make(map[string][]int) + } + ints, ok := b.toTranslateClear[i][val] + if !ok { + ints = make([]int, 0) + } + ints = append(ints, curPos) + b.toTranslateClear[i][val] = ints + } + case uint64: + b.clearRowIDs[i][curPos] = val + default: + return errors.Errorf("Clearing a value '%v' Type %[1]T is not currently supported (field '%s')", val, field.Name()) + } + // nil extend b.rowIDs so we don't run into a horrible bug + // where we skip doing clears because b.rowIDs doesn't have a + // value for this field + for len(b.rowIDs[i]) <= curPos { + b.rowIDs[i] = append(b.rowIDs[i], nilSentinel) + } + + } + + if len(b.ids) == cap(b.ids) { + return ErrBatchNowFull + } + if b.maxStaleness != time.Duration(0) { // set maxStaleness to 0 to disable staleness checking + if len(b.ids) == 1 { + b.staleTime = time.Now().Add(b.maxStaleness) + } else if time.Now().After(b.staleTime) { + return ErrBatchNowStale + } + } + return nil +} + +// ErrBatchNowFull, similar to io.EOF, is a marker error to notify the +// user of a batch that it is time to call Import. +var ErrBatchNowFull = errors.New("batch is now full - you cannot add any more records (though the one you just added was accepted)") + +// ErrBatchAlreadyFull is a real error saying that Batch.Add did not +// complete because the batch was full. +var ErrBatchAlreadyFull = errors.New("batch was already full, record was rejected") + +// ErrBatchNowStale indicates that the oldest record in the batch is older than +// the maxStaleness value of the batch. Like ErrBatchNowFull, the error does +// not mean the record was rejected. +var ErrBatchNowStale = errors.New("batch is stale and needs to be imported (however, record was accepted)") + +// Import does translation, creates the fragment files, and then, +// if we're not using split batch mode, imports everything to +// Pilosa. It then resets internal data structures for the next +// batch. If we are using split batch mode, it saves the fragment +// data to the batch, resets all other internal structures, and +// continues. split batch mode DOES NOT CURRENTLY SUPPORT MUTEX +// OR INT FIELDS! +func (b *Batch) Import() error { + start := time.Now() + trns, err := b.client.StartTransaction("", b.prevDuration*10, false, time.Hour) + if err != nil { + return errors.Wrap(err, "starting transaction") + } + defer func() { + trns, err := b.client.FinishTransaction(trns.ID) + if err != nil { + b.log.Printf("error finishing transaction: %v. trns: %+v", err, trns) + } + b.client.Stats.Timing(MetricBatchImportDurationSeconds, time.Since(start), 1.0) + }() + + size := len(b.ids) + transStart := time.Now() + // first we need to translate the toTranslate, then fill out the missing row IDs + err = b.doTranslation() + if err != nil { + return errors.Wrap(err, "doing Translation") + } + transTime := time.Now() + b.log.Printf("translating batch of %d took: %v", size, transTime.Sub(transStart)) + + frags, clearFrags, err := b.makeFragments(b.frags, b.clearFrags) + if err != nil { + return errors.Wrap(err, "making fragments (flush)") + } + makeTime := time.Now() + b.log.Printf("making fragments for batch of %d took %v", size, makeTime.Sub(transTime)) + + if b.splitBatchMode { + b.frags = frags + b.clearFrags = clearFrags + } else { + b.frags = make(fragments) + b.clearFrags = make(fragments) + // create bitmaps out of each field in b.rowIDs and import. Also + // import int data. + err = b.doImport(frags, clearFrags) + if err != nil { + return errors.Wrap(err, "doing import") + } + b.log.Printf("importing fragments took %v", time.Since(makeTime)) + } + + b.reset() + return nil +} + +// Flush is only applicable in split batch mode where it actually +// imports the stored data to Pilosa. Otherwise it simply returns +// nil. +func (b *Batch) Flush() error { + if !b.splitBatchMode { + return nil + } + start := time.Now() + + trns, err := b.client.StartTransaction("", b.prevDuration*10, false, time.Hour) + if err != nil { + return errors.Wrap(err, "starting transaction") + } + defer func() { + trns, err := b.client.FinishTransaction(trns.ID) + if err != nil { + b.log.Printf("error finishing transaction: %v. trns: %+v", err, trns) + } + b.client.Stats.Timing(MetricBatchFlushDurationSeconds, time.Since(start), 1.0) + }() + + importStart := time.Now() + err = b.doImport(b.frags, b.clearFrags) + if err != nil { + return errors.Wrap(err, "doing import (ImportFragments)") + } + + b.log.Debugf("superbatch import took %v", time.Since(importStart)) + + b.reset() + b.frags = make(fragments) + b.clearFrags = make(fragments) + return nil +} + +func (b *Batch) doTranslation() error { + keys := make([]string, 0) + + // translate column keys if there are any + + // TODO test. Also this implementation (using a set to de-dup + // keys) will likely have much worse performance than the previous + // one (two slices, one of keys one of ids) in the case that most + // of the keys are unique. + keySet := make(map[string]uint64) + for _, key := range b.toTranslateID { + if key != "" { + if _, ok := keySet[key]; ok { + continue + } + keys = append(keys, key) + keySet[key] = 0 + } + } + + if len(keys) > 0 { + start := time.Now() + ids, err := b.translateColumnKeys(b.index, keys) + if err != nil { + return errors.Wrap(err, "translating col keys") + } + if len(ids) != len(keys) { + return errors.Errorf("requested IDs for %d column keys but got %d back", len(keys), len(ids)) + } + b.log.Debugf("translating %d column keys took %v", len(keys), time.Since(start)) + b.addColTranslations(keys, ids) + for j, id := range ids { + keySet[keys[j]] = id + } + for index, ttkey := range b.toTranslateID { + if ttkey != "" { + b.ids[index] = keySet[ttkey] + } + } + } + // translate row keys + for i, tt := range b.toTranslate { + fieldName := b.header[i].Name() + keys = keys[:0] + + // make a slice of keys + for k := range tt { + keys = append(keys, k) + } + // append keys to clear so we can translate them all in one + // request. ttEnd is the index where clearing starts which we + // use later on. + ttEnd := len(keys) + ttc := b.toTranslateClear[i] + for k := range ttc { + keys = append(keys, k) + } + + if len(keys) == 0 { + continue + } + + // translate keys from Pilosa + start := time.Now() + ids, err := b.translateRowKeys(b.headerMap[fieldName], keys) + if err != nil { + return errors.Wrap(err, "translating row keys") + } + if len(ids) != len(keys) { + return errors.Errorf("requested IDs for %d row keys but got %d back", len(keys), len(ids)) + } + b.log.Debugf("translating %d row keys for %s took %v", len(keys), fieldName, time.Since(start)) + b.addRowTranslations(fieldName, keys, ids) + + switch b.header[i].Opts().Type() { + case FieldTypeInt: + // handle foreign key int fields — fill out b.values instead of b.rows + for j := 0; j < ttEnd; j++ { + key := keys[j] + id := ids[j] + for _, recordIdx := range tt[key] { + b.values[fieldName][recordIdx] = int64(id) + } + } + case FieldTypeDecimal: + return errors.Errorf("unexpected field type for translation: decimal") + default: + // fill out missing IDs in local batch records with translated IDs + rows := b.rowIDs[i] + for j := 0; j < ttEnd; j++ { + key := keys[j] + id := ids[j] + for _, recordIdx := range tt[key] { + rows[recordIdx] = id + } + } + // fill out missing IDs in clear lists. + clearRows := b.clearRowIDs[i] + for j := ttEnd; j < len(keys); j++ { + key := keys[j] + id := ids[j] + for _, recordIdx := range ttc[key] { + clearRows[recordIdx] = id + } + } + } + } + + for fieldName, tt := range b.toTranslateSets { + keys = keys[:0] + + for k := range tt { + keys = append(keys, k) + } + + if len(keys) == 0 { + continue + } + // translate keys from Pilosa + start := time.Now() + ids, err := b.translateRowKeys(b.headerMap[fieldName], keys) + if err != nil { + return errors.Wrap(err, "translating row keys (sets)") + } + if len(ids) != len(keys) { + return errors.Errorf("requested IDs for %d row (set) keys but got %d back", len(keys), len(ids)) + } + b.log.Debugf("translating %d row keys(sets) for %s took %v", len(keys), fieldName, time.Since(start)) + b.addRowTranslations(fieldName, keys, ids) + rowIDSets := b.rowIDSets[fieldName] + rowIDSets = rowIDSets[:cap(b.ids)] + b.rowIDSets[fieldName] = rowIDSets + for j, key := range keys { + rowID := ids[j] + for _, recordIdx := range tt[key] { + rowIDSets[recordIdx] = append(rowIDSets[recordIdx], rowID) + } + } + } + + return nil +} + +func (b *Batch) translateColumnKeys(index *Index, keys []string) ([]uint64, error) { + batchSize := b.keyTranslateBatchSize + if batchSize <= 0 { + batchSize = len(keys) + } + + ids := make([]uint64, 0, len(keys)) + for i := 0; i < len(keys); i += batchSize { + keySlice := keys[i:] + if len(keySlice) > batchSize { + keySlice = keySlice[:batchSize] + } + + idSlice, err := b.client.TranslateColumnKeys(b.index, keySlice) + if err != nil { + return nil, err + } else if len(idSlice) != len(keySlice) { + return nil, errors.Errorf("requested IDs slice for %d column keys but got %d back", len(keySlice), len(idSlice)) + } + ids = append(ids, idSlice...) + } + + return ids, nil +} + +func (b *Batch) translateRowKeys(field *Field, keys []string) ([]uint64, error) { + batchSize := b.keyTranslateBatchSize + if batchSize <= 0 { + batchSize = len(keys) + } + + ids := make([]uint64, 0, len(keys)) + for i := 0; i < len(keys); i += batchSize { + keySlice := keys[i:] + if len(keySlice) > batchSize { + keySlice = keySlice[:batchSize] + } + + idSlice, err := b.client.TranslateRowKeys(field, keySlice) + if err != nil { + return nil, err + } else if len(idSlice) != len(keySlice) { + return nil, errors.Errorf("requested IDs slice for %d row keys but got %d back", len(keySlice), len(idSlice)) + } + ids = append(ids, idSlice...) + } + + return ids, nil +} + +func (b *Batch) doImport(frags, clearFrags fragments) error { + + start := time.Now() + eg := egpool.Group{PoolSize: 20} + // TODO, currently this relies on upstream behavior of + // makeFragments to guarantee that any shard/field combination in + // clearFrags also has a shard/field in frags. We're only + // iterating over frags and then checking to see if clearFrags has + // the same keys. If we optimized makeFragments to skip adding + // things to frags which had no set bits (e.g. if we were only + // clearing things), then this code would need to be updated to + // ensure that it looked at the things in clearFrags which were + // *not* in frags. + for fragmentKey, viewMap := range frags { + field := fragmentKey.field + shard := fragmentKey.shard + viewMap := viewMap + + eg.Go(func() error { + clearViewMap := clearFrags.GetViewMap(shard, field) + if len(clearViewMap) > 0 { + start := time.Now() + err := b.client.ImportRoaringBitmap(b.index.Field(field), shard, clearViewMap, true) + if err != nil { + return errors.Wrapf(err, "import clearing clearing data for %s", field) + } + b.log.Debugf("imp-roar-clr %s,shard:%d,views:%d %v", field, shard, len(clearViewMap), time.Since(start)) + } + + start := time.Now() + err := b.client.ImportRoaringBitmap(b.index.Field(field), shard, viewMap, false) + b.log.Debugf("imp-roar %s,shard:%d,views:%d %v", field, shard, len(clearViewMap), time.Since(start)) + return errors.Wrapf(err, "importing data for %s", field) + }) + } + eg.Go(func() error { return b.importValueData() }) + eg.Go(func() error { return b.importMutexData() }) + + err := eg.Wait() + if err != nil { + if pferr := anyCause(ErrPreconditionFailed, eg.Errors()...); pferr != nil { + return pferr + } + return err + } + b.prevDuration = time.Since(start) + return nil +} + +func anyCause(cause error, errs ...error) error { + if cause == nil { + return nil + } + + for _, err := range errs { + if errors.Cause(err) == cause { + return err + } + } + return nil +} + +// this is kind of bad as it means we can never import column id +// ^uint64(0) which is a valid column ID. I think it's unlikely to +// matter much in practice (we could maybe special case it somewhere +// if needed though). +var nilSentinel = ^uint64(0) + +func (b *Batch) makeFragments(frags, clearFrags fragments) (fragments, fragments, error) { + shardWidth := b.index.ShardWidth() + if shardWidth == 0 { + shardWidth = DefaultShardWidth + } + emptyClearRows := make(map[int]uint64) + + // create _exists fragments if needed + if b.index.Opts().TrackExistence() { + var curBM *roaring.Bitmap + curShard := ^uint64(0) // impossible sentinel value for shard. + for _, col := range b.ids { + if col/shardWidth != curShard { + curShard = col / shardWidth + curBM = frags.GetOrCreate(curShard, "_exists", "") + } + curBM.DirectAdd(col % shardWidth) + } + } + + for i, rowIDs := range b.rowIDs { + if len(rowIDs) == 0 { + continue // this can happen when the values that came in for this field were string slices + } + clearRows := b.clearRowIDs[i] + if clearRows == nil { + clearRows = emptyClearRows + } + field := b.header[i] + opts := field.Opts() + if opts.Type() == FieldTypeMutex { + continue // we handle mutex fields separately — they can't use importRoaring + } + curShard := ^uint64(0) // impossible sentinel value for shard. + var curBM *roaring.Bitmap + var clearBM *roaring.Bitmap + for j := range b.ids { + col, row := b.ids[j], rowIDs[j] + if col/shardWidth != curShard { + curShard = col / shardWidth + curBM = frags.GetOrCreate(curShard, field.Name(), "") + clearBM = clearFrags.GetOrCreate(curShard, field.Name(), "") + } + if row != nilSentinel { + // TODO this is super ugly, but we want to avoid setting + // bits on the standard view in the specific case when + // there isn't one. Should probably refactor this whole + // loop to be more general w.r.t. views. Also... tests for + // the NoStandardView case would be great. + if !(opts.Type() == FieldTypeTime && opts.NoStandardView()) { + curBM.DirectAdd(row*shardWidth + (col % shardWidth)) + } + if opts.Type() == FieldTypeTime { + views, err := b.times[j].views(opts.TimeQuantum()) + if err != nil { + return nil, nil, errors.Wrap(err, "calculating views") + } + for _, view := range views { + tbm := frags.GetOrCreate(curShard, field.Name(), view) + tbm.DirectAdd(row*shardWidth + (col % shardWidth)) + } + } + } + + clearRow, ok := clearRows[j] + if ok { + clearBM.DirectAddN(clearRow*shardWidth + (col % shardWidth)) + // we're going to execute the clear before the set, so + // we want to make sure that at this point, the "set" + // fragments don't contain the bit that we're clearing + curBM.DirectRemoveN(clearRow*shardWidth + (col % shardWidth)) + } + } + } + + for fname, rowIDSets := range b.rowIDSets { + if len(rowIDSets) == 0 { + continue + } else if len(rowIDSets) < len(b.ids) { + // rowIDSets is guaranteed to have capacity == to b.ids, + // but if the last record had a nil for this field, it + // might not have the same length, so we re-slice it to + // ensure the lengths are the same. + rowIDSets = rowIDSets[:len(b.ids)] + } + field := b.headerMap[fname] + opts := field.Opts() + curShard := ^uint64(0) // impossible sentinel value for shard. + var curBM *roaring.Bitmap + for j := range b.ids { + col, rowIDs := b.ids[j], rowIDSets[j] + if len(rowIDs) == 0 { + continue + } + if col/shardWidth != curShard { + curShard = col / shardWidth + curBM = frags.GetOrCreate(curShard, fname, "") + } + // TODO this is super ugly, but we want to avoid setting + // bits on the standard view in the specific case when + // there isn't one. Should probably refactor this whole + // loop to be more general w.r.t. views. Also... tests for + // the NoStandardView case would be great. + if !(opts.Type() == FieldTypeTime && opts.NoStandardView()) { + for _, row := range rowIDs { + curBM.DirectAdd(row*shardWidth + (col % shardWidth)) + } + } + if opts.Type() == FieldTypeTime { + views, err := b.times[j].views(opts.TimeQuantum()) + if err != nil { + return nil, nil, errors.Wrap(err, "calculating views") + } + for _, view := range views { + tbm := frags.GetOrCreate(curShard, fname, view) + for _, row := range rowIDs { + tbm.DirectAdd(row*shardWidth + (col % shardWidth)) + } + } + } + } + } + return frags, clearFrags, nil +} + +// importValueData imports data for int fields. +func (b *Batch) importValueData() error { + shardWidth := b.index.ShardWidth() + if shardWidth == 0 { + shardWidth = DefaultShardWidth + } + eg := egpool.Group{PoolSize: 20} + + ids := make([]uint64, len(b.ids)) + for fieldName, bvalues := range b.values { + ids = ids[:len(b.ids)] + + // trim out null values from ids and values. + nullIndices := b.nullIndices[fieldName] + + i, n := uint64(0), 0 + for _, nullIndex := range nullIndices { + copy(ids[n:], b.ids[i:nullIndex]) + n += copy(bvalues[n:], bvalues[i:nullIndex]) + i = nullIndex + 1 + } + + copy(ids[n:], b.ids[i:]) + n += copy(bvalues[n:], bvalues[i:]) + ids, bvalues = ids[:n], bvalues[:n] + + // now do imports by shard + if len(ids) == 0 { + continue // TODO test this "all nil" case + } + curShard := ids[0] / shardWidth + startIdx := 0 + for i := 1; i <= len(ids); i++ { + var recordID uint64 + if i < len(ids) { + recordID = ids[i] + } else { + recordID = (curShard + 2) * shardWidth + } + + if recordID/shardWidth != curShard { + endIdx := i + shard := curShard + field := b.headerMap[fieldName] + path, data, err := b.client.EncodeImportValues(field, shard, bvalues[startIdx:endIdx], ids[startIdx:endIdx], false) + if err != nil { + return errors.Wrap(err, "encoding import values") + } + eg.Go(func() error { + start := time.Now() + err := b.client.DoImportValues(b.index.Name(), shard, path, data) + b.log.Debugf("imp-vals %s,shard:%d,data:%d %v", field, shard, len(data), time.Since(start)) + return errors.Wrapf(err, "importing values for %s", field) + }) + startIdx = i + curShard = recordID / shardWidth + } + } + } + err := eg.Wait() + if err != nil { + if pferr := anyCause(ErrPreconditionFailed, eg.Errors()...); pferr != nil { + return pferr + } + return err + } + return errors.Wrap(err, "importing value data") +} + +// TODO this should work for bools as well - just need to support them +// at batch creation time and when calling Add, I think. +func (b *Batch) importMutexData() error { + shardWidth := b.index.ShardWidth() + if shardWidth == 0 { + shardWidth = DefaultShardWidth + } + + eg := egpool.Group{PoolSize: 20} + ids := make([]uint64, 0, len(b.ids)) + for findex, rowIDs := range b.rowIDs { + field := b.header[findex] + if field.Opts().Type() != FieldTypeMutex { + continue + } + ids = ids[:0] + + // get slice of column ids for non-nil rowIDs and cut nil row + // IDs out of rowIDs. + idsIndex := 0 + for i, id := range b.ids { + rowID := rowIDs[i] + if rowID == nilSentinel { + continue + } + rowIDs[idsIndex] = rowID + ids = append(ids, id) + idsIndex++ + } + rowIDs = rowIDs[:idsIndex] + + if len(ids) == 0 { + continue + } + curShard := ids[0] / shardWidth + startIdx := 0 + for i := 1; i <= len(ids); i++ { + var recordID uint64 + if i < len(ids) { + recordID = ids[i] + } else { + recordID = (curShard + 2) * shardWidth + } + + if recordID/shardWidth != curShard { + endIdx := i + shard := curShard + field := field + path, data, err := b.client.EncodeImport(field, shard, rowIDs[startIdx:endIdx], ids[startIdx:endIdx], false) + if err != nil { + return errors.Wrap(err, "encoding mutex import") + } + eg.Go(func() error { + start := time.Now() + err := b.client.DoImport(b.index.Name(), shard, path, data) + b.log.Debugf("imp-mux %s,shard:%d,data:%d %v", field.Name(), shard, len(data), time.Since(start)) + return errors.Wrapf(err, "importing values for %s", field) + }) + startIdx = i + curShard = recordID / shardWidth + } + } + } + err := eg.Wait() + if err != nil { + if pferr := anyCause(ErrPreconditionFailed, eg.Errors()...); pferr != nil { + return pferr + } + return err + } + return errors.Wrap(err, "importing mutex data") +} + +// reset is called at the end of importing to ready the batch for the +// next round. Where possible it does not re-allocate memory. +func (b *Batch) reset() { + b.ids = b.ids[:0] + b.times = b.times[:0] + for i, rowIDs := range b.rowIDs { + b.rowIDs[i] = rowIDs[:0] + } + for _, tt := range b.toTranslate { + for k := range tt { + delete(tt, k) // TODO pool these slices + } + } + for _, tts := range b.toTranslateSets { + for k := range tts { + delete(tts, k) + } + } + for field, rowIDSet := range b.rowIDSets { + for i := range rowIDSet { + rowIDSet[i] = nil + } + b.rowIDSets[field] = rowIDSet[:0] + } + for _, rowIDs := range b.clearRowIDs { + for k := range rowIDs { + delete(rowIDs, k) + } + } + for _, clearMap := range b.toTranslateClear { + for k := range clearMap { + delete(clearMap, k) + } + } + for i := range b.toTranslateID { + b.toTranslateID[i] = "" + } + for k := range b.values { + delete(b.values, k) // TODO pool these slices + } + for k := range b.nullIndices { + delete(b.nullIndices, k) // TODO pool these slices + } + b.cycle++ + for k, trans := range b.colTranslations { + if trans.lastUsed-b.cycle > b.maxAge { + delete(b.colTranslations, k) + } + } + for field, rowTranslations := range b.rowTranslations { + for k, trans := range rowTranslations { + if trans.lastUsed-b.cycle > b.maxAge { + delete(rowTranslations, k) + } + } + + if len(rowTranslations) == 0 { + delete(b.rowTranslations, field) + } + } +} + +// map[shard][field][view]fragmentData +type fragments map[fragmentKey]map[string]*roaring.Bitmap + +type fragmentKey struct { + shard uint64 + field string +} + +func (f fragments) GetOrCreate(shard uint64, field, view string) *roaring.Bitmap { + key := fragmentKey{shard, field} + viewMap, ok := f[key] + if !ok { + viewMap = make(map[string]*roaring.Bitmap) + f[key] = viewMap + } + bm, ok := viewMap[view] + if !ok { + bm = roaring.NewBTreeBitmap() + viewMap[view] = bm + } + return bm +} + +func (f fragments) GetViewMap(shard uint64, field string) map[string]*roaring.Bitmap { + key := fragmentKey{shard, field} + viewMap, ok := f[key] + if !ok { + return nil + } + // Remove any views which have an empty bitmap. + // TODO: Ideally we would prevent allocating the empty bitmap to begin with, + // but the logic is a bit tricky, and since we don't want to spend too much + // time on it right now, we're leaving that for a future exercise. + for k, v := range viewMap { + if v.Count() == 0 { + delete(viewMap, k) + } + } + return viewMap +} diff --git a/client/batch_test.go b/client/batch_test.go new file mode 100644 index 000000000..8255562eb --- /dev/null +++ b/client/batch_test.go @@ -0,0 +1,1311 @@ +//+build integration + +// Copyright 2017 Pilosa Corp. +// +// 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 client + +import ( + "reflect" + "sort" + "strconv" + "testing" + "time" + + "github.com/pkg/errors" +) + +func TestStringSliceCombos(t *testing.T) { + client := DefaultClient() + schema := NewSchema() + idx := schema.Index("test-string-slicecombos") + fields := make([]*Field, 1) + fields[0] = idx.Field("a1", OptFieldKeys(true), OptFieldTypeSet(CacheTypeRanked, 100)) + err := client.SyncSchema(schema) + if err != nil { + t.Fatalf("syncing schema: %v", err) + } + defer func() { + err := client.DeleteIndex(idx) + if err != nil { + t.Logf("problem cleaning up from test: %v", err) + } + }() + + b, err := NewBatch(client, 5, idx, fields) + if err != nil { + t.Fatalf("creating new batch: %v", err) + } + + records := []Row{ + {ID: uint64(0), Values: []interface{}{[]string{"a", "b", "c"}}}, + {ID: uint64(1), Values: []interface{}{[]string{"z"}}}, + {ID: uint64(2), Values: []interface{}{[]string{}}}, + {ID: uint64(3), Values: []interface{}{[]string{"q", "r", "s", "t", "c"}}}, + {ID: uint64(4), Values: []interface{}{nil}}, + {ID: uint64(5), Values: []interface{}{[]string{"a", "b", "c"}}}, + {ID: uint64(6), Values: []interface{}{[]string{"a", "b", "c"}}}, + {ID: uint64(7), Values: []interface{}{[]string{"z"}}}, + {ID: uint64(8), Values: []interface{}{[]string{}}}, + {ID: uint64(9), Values: []interface{}{[]string{"q", "r", "s", "t"}}}, + {ID: uint64(10), Values: []interface{}{nil}}, + {ID: uint64(11), Values: []interface{}{[]string{"a", "b", "c"}}}, + {ID: uint64(12), Values: []interface{}{[]string{}}}, + {ID: uint64(13), Values: []interface{}{[]string{}}}, + } + + err = ingestRecords(records, b) + if err != nil { + t.Fatalf("importing: %v", err) + } + + a1 := fields[0] + + result := tq(t, client, a1.TopN(10)) + rez := sortableCRI(result.CountItems()) + sort.Sort(rez) + exp := sortableCRI{ + {Key: "a", Count: 4}, + {Key: "b", Count: 4}, + {Key: "c", Count: 5}, + {Key: "q", Count: 2}, + {Key: "r", Count: 2}, + {Key: "s", Count: 2}, + {Key: "t", Count: 2}, + {Key: "z", Count: 2}, + } + sort.Sort(exp) + errorIfNotEqual(t, exp, rez) + + result = tq(t, client, a1.Row("a")) + errorIfNotEqual(t, result.Row().Columns, []uint64{0, 5, 6, 11}) + result = tq(t, client, a1.Row("b")) + errorIfNotEqual(t, result.Row().Columns, []uint64{0, 5, 6, 11}) + result = tq(t, client, a1.Row("c")) + errorIfNotEqual(t, result.Row().Columns, []uint64{0, 3, 5, 6, 11}) + result = tq(t, client, a1.Row("z")) + errorIfNotEqual(t, result.Row().Columns, []uint64{1, 7}) + result = tq(t, client, a1.Row("q")) + errorIfNotEqual(t, result.Row().Columns, []uint64{3, 9}) + result = tq(t, client, a1.Row("r")) + errorIfNotEqual(t, result.Row().Columns, []uint64{3, 9}) + result = tq(t, client, a1.Row("s")) + errorIfNotEqual(t, result.Row().Columns, []uint64{3, 9}) + result = tq(t, client, a1.Row("t")) + errorIfNotEqual(t, result.Row().Columns, []uint64{3, 9}) + + result = tq(t, client, idx.RawQuery("Count(All())")) + errorIfNotEqual(t, result.Count(), int64(14)) +} + +func errorIfNotEqual(t *testing.T, exp, got interface{}) { + t.Helper() + if !reflect.DeepEqual(exp, got) { + t.Errorf("unequal exp/got:\n%v\n%v", exp, got) + } +} + +type sortableCRI []CountResultItem + +func (s sortableCRI) Len() int { return len(s) } +func (s sortableCRI) Less(i, j int) bool { + if s[i].Count != s[j].Count { + return s[i].Count > s[j].Count + } + if s[i].ID != s[j].ID { + return s[i].ID < s[j].ID + } + if s[i].Key != s[j].Key { + return s[i].Key < s[j].Key + } + return true +} +func (s sortableCRI) Swap(i, j int) { + s[i], s[j] = s[j], s[i] +} + +func tq(t *testing.T, client *Client, query PQLQuery) QueryResult { + resp, err := client.Query(query) + if err != nil { + t.Fatalf("querying: %v", err) + } + return resp.Results()[0] +} + +func ingestRecords(records []Row, batch *Batch) error { + for _, rec := range records { + err := batch.Add(rec) + if err == ErrBatchNowFull { + err = batch.Import() + if err != nil { + return errors.Wrap(err, "importing batch") + } + } else if err != nil { + return errors.Wrap(err, "while adding record") + } + } + if batch.Len() > 0 { + err := batch.Import() + if err != nil { + return errors.Wrap(err, "importing batch") + } + } + return nil +} + +func TestImportBatchInts(t *testing.T) { + client := DefaultClient() + schema := NewSchema() + idx := schema.Index("gopilosatest-blah") + field := idx.Field("anint", OptFieldTypeInt()) + err := client.SyncSchema(schema) + if err != nil { + t.Fatalf("syncing schema: %v", err) + } + + b, err := NewBatch(client, 3, idx, []*Field{field}) + if err != nil { + t.Fatalf("getting batch: %v", err) + } + + r := Row{Values: make([]interface{}, 1)} + + for i := uint64(0); i < 3; i++ { + r.ID = i + r.Values[0] = int64(i) + err := b.Add(r) + if err != nil && err != ErrBatchNowFull { + t.Fatalf("adding to batch: %v", err) + } + } + err = b.Import() + if err != nil { + t.Fatalf("importing: %v", err) + } + + r.ID = uint64(0) + r.Values[0] = nil + err = b.Add(r) + if err != nil { + t.Fatalf("adding after import: %v", err) + } + r.ID = uint64(1) + r.Values[0] = int64(7) + err = b.Add(r) + if err != nil { + t.Fatalf("adding second after import: %v", err) + } + + err = b.Import() + if err != nil { + t.Fatalf("second import: %v", err) + } + + resp, err := client.Query(idx.BatchQuery(field.Equals(0), field.Equals(7), field.Equals(2))) + if err != nil { + t.Fatalf("querying: %v", err) + } + + for i, result := range resp.Results() { + if !reflect.DeepEqual(result.Row().Columns, []uint64{uint64(i)}) { + t.Errorf("expected %v for %d, but got %v", []uint64{uint64(i)}, i, result.Row().Columns) + } + } +} + +func TestTrimNull(t *testing.T) { + client := DefaultClient() + schema := NewSchema() + idx := schema.Index("gopilosatest-null") + field := idx.Field("empty", OptFieldTypeInt()) + err := client.SyncSchema(schema) + if err != nil { + t.Fatalf("syncing schema: %v", err) + } + defer func() { + err := client.DeleteIndex(idx) + if err != nil { + t.Logf("problem cleaning up from test: %v", err) + } + }() + b, err := NewBatch(client, 3, idx, []*Field{field}) + if err != nil { + t.Fatalf("getting batch: %v", err) + } + b.nullIndices = make(map[string][]uint64, 1) + b.nullIndices[field.Name()] = []uint64{0, 1, 2} + r := Row{Values: make([]interface{}, 1)} + for i := 0; i < 3; i++ { + r.ID = uint64(i) + r.Values[0] = int64(i) + err := b.Add(r) + if err != nil && err != ErrBatchNowFull { + t.Fatalf("adding to batch: %v", err) + } + } + err = b.Import() + if err != nil { + t.Fatalf("importing: %v", err) + } + resp, err := client.Query(idx.BatchQuery(field.Equals(0), field.Equals(1), field.Equals(2))) + if err != nil { + t.Fatalf("querying: %v", err) + } + for i, result := range resp.Results() { + if !reflect.DeepEqual(result.Row().Columns, []uint64(nil)) { + t.Errorf("expected %#v for %d, but got %#v", []uint64(nil), i, result.Row().Columns) + } + } + + b, err = NewBatch(client, 4, idx, []*Field{field}) + if err != nil { + t.Fatalf("getting batch: %v", err) + } + r = Row{Values: make([]interface{}, 1)} + for i := 10; i < 40; i += 10 { + r.ID = uint64(i) + r.Values[0] = int64(i) + err := b.Add(r) + if err != nil && err != ErrBatchNowFull { + t.Fatalf("adding to batch: %v", err) + } + } + + r.ID = uint64(40) + r.Values[0] = nil + err = b.Add(r) + if err != nil && err != ErrBatchNowFull { + t.Fatalf("adding to batch: %v", err) + } + err = b.Import() + if err != nil { + t.Fatalf("importing: %v", err) + } + + resp, err = client.Query(idx.BatchQuery(field.Equals(10), field.Equals(40), field.Equals(20), field.Equals(30))) + if err != nil { + t.Fatalf("querying: %v", err) + } + for i, result := range resp.Results() { + if 1 == i { + if !reflect.DeepEqual(result.Row().Columns, []uint64(nil)) { + t.Errorf("expected %#v for %d, but got %#v", []uint64(nil), i, result.Row().Columns) + } + } else { + if !reflect.DeepEqual(result.Row().Columns, []uint64{result.Row().Columns[0]}) { + t.Errorf("expected %#v for %d, but got %#v", []uint64{result.Row().Columns[0]}, i, result.Row().Columns) + } + } + } + +} + +func TestStringSliceEmptyAndNil(t *testing.T) { + client := DefaultClient() + schema := NewSchema() + idx := schema.Index("test-string-slice-nil") + fields := make([]*Field, 1) + fields[0] = idx.Field("strslice", OptFieldKeys(true), OptFieldTypeSet(CacheTypeRanked, 100)) + err := client.SyncSchema(schema) + if err != nil { + t.Fatalf("syncing schema: %v", err) + } + defer func() { + err := client.DeleteIndex(idx) + if err != nil { + t.Logf("problem cleaning up from test: %v", err) + } + }() + + // first create a batch and test adding a single value with empty + // string - this failed with a translation error at one point, and + // how we catch it and treat it like a nil. + b, err := NewBatch(client, 2, idx, fields) + if err != nil { + t.Fatalf("creating new batch: %v", err) + } + r := Row{Values: make([]interface{}, len(fields))} + r.ID = uint64(1) + r.Values[0] = "" + err = b.Add(r) + if err != nil { + t.Fatalf("adding: %v", err) + } + err = b.Import() + if err != nil { + t.Fatalf("importing: %v", err) + } + + // now create a batch and add a mixture of string slice values + b, err = NewBatch(client, 6, idx, fields) + if err != nil { + t.Fatalf("creating new batch: %v", err) + } + r = Row{Values: make([]interface{}, len(fields))} + r.ID = uint64(0) + r.Values[0] = []string{"a"} + err = b.Add(r) + if err != nil { + t.Fatalf("adding to batch: %v", err) + } + + r.ID = uint64(1) + r.Values[0] = nil + err = b.Add(r) + if err != nil { + t.Fatalf("adding batch with nil stringslice to r: %v", err) + } + + r.ID = uint64(2) + r.Values[0] = []string{"a", "b", "z"} + err = b.Add(r) + if err != nil { + t.Fatalf("adding batch with idslice to r: %v", err) + } + + r.ID = uint64(3) + r.Values[0] = []string{"b", "c"} + err = b.Add(r) + if err != nil { + t.Fatalf("adding batch with stringslice to r: %v", err) + } + + r.ID = uint64(4) + r.Values[0] = []string{} + err = b.Add(r) + if err != nil { + t.Fatalf("adding batch with stringslice to r: %v", err) + } + + err = b.Import() + if err != nil { + t.Fatalf("importing: %v", err) + } + + rows := []interface{}{"a", "b", "c", "z"} + resp, err := client.Query(idx.BatchQuery(fields[0].Row(rows[0]), fields[0].Row(rows[1]), fields[0].Row(rows[2]), fields[0].Row(rows[3]))) + if err != nil { + t.Fatalf("querying: %v", err) + } + + // TODO test is flaky because we can't guarantee what a,b,c map to + expectations := [][]uint64{{0, 2}, {2, 3}, {3}, {2}} + for i, re := range resp.Results() { + if !reflect.DeepEqual(re.Row().Columns, expectations[i]) { + t.Errorf("expected row %v to have columns %v, but got %v", rows[i], expectations[i], re.Row().Columns) + } + } + +} + +func TestStringSlice(t *testing.T) { + client := DefaultClient() + schema := NewSchema() + idx := schema.Index("test-string-slice") + fields := make([]*Field, 1) + fields[0] = idx.Field("strslice", OptFieldKeys(true), OptFieldTypeSet(CacheTypeRanked, 100)) + err := client.SyncSchema(schema) + if err != nil { + t.Fatalf("syncing schema: %v", err) + } + defer func() { + err := client.DeleteIndex(idx) + if err != nil { + t.Logf("problem cleaning up from test: %v", err) + } + }() + + b, err := NewBatch(client, 3, idx, fields) + if err != nil { + t.Fatalf("creating new batch: %v", err) + } + + rowmap := map[string]uint64{ + "c": 9, + "d": 10, + "f": 13, + } + b.rowTranslations["strslice"] = make(map[string]agedTranslation) + for k, id := range rowmap { + b.rowTranslations["strslice"][k] = agedTranslation{ + id: id, + } + } + + r := Row{Values: make([]interface{}, len(fields))} + r.ID = uint64(0) + r.Values[0] = []string{"a"} + err = b.Add(r) + if err != nil { + t.Fatalf("adding to batch: %v", err) + } + if got := b.toTranslateSets["strslice"]["a"]; !reflect.DeepEqual(got, []int{0}) { + t.Fatalf("expected []int{0}, got: %v", got) + } + + r.ID = uint64(1) + r.Values[0] = []string{"a", "b", "c"} + err = b.Add(r) + if err != nil { + t.Fatalf("adding to batch: %v", err) + } + if got := b.toTranslateSets["strslice"]["a"]; !reflect.DeepEqual(got, []int{0, 1}) { + t.Fatalf("expected []int{0,1}, got: %v", got) + } + if got := b.toTranslateSets["strslice"]["b"]; !reflect.DeepEqual(got, []int{1}) { + t.Fatalf("expected []int{1}, got: %v", got) + } + if got, ok := b.toTranslateSets["strslice"]["c"]; ok { + t.Fatalf("should be nothing at c, got: %v", got) + } + if got := b.rowIDSets["strslice"][1]; !reflect.DeepEqual(got, []uint64{9}) { + t.Fatalf("expected c to map to rowID 9 but got %v", got) + } + + r.ID = uint64(2) + r.Values[0] = []string{"d", "e", "f"} + err = b.Add(r) + if err != ErrBatchNowFull { + t.Fatalf("adding to batch: %v", err) + } + if got, ok := b.toTranslateSets["strslice"]["d"]; ok { + t.Fatalf("should be nothing at d, got: %v", got) + } + if got, ok := b.toTranslateSets["strslice"]["f"]; ok { + t.Fatalf("should be nothing at f, got: %v", got) + } + if got := b.toTranslateSets["strslice"]["e"]; !reflect.DeepEqual(got, []int{2}) { + t.Fatalf("expected []int{2}, got: %v", got) + } + if got := b.rowIDSets["strslice"][2]; !reflect.DeepEqual(got, []uint64{10, 13}) { + t.Fatalf("expected c to map to rowID 9 but got %v", got) + } + + err = b.doTranslation() + if err != nil { + t.Fatalf("translating: %v", err) + } + + if got0 := b.rowIDSets["strslice"][0]; len(got0) != 1 { + t.Errorf("after translation, rec 0, wrong len: %v", got0) + } else if got1 := b.rowIDSets["strslice"][1]; len(got1) != 3 || got1[0] != 9 || (got1[1] != got0[0] && got1[2] != got0[0]) { + t.Errorf("after translation, rec 1: %v, rec 0: %v", got1, got0) + } else if got2 := b.rowIDSets["strslice"][2]; len(got2) != 3 || got2[0] != 10 || got2[1] != 13 || got2[2] == got1[2] || got2[2] == got0[0] { + t.Errorf("after translation, rec 2: %v", got2) + } + + frags, clearFrags, err := b.makeFragments(make(fragments), make(fragments)) + if err != nil { + t.Errorf("making fragments: %v", err) + } + + err = b.doImport(frags, clearFrags) + if err != nil { + t.Fatalf("doing import: %v", err) + } + + resp, err := client.Query(idx.BatchQuery(fields[0].Row("a"))) + if err != nil { + t.Fatalf("querying: %v", err) + } + result := resp.Result() + if !reflect.DeepEqual(result.Row().Columns, []uint64{0, 1}) { + t.Fatalf("expected a to be [0,1], got %v", result.Row().Columns) + } +} + +func TestSingleClearBatchRegression(t *testing.T) { + client := DefaultClient() + schema := NewSchema() + idx := schema.Index("gopilosatest-blah") + numFields := 1 + fields := make([]*Field, numFields) + fields[0] = idx.Field("zero", OptFieldKeys(true)) + + err := client.SyncSchema(schema) + if err != nil { + t.Fatalf("syncing schema: %v", err) + } + defer func() { + err := client.DeleteIndex(idx) + if err != nil { + t.Logf("problem cleaning up from test: %v", err) + } + }() + + _, err = client.Query(fields[0].Set("row1", 1)) + if err != nil { + t.Fatalf("setting bit: %v", err) + } + + b, err := NewBatch(client, 1, idx, fields) + if err != nil { + t.Fatalf("getting new batch: %v", err) + } + r := Row{ID: uint64(1), Values: make([]interface{}, numFields), Clears: make(map[int]interface{})} + r.Values[0] = nil + r.Clears[0] = "row1" + err = b.Add(r) + if err != ErrBatchNowFull { + t.Fatalf("wrong error from batch add: %v", err) + } + + err = b.Import() + if err != nil { + t.Fatalf("error importing: %v", err) + } + + resp, err := client.Query(fields[0].Row("row1")) + if err != nil { + t.Fatalf("error querying: %v", err) + } + result := resp.Results()[0].Row().Columns + if len(result) != 0 { + t.Fatalf("unexpected values in row: result %+v", result) + } + +} + +func TestBatches(t *testing.T) { + client := DefaultClient() + schema := NewSchema() + idx := schema.Index("gopilosatest-blah") + numFields := 5 + fields := make([]*Field, numFields) + fields[0] = idx.Field("zero", OptFieldKeys(true)) + fields[1] = idx.Field("one", OptFieldKeys(true)) + fields[2] = idx.Field("two", OptFieldKeys(true)) + fields[3] = idx.Field("three", OptFieldTypeInt()) + fields[4] = idx.Field("four", OptFieldTypeTime(TimeQuantumYearMonthDay)) + err := client.SyncSchema(schema) + if err != nil { + t.Fatalf("syncing schema: %v", err) + } + defer func() { + err := client.DeleteIndex(idx) + if err != nil { + t.Logf("problem cleaning up from test: %v", err) + } + }() + b, err := NewBatch(client, 10, idx, fields) + if err != nil { + t.Fatalf("getting new batch: %v", err) + } + r := Row{Values: make([]interface{}, numFields), Clears: make(map[int]interface{})} + r.Time.Set(time.Date(2019, time.January, 2, 15, 45, 0, 0, time.UTC)) + + for i := 0; i < 9; i++ { + r.ID = uint64(i) + if i%2 == 0 { + r.Values[0] = "a" + r.Values[1] = "b" + r.Values[2] = "c" + r.Values[3] = int64(99) + r.Values[4] = uint64(1) + r.Time.SetMonth("01") + } else { + r.Values[0] = "x" + r.Values[1] = "y" + r.Values[2] = "z" + r.Values[3] = int64(-10) + r.Values[4] = uint64(1) + r.Time.SetMonth("02") + } + if i == 8 { + r.Values[0] = nil + r.Clears[1] = uint64(97) + r.Clears[2] = "c" + r.Values[3] = nil + r.Values[4] = nil + } + err := b.Add(r) + if err != nil { + t.Fatalf("unexpected err adding record: %v", err) + } + + } + + if len(b.toTranslate[0]) != 2 { + t.Fatalf("wrong number of keys in toTranslate[0]") + } + for k, ints := range b.toTranslate[0] { + if k == "a" { + if !reflect.DeepEqual(ints, []int{0, 2, 4, 6}) { + t.Fatalf("wrong ints for key a in field zero: %v", ints) + } + } else if k == "x" { + if !reflect.DeepEqual(ints, []int{1, 3, 5, 7}) { + t.Fatalf("wrong ints for key x in field zero: %v", ints) + } + + } else { + t.Fatalf("unexpected key %s", k) + } + } + if !reflect.DeepEqual(b.toTranslateClear, map[int]map[string][]int{2: {"c": {8}}}) { + t.Errorf("unexpected toTranslateClear: %+v", b.toTranslateClear) + } + if !reflect.DeepEqual(b.clearRowIDs, map[int]map[int]uint64{1: {8: 97}, 2: {}}) { + t.Errorf("unexpected clearRowIDs: %+v", b.clearRowIDs) + } + + if !reflect.DeepEqual(b.values["three"], []int64{99, -10, 99, -10, 99, -10, 99, -10, 0}) { + t.Fatalf("unexpected values: %v", b.values["three"]) + } + if !reflect.DeepEqual(b.nullIndices["three"], []uint64{8}) { + t.Fatalf("unexpected nullIndices: %v", b.nullIndices["three"]) + } + + if len(b.toTranslate[1]) != 2 { + t.Fatalf("wrong number of keys in toTranslate[1]") + } + for k, ints := range b.toTranslate[1] { + if k == "b" { + if !reflect.DeepEqual(ints, []int{0, 2, 4, 6, 8}) { + t.Fatalf("wrong ints for key b in field one: %v", ints) + } + } else if k == "y" { + if !reflect.DeepEqual(ints, []int{1, 3, 5, 7}) { + t.Fatalf("wrong ints for key y in field one: %v", ints) + } + + } else { + t.Fatalf("unexpected key %s", k) + } + } + + if len(b.toTranslate[2]) != 2 { + t.Fatalf("wrong number of keys in toTranslate[2]") + } + for k, ints := range b.toTranslate[2] { + if k == "c" { + if !reflect.DeepEqual(ints, []int{0, 2, 4, 6, 8}) { + t.Fatalf("wrong ints for key c in field two: %v", ints) + } + } else if k == "z" { + if !reflect.DeepEqual(ints, []int{1, 3, 5, 7}) { + t.Fatalf("wrong ints for key z in field two: %v", ints) + } + + } else { + t.Fatalf("unexpected key %s", k) + } + } + + err = b.Add(r) + if err != ErrBatchNowFull { + t.Fatalf("should have gotten full batch error, but got %v", err) + } + + err = b.Add(r) + if err != ErrBatchAlreadyFull { + t.Fatalf("should have gotten already full batch error, but got %v", err) + } + + if !reflect.DeepEqual(b.values["three"], []int64{99, -10, 99, -10, 99, -10, 99, -10, 0, 0}) { + t.Fatalf("unexpected values: %v", b.values["three"]) + } + + err = b.doTranslation() + if err != nil { + t.Fatalf("doing translation: %v", err) + } + + for fidx, rowIDs := range b.rowIDs { + // we don't know which key will get translated first, but we do know the pattern + if fidx == 0 { + if !reflect.DeepEqual(rowIDs, []uint64{1, 2, 1, 2, 1, 2, 1, 2, nilSentinel, nilSentinel}) && + !reflect.DeepEqual(rowIDs, []uint64{2, 1, 2, 1, 2, 1, 2, 1, nilSentinel, nilSentinel}) { + t.Fatalf("unexpected row ids for field %d: %v", fidx, rowIDs) + } + + } else if fidx == 4 { + if !reflect.DeepEqual(rowIDs, []uint64{1, 1, 1, 1, 1, 1, 1, 1, nilSentinel, nilSentinel}) { + t.Fatalf("unexpected rowids for time field") + } + } else if fidx == 3 { + if len(rowIDs) != 0 { + t.Fatalf("expected no rowIDs for int field, but got: %v", rowIDs) + } + } else { + if !reflect.DeepEqual(rowIDs, []uint64{1, 2, 1, 2, 1, 2, 1, 2, 1, nilSentinel}) && !reflect.DeepEqual(rowIDs, []uint64{2, 1, 2, 1, 2, 1, 2, 1, 2, nilSentinel}) { + t.Fatalf("unexpected row ids for field %d: %v", fidx, rowIDs) + } + } + } + + if !reflect.DeepEqual(b.clearRowIDs[1], map[int]uint64{8: 97}) { + t.Errorf("unexpected clearRowIDs after translation: %+v", b.clearRowIDs[1]) + } + if !reflect.DeepEqual(b.clearRowIDs[2], map[int]uint64{8: 2}) && !reflect.DeepEqual(b.clearRowIDs[2], map[int]uint64{8: 1}) { + t.Errorf("unexpected clearRowIDs: after translation%+v", b.clearRowIDs[2]) + } + + frags, clearFrags, err := b.makeFragments(make(fragments), make(fragments)) + if err != nil { + t.Errorf("making fragments: %v", err) + } + + err = b.doImport(frags, clearFrags) + if err != nil { + t.Fatalf("doing import: %v", err) + } + + b.reset() + + for i := 9; i < 19; i++ { + r.ID = uint64(i) + if i%2 == 0 { + r.Values[0] = "a" + r.Values[1] = "b" + r.Values[2] = "c" + r.Values[3] = int64(99) + r.Values[4] = uint64(1) + } else { + r.Values[0] = "x" + r.Values[1] = "y" + r.Values[2] = "z" + r.Values[3] = int64(-10) + r.Values[4] = uint64(2) + } + err := b.Add(r) + if i != 18 && err != nil { + t.Fatalf("unexpected err adding record: %v", err) + } + if i == 18 && err != ErrBatchNowFull { + t.Fatalf("unexpected err: %v", err) + } + } + + // should do nothing + err = b.doTranslation() + if err != nil { + t.Fatalf("doing translation: %v", err) + } + + frags, clearFrags, err = b.makeFragments(make(fragments), make(fragments)) + if err != nil { + t.Errorf("making fragments: %v", err) + } + + err = b.doImport(frags, clearFrags) + if err != nil { + t.Fatalf("doing import: %v", err) + } + + for fidx, rowIDs := range b.rowIDs { + if fidx == 3 { + if len(rowIDs) != 0 { + t.Fatalf("expected no rowIDs for int field, but got: %v", rowIDs) + } + continue + } + // we don't know which key will get translated first, but we do know the pattern + if !reflect.DeepEqual(rowIDs, []uint64{1, 2, 1, 2, 1, 2, 1, 2, 1, 2}) && !reflect.DeepEqual(rowIDs, []uint64{2, 1, 2, 1, 2, 1, 2, 1, 2, 1}) { + t.Fatalf("unexpected row ids for field %d: %v", fidx, rowIDs) + } + } + + b.reset() + + for i := 19; i < 29; i++ { + r.ID = uint64(i) + if i%2 == 0 { + r.Values[0] = "d" + r.Values[1] = "e" + r.Values[2] = "f" + r.Values[3] = int64(100) + r.Values[4] = uint64(3) + } else { + r.Values[0] = "u" + r.Values[1] = "v" + r.Values[2] = "w" + r.Values[3] = int64(0) + r.Values[4] = uint64(4) + } + err := b.Add(r) + if i != 28 && err != nil { + t.Fatalf("unexpected err adding record: %v", err) + } + if i == 28 && err != ErrBatchNowFull { + t.Fatalf("unexpected err: %v", err) + } + } + + err = b.doTranslation() + if err != nil { + t.Fatalf("doing translation: %v", err) + } + + frags, clearFrags, err = b.makeFragments(make(fragments), make(fragments)) + if err != nil { + t.Errorf("making fragments: %v", err) + } + + err = b.doImport(frags, clearFrags) + if err != nil { + t.Fatalf("doing import: %v", err) + } + + for fidx, rowIDs := range b.rowIDs { + // we don't know which key will get translated first, but we do know the pattern + if fidx == 3 { + if len(rowIDs) != 0 { + t.Fatalf("expected no rowIDs for int field, but got: %v", rowIDs) + } + continue + } + if !reflect.DeepEqual(rowIDs, []uint64{3, 4, 3, 4, 3, 4, 3, 4, 3, 4}) && !reflect.DeepEqual(rowIDs, []uint64{4, 3, 4, 3, 4, 3, 4, 3, 4, 3}) { + t.Fatalf("unexpected row ids for field %d: %v", fidx, rowIDs) + } + } + + frags, _, err = b.makeFragments(make(fragments), make(fragments)) + if err != nil { + t.Fatalf("making fragments: %v", err) + } + + var n int + for key := range frags { + if key.shard == 0 { + n++ + } + } + if n != 5 { // zero, one, two, four (three is an int field so not in fragments) + _exists + t.Fatalf("there should be 5 views, but have %d", n) + } + + resp, err := client.Query(idx.BatchQuery(fields[0].Row("a"), + fields[1].Row("b"), + fields[2].Row("c"), + fields[3].Equals(99))) + if err != nil { + t.Fatalf("querying: %v", err) + } + + results := resp.Results() + for _, j := range []int{0, 2, 3} { + cols := results[j].Row().Columns + if !reflect.DeepEqual(cols, []uint64{0, 2, 4, 6, 10, 12, 14, 16, 18}) { + t.Fatalf("unexpected columns for a: %v", cols) + } + } + res := results[1] + cols := res.Row().Columns + if !reflect.DeepEqual(cols, []uint64{0, 2, 4, 6, 8, 10, 12, 14, 16, 18}) { + t.Fatalf("unexpected columns for field 1 row b: %v", cols) + } + + resp, err = client.Query(idx.BatchQuery(fields[0].Row("d"), + fields[1].Row("e"), + fields[2].Row("f"))) + if err != nil { + t.Fatalf("querying: %v", err) + } + + results = resp.Results() + for _, res := range results { + cols := res.Row().Columns + if !reflect.DeepEqual(cols, []uint64{20, 22, 24, 26, 28}) { + t.Fatalf("unexpected columns: %v", cols) + } + } + + resp, err = client.Query(idx.BatchQuery(fields[3].GT(-11), + fields[3].Equals(0), + fields[3].Equals(100), + fields[4].Range(1, time.Date(2019, time.January, 1, 0, 0, 0, 0, time.UTC), time.Date(2019, time.January, 29, 0, 0, 0, 0, time.UTC)), + fields[4].Range(1, time.Date(2019, time.February, 1, 0, 0, 0, 0, time.UTC), time.Date(2019, time.February, 29, 0, 0, 0, 0, time.UTC)))) + if err != nil { + t.Fatalf("querying: %v", err) + } + results = resp.Results() + cols = results[0].Row().Columns + if !reflect.DeepEqual(cols, []uint64{0, 1, 2, 3, 4, 5, 6, 7, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28}) { + t.Fatalf("all columns (but 8) should be greater than -11, but got: %v", cols) + } + cols = results[1].Row().Columns + if !reflect.DeepEqual(cols, []uint64{19, 21, 23, 25, 27}) { + t.Fatalf("wrong cols for ==0: %v", cols) + } + cols = results[2].Row().Columns + if !reflect.DeepEqual(cols, []uint64{20, 22, 24, 26, 28}) { + t.Fatalf("wrong cols for ==100: %v", cols) + } + cols = results[3].Row().Columns + exp := []uint64{0, 2, 4, 6, 10, 12, 14, 16, 18} + if !reflect.DeepEqual(cols, exp) { + t.Fatalf("wrong cols for January: got/want\n%v\n%v", cols, exp) + } + cols = results[4].Row().Columns + exp = []uint64{1, 3, 5, 7} + if !reflect.DeepEqual(cols, exp) { + t.Fatalf("wrong cols for January: got/want\n%v\n%v", cols, exp) + } + + b.reset() + r.ID = uint64(0) + r.Values[0] = "x" + r.Values[1] = "b" + r.Clears[0] = "a" + r.Clears[1] = "b" // b should get cleared + err = b.Add(r) + if err != nil { + t.Fatalf("adding with clears: %v", err) + } + err = b.Import() + if err != nil { + t.Fatalf("importing w/clears: %v", err) + } + resp, err = client.Query(idx.BatchQuery( + fields[0].Row("a"), + fields[0].Row("x"), + fields[1].Row("b"), + )) + if err != nil { + t.Fatalf("querying after clears: %v", err) + } + if arow := resp.Results()[0].Row().Columns; arow[0] == 0 { + t.Errorf("shouldn't have id 0 in row a after clearing! %v", arow) + } + if xrow := resp.Results()[1].Row().Columns; xrow[0] != 0 { + t.Errorf("should have id 0 in row x after setting %v", xrow) + } + if brow := resp.Results()[2].Row().Columns; brow[0] == 0 { + t.Errorf("shouldn't have id 0 in row b after clearing! %v", brow) + } + + // TODO test importing across multiple shards +} + +func TestBatchesStringIDs(t *testing.T) { + client := DefaultClient() + schema := NewSchema() + idx := schema.Index("gopilosatest-blah", OptIndexKeys(true)) + fields := make([]*Field, 1) + fields[0] = idx.Field("zero", OptFieldKeys(true)) + err := client.SyncSchema(schema) + if err != nil { + t.Fatalf("syncing schema: %v", err) + } + defer func() { + err := client.DeleteIndex(idx) + if err != nil { + t.Logf("problem cleaning up from test: %v", err) + } + }() + + b, err := NewBatch(client, 3, idx, fields) + if err != nil { + t.Fatalf("getting new batch: %v", err) + } + + r := Row{Values: make([]interface{}, 1)} + + for i := 0; i < 3; i++ { + r.ID = strconv.Itoa(i) + if i%2 == 0 { + r.Values[0] = "a" + } else { + r.Values[0] = "x" + } + err := b.Add(r) + if err != nil && err != ErrBatchNowFull { + t.Fatalf("unexpected err adding record: %v", err) + } + } + + if len(b.toTranslateID) != 3 { + t.Fatalf("id translation table unexpected size: %v", b.toTranslateID) + } + for i, k := range b.toTranslateID { + if ik, err := strconv.Atoi(k); err != nil || ik != i { + t.Errorf("unexpected toTranslateID key %s at index %d", k, i) + } + } + + err = b.doTranslation() + if err != nil { + t.Fatalf("translating: %v", err) + } + + // the ids are based off what the strings hash to, and are at the + // very beginning of a few different shards. this could change if + // Pilosa's hashing algorithm changes. + if err := isPermutationOfInt(b.ids, []uint64{44040193, 45088769, 41943041}); err != nil { + t.Fatalf("wrong ids: %v. exp/got:\n%v\n%v", err, []uint64{44040193, 45088769, 41943041}, b.ids) + } + + err = b.Import() + if err != nil { + t.Fatalf("importing: %v", err) + } + + resp, err := client.Query(idx.BatchQuery(fields[0].Row("a"), fields[0].Row("x"))) + if err != nil { + t.Fatalf("querying: %v", err) + } + + results := resp.Results() + for i, res := range results { + cols := res.Row().Keys + if i == 0 && !reflect.DeepEqual(cols, []string{"0", "2"}) && !reflect.DeepEqual(cols, []string{"2", "0"}) { + t.Fatalf("unexpected columns: %v", cols) + } + if i == 1 && !reflect.DeepEqual(cols, []string{"1"}) { + t.Fatalf("unexpected columns: %v", cols) + } + } + + b.reset() + + r.ID = "1" + r.Values[0] = "a" + err = b.Add(r) + if err != nil { + t.Fatalf("unexpected err adding record: %v", err) + } + + r.ID = "3" + r.Values[0] = "z" + err = b.Add(r) + if err != nil { + t.Fatalf("unexpected err adding record: %v", err) + } + + err = b.Import() + if err != nil { + t.Fatalf("importing: %v", err) + } + + resp, err = client.Query(idx.BatchQuery(fields[0].Row("a"), fields[0].Row("z"))) + if err != nil { + t.Fatalf("querying: %v", err) + } + + results = resp.Results() + for i, res := range results { + cols := res.Row().Keys + if err := isPermutationOf(cols, []string{"0", "1", "2"}); i == 0 && err != nil { + t.Fatalf("unexpected columns: %v: %v", cols, err) + } + if i == 1 && !reflect.DeepEqual(cols, []string{"3"}) { + t.Fatalf("unexpected columns: %v", cols) + } + } + +} + +func isPermutationOf(one, two []string) error { + if len(one) != len(two) { + return errors.Errorf("different lengths %d and %d", len(one), len(two)) + } +outer: + for _, vOne := range one { + for j, vTwo := range two { + if vOne == vTwo { + two = append(two[:j], two[j+1:]...) + continue outer + } + } + return errors.Errorf("%s in one but not two", vOne) + } + if len(two) != 0 { + return errors.Errorf("vals in two but not one: %v", two) + } + return nil +} + +func isPermutationOfInt(one, two []uint64) error { + if len(one) != len(two) { + return errors.Errorf("different lengths %d and %d", len(one), len(two)) + } +outer: + for _, vOne := range one { + for j, vTwo := range two { + if vOne == vTwo { + two = append(two[:j], two[j+1:]...) + continue outer + } + } + return errors.Errorf("%d in one but not two", vOne) + } + if len(two) != 0 { + return errors.Errorf("vals in two but not one: %v", two) + } + return nil +} + +func TestQuantizedTime(t *testing.T) { + cases := []struct { + name string + time time.Time + year string + month string + day string + hour string + quantum TimeQuantum + reset bool + exp []string + expErr string + }{ + { + name: "no time quantum", + expErr: "", + }, + { + name: "no time quantum with data", + year: "2017", + exp: []string{}, + expErr: "", + }, + { + name: "no data", + quantum: TimeQuantumYear, + exp: nil, + expErr: "", + }, + { + name: "timestamp", + time: time.Date(2013, time.October, 16, 17, 34, 43, 0, time.FixedZone("UTC-5", -5*60*60)), + quantum: "YMDH", + exp: []string{"2013", "201310", "20131016", "2013101617"}, + }, + { + name: "timestamp-less-granular", + time: time.Date(2013, time.October, 16, 17, 34, 43, 0, time.FixedZone("UTC-5", -5*60*60)), + quantum: "YM", + exp: []string{"2013", "201310"}, + }, + { + name: "timestamp-mid-granular", + time: time.Date(2013, time.October, 16, 17, 34, 43, 0, time.FixedZone("UTC-5", -5*60*60)), + quantum: "MD", + exp: []string{"201310", "20131016"}, + }, + { + name: "justyear", + year: "2013", + quantum: "Y", + exp: []string{"2013"}, + }, + { + name: "justyear-wantmonth", + year: "2013", + quantum: "YM", + expErr: "no data set for month", + }, + { + name: "timestamp-changeyear", + time: time.Date(2013, time.October, 16, 17, 34, 43, 0, time.FixedZone("UTC-5", -5*60*60)), + year: "2019", + quantum: "YMDH", + exp: []string{"2019", "201910", "20191016", "2019101617"}, + }, + { + name: "yearmonthdayhour", + year: "2013", + month: "10", + day: "16", + hour: "17", + quantum: "YMDH", + exp: []string{"2013", "201310", "20131016", "2013101617"}, + }, + { + name: "timestamp-changehour", + time: time.Date(2013, time.October, 16, 17, 34, 43, 0, time.FixedZone("UTC-5", -5*60*60)), + hour: "05", + quantum: "MDH", + exp: []string{"201310", "20131016", "2013101605"}, + }, + { + name: "timestamp", + time: time.Date(2013, time.October, 16, 17, 34, 43, 0, time.FixedZone("UTC-5", -5*60*60)), + quantum: "YMDH", + reset: true, + exp: nil, + }, + } + + for i, test := range cases { + t.Run(test.name+strconv.Itoa(i), func(t *testing.T) { + tq := QuantizedTime{} + var zt time.Time + if zt != test.time { + tq.Set(test.time) + } + if test.year != "" { + tq.SetYear(test.year) + } + if test.month != "" { + tq.SetMonth(test.month) + } + if test.day != "" { + tq.SetDay(test.day) + } + if test.hour != "" { + tq.SetHour(test.hour) + } + if test.reset { + tq.Reset() + } + + views, err := tq.views(test.quantum) + if !reflect.DeepEqual(views, test.exp) { + t.Errorf("unexpected views, got/want:\n%v\n%v\n", views, test.exp) + } + if (err != nil && err.Error() != test.expErr) || (err == nil && test.expErr != "") { + t.Errorf("unexpected error, got/want:\n%v\n%s\n", err, test.expErr) + } + }) + } + +} + +func TestBatchStaleness(t *testing.T) { + client := DefaultClient() + schema := NewSchema() + idx := schema.Index("gopilosatest-blah") + field := idx.Field("anint", OptFieldTypeInt()) + err := client.SyncSchema(schema) + if err != nil { + t.Fatalf("syncing schema: %v", err) + } + defer func() { + err := client.DeleteIndex(idx) + if err != nil { + t.Logf("problem cleaning up from test: %v", err) + } + }() + + b, err := NewBatch(client, 3, idx, []*Field{field}, OptMaxStaleness(time.Millisecond)) + if err != nil { + t.Fatalf("getting batch: %v", err) + } + + r := Row{ID: uint64(0), Values: []interface{}{int64(0)}} + err = b.Add(r) + if err != nil && err != ErrBatchNowFull { + t.Fatalf("adding to batch: %v", err) + } + + // sleep so batch becomes stale + time.Sleep(time.Millisecond) + + r = Row{ID: uint64(1), Values: []interface{}{int64(0)}} + err = b.Add(r) + if err != ErrBatchNowStale { + t.Fatal("batch expected to be stale") + } +} diff --git a/client/client.go b/client/client.go new file mode 100644 index 000000000..846daebb4 --- /dev/null +++ b/client/client.go @@ -0,0 +1,1811 @@ +// Copyright 2017 Pilosa Corp. +// +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions +// are met: +// +// 1. Redistributions of source code must retain the above copyright +// notice, this list of conditions and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright +// notice, this list of conditions and the following disclaimer in the +// documentation and/or other materials provided with the distribution. +// +// 3. Neither the name of the copyright holder nor the names of its +// contributors may be used to endorse or promote products derived +// from this software without specific prior written permission. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND +// CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, +// INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF +// MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR +// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, +// BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH +// DAMAGE. + +package client + +import ( + "bytes" + "crypto/tls" + "encoding/json" + "fmt" + "hash/fnv" + "io" + "log" + "math/rand" + "net" + "net/http" + "net/url" + "os" + "strconv" + "strings" + "sync" + "time" + + "github.com/golang/protobuf/proto" //nolint:staticcheck + "github.com/opentracing/opentracing-go" + "github.com/pilosa/pilosa/v2" + "github.com/pilosa/pilosa/v2/internal" + pnet "github.com/pilosa/pilosa/v2/net" + "github.com/pilosa/pilosa/v2/pql" + "github.com/pilosa/pilosa/v2/roaring" + "github.com/pilosa/pilosa/v2/stats" + "github.com/pkg/errors" + "golang.org/x/sync/errgroup" +) + +// PQLVersion is the version of PQL expected by the client +const PQLVersion = "1.0" + +// DefaultShardWidth is used if an index doesn't have it defined. +const DefaultShardWidth = 1 << 20 + +const maxHosts = 10 + +// Client is the HTTP client for Pilosa server. +type Client struct { + cluster *Cluster + client *http.Client + logger *log.Logger + primaryURI *pnet.URI + primaryLock *sync.RWMutex + manualFragmentNode *fragmentNode + manualServerURI *pnet.URI + tracer opentracing.Tracer + Stats stats.StatsClient + // An exponential backoff algorithm retries requests exponentially (if an HTTP request fails), + // increasing the waiting time between retries up to a maximum backoff time. + maxBackoff time.Duration + maxRetries int + + nat map[pnet.URI]pnet.URI + + shardNodes shardNodes + tick *time.Ticker + done chan struct{} + + // TODO — remove this! Pilosa should have API which gives us the node->partition mapping + Hasher Hasher +} + +func (c *Client) getURIsForShard(index string, shard uint64) ([]*pnet.URI, error) { + uris, ok := c.shardNodes.Get(index, shard) + if ok { + return uris, nil + } + fragmentNodes, err := c.fetchFragmentNodes(index, shard) + if err != nil { + return nil, errors.Wrap(err, "trying to look up nodes for shard") + } + uris = make([]*pnet.URI, 0, len(fragmentNodes)) + for _, fn := range fragmentNodes { + uris = append(uris, fn.URI()) + } + c.shardNodes.Put(index, shard, uris) + return uris, nil +} + +func (c *Client) runChangeDetection() { + for { + select { + case <-c.tick.C: + c.detectClusterChanges() + case <-c.done: + return + } + } +} + +func (c *Client) Close() error { + c.tick.Stop() + close(c.done) + return nil +} + +// detectClusterChanges chooses a random index and shard from the +// shardNodes cache and deletes it. It then looks it up from Pilosa to +// see if it still matches, and if not it drops the whole cache. +func (c *Client) detectClusterChanges() { + c.shardNodes.mu.Lock() + needsUnlock := true + // we rely on Go's random map iteration order to get a random + // element. If it doesn't end up being random, it shouldn't + // actually matter. + for index, shardMap := range c.shardNodes.data { + for shard, uris := range shardMap { + delete(shardMap, shard) + c.shardNodes.data[index] = shardMap + c.shardNodes.mu.Unlock() + needsUnlock = false + newURIs, err := c.getURIsForShard(index, shard) // refetch URIs from server. + if err != nil { + c.logger.Printf("problem invalidating shard node cache: %v", err) + return + } + if len(uris) != len(newURIs) { + c.logger.Printf("invalidating shard node cache old: %v, new: %v", uris, newURIs) + c.shardNodes.Invalidate() + return + } + for i := range uris { + u1, u2 := uris[i], newURIs[i] + if *u1 != *u2 { + c.logger.Printf("invalidating shard node cache, uri mismatch at %d old: %v, new: %v", i, uris, newURIs) + c.shardNodes.Invalidate() + return + } + } + break + } + break + } + if needsUnlock { + c.shardNodes.mu.Unlock() + } +} + +// DefaultClient creates a client with the default address and options. +func DefaultClient() *Client { + return newClientWithCluster(NewClusterWithHost(pnet.DefaultURI()), nil) +} + +func newClientFromAddresses(addresses []string, options *ClientOptions) (*Client, error) { + uris := make([]*pnet.URI, len(addresses)) + for i, address := range addresses { + uri, err := pnet.NewURIFromAddress(address) + if err != nil { + return nil, err + } + uris[i] = uri + } + cluster := NewClusterWithHost(uris...) + client := newClientWithCluster(cluster, options) + return client, nil +} + +func newClientWithCluster(cluster *Cluster, options *ClientOptions) *Client { + client := newClientWithOptions(options) + client.cluster = cluster + return client +} + +func newClientWithURI(uri *pnet.URI, options *ClientOptions) *Client { + client := newClientWithOptions(options) + if options.manualServerAddress { + fragmentNode := newFragmentNodeFromURI(uri) + client.manualFragmentNode = &fragmentNode + client.manualServerURI = uri + client.cluster = NewClusterWithHost() + } + client.cluster = NewClusterWithHost(uri) + return client +} + +func newClientWithOptions(options *ClientOptions) *Client { + if options == nil { + options = &ClientOptions{} + } + options = options.withDefaults() + + c := &Client{ + client: newHTTPClient(options.withDefaults()), + logger: log.New(os.Stderr, "", log.Ldate|log.Ltime|log.Lmicroseconds), + primaryLock: &sync.RWMutex{}, + + shardNodes: newShardNodes(), + tick: time.NewTicker(time.Minute), + done: make(chan struct{}), + + nat: options.nat, + + // TODO get rid of this. Pilosa should have api to expose node->partition mapping. + Hasher: &jmphasher{}, + } + + if options.tracer == nil { + c.tracer = NoopTracer{} + } else { + c.tracer = options.tracer + } + if options.stats == nil { + c.Stats = stats.NopStatsClient + } else { + c.Stats = options.stats + } + + c.maxRetries = *options.retries + c.maxBackoff = 2 * time.Minute + go c.runChangeDetection() + return c + +} + +// NewClient creates a client with the given address, URI, or cluster and options. +func NewClient(addrURIOrCluster interface{}, options ...ClientOption) (*Client, error) { + var cluster *Cluster + clientOptions := &ClientOptions{ + nat: make(map[pnet.URI]pnet.URI), + } + err := clientOptions.addOptions(options...) + if err != nil { + return nil, err + } + + switch u := addrURIOrCluster.(type) { + case string: + uri, err := pnet.NewURIFromAddress(u) + if err != nil { + return nil, err + } + return newClientWithURI(uri, clientOptions), nil + case []string: + if len(u) == 1 { + uri, err := pnet.NewURIFromAddress(u[0]) + if err != nil { + return nil, err + } + return newClientWithURI(uri, clientOptions), nil + } else if clientOptions.manualServerAddress { + return nil, ErrSingleServerAddressRequired + } + return newClientFromAddresses(u, clientOptions) + case *pnet.URI: + uriCopy := *u + return newClientWithURI(&uriCopy, clientOptions), nil + case []*pnet.URI: + if len(u) == 1 { + uriCopy := *u[0] + return newClientWithURI(&uriCopy, clientOptions), nil + } else if clientOptions.manualServerAddress { + return nil, ErrSingleServerAddressRequired + } + cluster = NewClusterWithHost(u...) + case *Cluster: + cluster = u + case nil: + cluster = NewClusterWithHost() + default: + return nil, ErrAddrURIClusterExpected + } + + return newClientWithCluster(cluster, clientOptions), nil +} + +// Query runs the given query against the server with the given options. +// Pass nil for default options. +func (c *Client) Query(query PQLQuery, options ...interface{}) (*QueryResponse, error) { + span := c.tracer.StartSpan("Client.Query") + defer span.Finish() + + if err := query.Error(); err != nil { + return nil, err + } + queryOptions := &QueryOptions{} + err := queryOptions.addOptions(options...) + if err != nil { + return nil, err + } + serializedQuery := query.Serialize() + reqData, err := makeRequestData(serializedQuery.String(), queryOptions) + if err != nil { + return nil, errors.Wrap(err, "making request data") + } + path := fmt.Sprintf("/index/%s/query", query.Index().name) + _, respData, err := c.HTTPRequest("POST", path, reqData, defaultProtobufHeaders()) + if err != nil { + return nil, err + } + iqr := &internal.QueryResponse{} + err = proto.Unmarshal(respData, iqr) + if err != nil { + return nil, err + } + queryResponse, err := newQueryResponseFromInternal(iqr) + if err != nil { + return nil, err + } + return queryResponse, nil +} + +// CreateIndex creates an index on the server using the given Index struct. +func (c *Client) CreateIndex(index *Index) error { + span := c.tracer.StartSpan("Client.CreateIndex") + defer span.Finish() + + data := []byte(index.options.String()) + path := fmt.Sprintf("/index/%s", index.name) + status, body, err := c.HTTPRequest("POST", path, data, nil) + if err != nil { + return errors.Wrapf(err, "creating index: %s", index.name) + } + var resp struct { + CreatedAt int64 `json:"createdAt,omitempty"` + } + if err := json.Unmarshal(body, &resp); err == nil && resp.CreatedAt != 0 { + index.createdAt = resp.CreatedAt + } + + if status == http.StatusConflict { + return ErrIndexExists + } + return nil +} + +// CreateField creates a field on the server using the given Field struct. +func (c *Client) CreateField(field *Field) error { + span := c.tracer.StartSpan("Client.CreateField") + defer span.Finish() + + data := []byte(field.options.String()) + path := fmt.Sprintf("/index/%s/field/%s", field.index.name, field.name) + status, body, err := c.HTTPRequest("POST", path, data, nil) + if err != nil { + return errors.Wrapf(err, "creating field: %s in index: %s", field.name, field.index.name) + } + var resp struct { + CreatedAt int64 `json:"createdAt,omitempty"` + } + if err := json.Unmarshal(body, &resp); err == nil && resp.CreatedAt != 0 { + field.createdAt = resp.CreatedAt + } + + if status == http.StatusConflict { + return ErrFieldExists + } + return nil +} + +// EnsureIndex creates an index on the server if it does not exist. +func (c *Client) EnsureIndex(index *Index) error { + err := c.CreateIndex(index) + if err == ErrIndexExists { + return nil + } + return errors.Wrap(err, "creating index") +} + +func (c *Client) SyncIndex(index *Index) error { + err := c.EnsureIndex(index) + if err != nil { + return errors.Wrapf(err, "ensuring index exists") + } + + for name, field := range index.fields { + if name == "_exists" { + continue + } + err = c.EnsureField(field) + if err != nil { + return errors.Wrapf(err, "ensuring field") + } + } + + return nil +} + +// EnsureField creates a field on the server if it doesn't exists. +func (c *Client) EnsureField(field *Field) error { + err := c.CreateField(field) + if err == ErrFieldExists { + return nil + } + return err +} + +// DeleteIndex deletes an index on the server. +func (c *Client) DeleteIndex(index *Index) error { + if index != nil { + return c.DeleteIndexByName(index.Name()) + } + return nil +} + +// DeleteIndexByName deletes the named index on the server. +func (c *Client) DeleteIndexByName(index string) error { + span := c.tracer.StartSpan("Client.DeleteIndex") + defer span.Finish() + + path := fmt.Sprintf("/index/%s", index) + _, _, err := c.HTTPRequest("DELETE", path, nil, nil) + return err +} + +// DeleteField deletes a field on the server. +func (c *Client) DeleteField(field *Field) error { + span := c.tracer.StartSpan("Client.DeleteField") + defer span.Finish() + + path := fmt.Sprintf("/index/%s/field/%s", field.index.name, field.name) + _, _, err := c.HTTPRequest("DELETE", path, nil, nil) + return err +} + +// SyncSchema updates a schema with the indexes and fields on the server and +// creates the indexes and fields in the schema on the server side. +// This function does not delete indexes and the fields on the server side nor in the schema. +func (c *Client) SyncSchema(schema *Schema) error { + span := c.tracer.StartSpan("Client.SyncSchema") + defer span.Finish() + schema.mu.Lock() + defer schema.mu.Unlock() + serverSchema, err := c.Schema() + if err != nil { + return err + } + serverSchema.mu.RLock() + defer serverSchema.mu.RUnlock() + + return c.syncSchema(schema, serverSchema) +} + +func (c *Client) syncSchema(schema *Schema, serverSchema *Schema) error { + var err error + + // find out local - remote schema + diffSchema := schema.diff(serverSchema) + // create the indexes and fields which doesn't exist on the server side + for indexName, index := range diffSchema.indexes { + if _, ok := serverSchema.indexes[indexName]; !ok { + err = c.EnsureIndex(index) + if err != nil { + return errors.Wrap(err, "ensuring index") + } + } + for name, field := range index.fields { + if name == "_exists" { + continue + } + err = c.EnsureField(field) + if err != nil { + return errors.Wrapf(err, "ensuring field") + } + } + } + + // find out remote - local schema + diffSchema = serverSchema.diff(schema) + for indexName, index := range diffSchema.indexes { + if localIndex, ok := schema.indexes[indexName]; !ok { + schema.indexes[indexName] = index + } else { + for fieldName, field := range index.fields { + localIndex.fields[fieldName] = field + } + } + } + + return nil +} + +// Schema returns the indexes and fields on the server. +func (c *Client) Schema() (*Schema, error) { + span := c.tracer.StartSpan("Client.Schema") + defer span.Finish() + + var indexes []SchemaIndex + indexes, err := c.readSchema() + if err != nil { + return nil, err + } + schema := NewSchema() + for _, indexInfo := range indexes { + index := schema.indexWithOptions(indexInfo.Name, indexInfo.CreatedAt, indexInfo.ShardWidth, indexInfo.Options.asIndexOptions()) + + for _, fieldInfo := range indexInfo.Fields { + index.fieldWithOptions(fieldInfo.Name, fieldInfo.CreatedAt, fieldInfo.Options.asFieldOptions()) + } + } + return schema, nil +} + +// Import imports data for a single shard using the regular import +// endpoint rather than import-roaring. This is good for e.g. mutex or +// bool fields where import-roaring is not supported. +func (c *Client) Import(field *Field, shard uint64, vals, ids []uint64, clear bool) error { + path, data, err := c.EncodeImport(field, shard, vals, ids, clear) + if err != nil { + return errors.Wrap(err, "encoding import request") + } + err = c.DoImport(field.index.Name(), shard, path, data) + return errors.Wrap(err, "doing import") +} + +// EncodeImport computes the HTTP path and payload for an import +// request. It is typically followed by a call to DoImport. +func (c *Client) EncodeImport(field *Field, shard uint64, vals, ids []uint64, clear bool) (path string, data []byte, err error) { + msg := &internal.ImportRequest{ + Index: field.index.Name(), + IndexCreatedAt: field.index.CreatedAt(), + Field: field.Name(), + FieldCreatedAt: field.CreatedAt(), + Shard: shard, + RowIDs: vals, + ColumnIDs: ids, + } + data, err = proto.Marshal(msg) + if err != nil { + return "", nil, errors.Wrap(err, "marshaling Import to protobuf") + } + path = fmt.Sprintf("/index/%s/field/%s/import?clear=%s&ignoreKeyCheck=true", field.index.Name(), field.Name(), strconv.FormatBool(clear)) + return path, data, nil +} + +// DoImport takes a path and data payload (normally from EncodeImport +// or EncodeImportValues), logs the import, finds all nodes which own +// this shard, and concurrently imports to those nodes. +func (c *Client) DoImport(index string, shard uint64, path string, data []byte) error { + // Unlike ImportRoaring, Pilosa does not forward requests to the + // .../import endpoint to all replicas, so we must do that + // here. Yes this is odd. To make it worse, if a ../import request + // was made with keys that needed to be translated server side, + // Pilosa would handle sending the translated data to all the + // appropriate nodes and replicas. + + uris, err := c.getURIsForShard(index, shard) + if err != nil { + return errors.Wrap(err, "getting uris") + } + + eg := errgroup.Group{} + for _, uri := range uris { + uri := uri + eg.Go(func() error { + return c.importData(uri, path, data) + }) + } + return errors.Wrap(eg.Wait(), "importing to nodes") +} + +// EncodeImportValues computes the HTTP path and payload for an +// import-values request. It is typically followed by a call to +// DoImportValues. +func (c *Client) EncodeImportValues(field *Field, shard uint64, vals []int64, ids []uint64, clear bool) (path string, data []byte, err error) { + msg := &internal.ImportValueRequest{ + Index: field.index.Name(), + IndexCreatedAt: field.index.CreatedAt(), + Field: field.Name(), + FieldCreatedAt: field.CreatedAt(), + Shard: shard, + ColumnIDs: ids, + Values: vals, + } + data, err = proto.Marshal(msg) + if err != nil { + return "", nil, errors.Wrap(err, "marshaling ImportValue to protobuf") + } + path = fmt.Sprintf("/index/%s/field/%s/import?clear=%s&ignoreKeyCheck=true", field.index.Name(), field.Name(), strconv.FormatBool(clear)) + return path, data, nil +} + +// ImportValues takes the given integer values and column ids (which +// must all be in the given shard) and imports them into the given +// index,field,shard on all nodes which should hold that shard. It +// assumes that the ids have been translated from keys if necessary +// and so tells Pilosa to ignore checking if the index uses column +// keys. ImportValues wraps EncodeImportValues and DoImportValues — +// these are broken out and exported so that performance conscious +// users can re-use the same vals and ids byte buffers for local +// encoding, while performing the imports concurrently. +func (c *Client) ImportValues(field *Field, shard uint64, vals []int64, ids []uint64, clear bool) error { + path, data, err := c.EncodeImportValues(field, shard, vals, ids, clear) + if err != nil { + return errors.Wrap(err, "encoding import-values request") + } + err = c.DoImportValues(field.index.Name(), shard, path, data) + return errors.Wrap(err, "doing import values") +} + +// DoImportValues is deprecated. Use DoImport. +func (c *Client) DoImportValues(index string, shard uint64, path string, data []byte) error { + return c.DoImport(index, shard, path, data) +} + +func (c *Client) fetchFragmentNodes(indexName string, shard uint64) ([]fragmentNode, error) { + if c.manualFragmentNode != nil { + return []fragmentNode{*c.manualFragmentNode}, nil + } + path := fmt.Sprintf("/internal/fragment/nodes?shard=%d&index=%s", shard, indexName) + _, body, err := c.HTTPRequest("GET", path, []byte{}, nil) + if err != nil { + return nil, err + } + fragmentNodes := []fragmentNode{} + var fragmentNodeURIs []fragmentNodeRoot + err = json.Unmarshal(body, &fragmentNodeURIs) + if err != nil { + return nil, errors.Wrap(err, "unmarshaling fragment node URIs") + } + for _, nodeURI := range fragmentNodeURIs { + fragmentNodes = append(fragmentNodes, nodeURI.URI) + } + return fragmentNodes, nil +} + +func (c *Client) fetchPrimaryNode() (fragmentNode, error) { + if c.manualFragmentNode != nil { + return *c.manualFragmentNode, nil + } + status, err := c.Status() + if err != nil { + return fragmentNode{}, err + } + for _, node := range status.Nodes { + if node.IsPrimary { + nodeURI := node.URI.URI().Translate(c.nat) + return fragmentNode{ //nolint:gosimple + Scheme: nodeURI.Scheme, + Host: nodeURI.Host, + Port: nodeURI.Port, + }, nil + } + } + return fragmentNode{}, errors.New("Primary node not found") +} + +func (c *Client) importData(uri *pnet.URI, path string, data []byte) error { + if status, _, err := c.doRequest(uri, "POST", path, defaultProtobufHeaders(), data); err != nil { + return errors.Wrapf(err, "import to %s", uri.HostPort()) + } else if status == http.StatusPreconditionFailed { + return ErrPreconditionFailed + } + + return nil +} + +// ImportRoaringBitmap can import pre-made bitmaps for a number of +// different views into the given field/shard. If the view name in the +// map is an empty string, the standard view will be used. +func (c *Client) ImportRoaringBitmap(field *Field, shard uint64, views map[string]*roaring.Bitmap, clear bool) error { + uris, err := c.getURIsForShard(field.index.Name(), shard) + if err != nil { + return errors.Wrap(err, "getting URIs for import") + } + err = c.importRoaringBitmap(uris[0], field, shard, views, &ImportOptions{clear: clear}) + return errors.Wrap(err, "importing bitmap") +} + +func (c *Client) importRoaringBitmap(uri *pnet.URI, field *Field, shard uint64, views viewImports, options *ImportOptions) error { + protoViews := []*internal.ImportRoaringRequestView{} + for name, bmp := range views { + buf := &bytes.Buffer{} + _, err := bmp.WriteTo(buf) + if err != nil { + return errors.Wrap(err, "marshalling bitmap") + } + protoViews = append(protoViews, &internal.ImportRoaringRequestView{ + Name: name, + Data: buf.Bytes(), + }) + } + params := url.Values{} + params.Add("clear", strconv.FormatBool(options.clear)) + path := makeRoaringImportPath(field, shard, params) + req := &internal.ImportRoaringRequest{ + Clear: options.clear, + Views: protoViews, + IndexCreatedAt: field.index.CreatedAt(), + FieldCreatedAt: field.CreatedAt(), + } + data, err := proto.Marshal(req) + if err != nil { + return err + } + + status, _, err := c.doRequest(uri, "POST", path, defaultProtobufHeaders(), data) + if err != nil { + return errors.Wrapf(err, "roaring import to %s, status: %d", uri.HostPort(), status) + } + if status == http.StatusPreconditionFailed { + return ErrPreconditionFailed + } + + return nil +} + +// ExportField exports columns for a field. +func (c *Client) ExportField(field *Field) (io.Reader, error) { + span := c.tracer.StartSpan("Client.ExportField") + defer span.Finish() + + var shardsMax map[string]uint64 + var err error + + status, err := c.Status() + if err != nil { + return nil, err + } + shardsMax, err = c.shardsMax() + if err != nil { + return nil, err + } + status.indexMaxShard = shardsMax + shardURIs, err := c.statusToNodeShardsForIndex(status, field.index.Name()) + if err != nil { + return nil, err + } + + return newExportReader(c, shardURIs, field), nil +} + +// Info returns the server's configuration/host information. +func (c *Client) Info() (Info, error) { + span := c.tracer.StartSpan("Client.Info") + defer span.Finish() + + _, data, err := c.HTTPRequest("GET", "/info", nil, nil) + if err != nil { + return Info{}, errors.Wrap(err, "requesting /info") + } + info := Info{} + err = json.Unmarshal(data, &info) + if err != nil { + return Info{}, errors.Wrap(err, "unmarshaling /info data") + } + return info, nil +} + +// Status returns the server's status. +func (c *Client) Status() (Status, error) { + span := c.tracer.StartSpan("Client.Status") + defer span.Finish() + + _, data, err := c.HTTPRequest("GET", "/status", nil, nil) + if err != nil { + return Status{}, errors.Wrap(err, "requesting /status") + } + status := Status{} + err = json.Unmarshal(data, &status) + if err != nil { + return Status{}, errors.Wrap(err, "unmarshaling /status data") + } + return status, nil +} + +func (c *Client) readSchema() ([]SchemaIndex, error) { + _, data, err := c.HTTPRequest("GET", "/schema", nil, nil) + if err != nil { + return nil, errors.Wrap(err, "requesting /schema") + } + schemaInfo := SchemaInfo{} + err = json.Unmarshal(data, &schemaInfo) + if err != nil { + return nil, errors.Wrap(err, "unmarshaling /schema data") + } + return schemaInfo.Indexes, nil +} + +func (c *Client) shardsMax() (map[string]uint64, error) { + _, data, err := c.HTTPRequest("GET", "/internal/shards/max", nil, nil) + if err != nil { + return nil, errors.Wrap(err, "requesting /internal/shards/max") + } + m := map[string]map[string]uint64{} + err = json.Unmarshal(data, &m) + if err != nil { + return nil, errors.Wrap(err, "unmarshaling /internal/shards/max data") + } + return m["standard"], nil +} + +// HTTPRequest sends an HTTP request to the Pilosa server (used by idk) +// nolint: deadcode +func (c *Client) HTTPRequest(method string, path string, data []byte, headers map[string]string) (status int, body []byte, err error) { + span := c.tracer.StartSpan("Client.HTTPRequest") + + status, body, err = c.httpRequest(method, path, data, headers, false) + span.Finish() + return +} + +// httpRequest makes a request to the cluster - use this when you want the +// client to choose a host, and it doesn't matter if the request goes to a +// specific host +func (c *Client) httpRequest(method string, path string, data []byte, headers map[string]string, usePrimary bool) (int, []byte, error) { + if data == nil { + data = []byte{} + } + + var ( + status int + body []byte + err error + ) + // try at most maxHosts non-failed hosts; protect against broken cluster.removeHost + for i := 0; i < maxHosts; i++ { + host, herr := c.host(usePrimary) + if herr != nil { + return status, nil, errors.Wrapf(herr, "getting host, previous err: %v", err) + } + // doRequest implements expotential backoff + status, body, err = c.doRequest(host, method, path, c.augmentHeaders(headers), data) + if err == nil { + break + } + if c.manualServerURI == nil { + if usePrimary { + c.primaryLock.Lock() + c.primaryURI = nil + c.primaryLock.Unlock() + } else { + c.logger.Printf("removing host (%s) due to '%v'\n", host.Normalize(), err) + c.cluster.RemoveHost(host) + } + } + } + + if err != nil { + err = errors.Wrap(err, ErrTriedMaxHosts.Error()) + } + + return status, body, err +} + +// host returns the first URI that applies, in this order: +// - a non-nil manualServerURI +// - primary URI (if usePrimary = true) +// - the next host from the node list (round-robin) +func (c *Client) host(usePrimary bool) (*pnet.URI, error) { + if c.manualServerURI != nil { + return c.manualServerURI, nil + } + var host *pnet.URI + if usePrimary { + c.primaryLock.RLock() + host = c.primaryURI + c.primaryLock.RUnlock() + if host == nil { + c.primaryLock.Lock() + if c.primaryURI == nil { + node, err := c.fetchPrimaryNode() + if err != nil { + c.primaryLock.Unlock() + return nil, errors.Wrap(err, "fetching primary node") + } + host = pnet.URIFromAddress(fmt.Sprintf("%s://%s:%d", node.Scheme, node.Host, node.Port)) + } else { + host = c.primaryURI + } + c.primaryURI = host + c.primaryLock.Unlock() + } + } else { + // get a host from the cluster + host = c.cluster.Host() + if host == nil { + return nil, ErrEmptyCluster + } + } + return host, nil +} + +// doRequest creates and performs an http request. +func (c *Client) doRequest(host *pnet.URI, method, path string, headers map[string]string, data []byte) (int, []byte, error) { + var ( + req *http.Request + resp *http.Response + err error + sleepTime time.Duration + rand = rand.New(rand.NewSource(time.Now().UnixNano())) + ) + + for retry := 0; ; { + if req, err = buildRequest(host, method, path, headers, data); err != nil { + return 0, nil, errors.Wrap(err, "building request") + } + if resp, err = c.client.Do(req); err != nil { + return 0, nil, errors.Wrap(err, "sending request") + } + if warning := resp.Header.Get("warning"); warning != "" { + c.logger.Println(warning) + } + + buf := bytes.NewBuffer(make([]byte, 0, 1+resp.ContentLength)) + _, err = buf.ReadFrom(resp.Body) + _ = resp.Body.Close() + if err != nil { + return resp.StatusCode, nil, errors.Wrap(err, "reading response body") + } + + switch { + case resp.StatusCode >= 200 && resp.StatusCode < 300: + // [200, 300): OK + return resp.StatusCode, buf.Bytes(), nil + + case resp.StatusCode == 409: + // 409 Conflict + return resp.StatusCode, buf.Bytes(), nil + + case resp.StatusCode == 412: + // 412 Precondition Failed + return resp.StatusCode, buf.Bytes(), nil + + case resp.StatusCode == 429: + // 429 Too Many Requests + // A Retry-After header might be included to this response indicating how long to wait before making a new request. + if ms, _ := strconv.Atoi(resp.Header.Get("Retry-After")); ms > 0 { + sleepTime = time.Duration(ms) * time.Millisecond + } else { + sleepTime = time.Duration(1< 400 && resp.StatusCode < 500: + // Pilosa nodes sometimes return 400, we retry in that case. + // (400, 500): No need to retry in other 4xx cases. + return resp.StatusCode, nil, errors.New(strings.TrimSpace(buf.String())) + + case resp.StatusCode == 503: + // This indicates that Pilosa is not ready to service this request, + // typically during startup. In this case, it's ok to give Pilosa + // some time and try again. + sleepTime = time.Duration(1<= c.maxRetries { + // If the error here is nil, we still want to return an error because + // we've hit the max retries limit. If an error exists, wrap it. + errMsg := fmt.Sprintf("max retries (%d) exceeded", c.maxRetries) + if err == nil { + return resp.StatusCode, nil, errors.New(errMsg) + } + return resp.StatusCode, nil, errors.Wrap(err, errMsg) + } + // The client can continue retrying after it has reached the maxBackoff time. + if sleepTime > c.maxBackoff { + return resp.StatusCode, nil, errors.Wrapf(err, "max backoff (%s) time exceeded", c.maxBackoff) + } + retry++ + c.logger.Printf("request failed with: '%v' status: %d, retrying %d after %v ", err, resp.StatusCode, retry, sleepTime) + time.Sleep(sleepTime) + } + // Unreachable code +} + +// statusToNodeShardsForIndex finds the hosts which contains shards for the given index +func (c *Client) statusToNodeShardsForIndex(status Status, indexName string) (map[uint64]*pnet.URI, error) { + result := make(map[uint64]*pnet.URI) + if maxShard, ok := status.indexMaxShard[indexName]; ok { + for shard := 0; shard <= int(maxShard); shard++ { + fragmentNodes, err := c.fetchFragmentNodes(indexName, uint64(shard)) + if err != nil { + return nil, err + } + if len(fragmentNodes) == 0 { + return nil, ErrNoFragmentNodes + } + node := fragmentNodes[0] + uri := &pnet.URI{ + Host: node.Host, + Port: node.Port, + Scheme: node.Scheme, + } + + result[uint64(shard)] = uri + } + } else { + return nil, ErrNoShard + } + return result, nil +} + +func (c *Client) augmentHeaders(headers map[string]string) map[string]string { + if headers == nil { + headers = map[string]string{} + } + + // TODO: move the following block to NewClient once cluster-resize support branch is merged. + version := strings.TrimPrefix(Version, "v") + + headers["User-Agent"] = fmt.Sprintf("pilosa/client/%s", version) + return headers +} + +func (c *Client) TranslateRowKeys(field *Field, keys []string) ([]uint64, error) { + req := &internal.TranslateKeysRequest{ + Index: field.index.name, + Field: field.name, + Keys: keys, + } + return c.translateKeys(req) +} + +func (c *Client) TranslateColumnKeys(index *Index, keys []string) ([]uint64, error) { + // If a manual server URI override has been provided, there's no + // point in partitioning the translation request on the client + // because every request is going to be sent to the manual URI. + if c.manualServerURI != nil { + req := &internal.TranslateKeysRequest{ + Index: index.name, + Keys: keys, + } + return c.translateKeys(req) + } + + // Get the list of hosts from the server. + // TODO: it's not ideal to request the list of nodes from the server + // on every call to TranslateColumnKeys(), but if we cache that list + // on the client, we risk calculating the partition distribution based + // on a stale node list. This TODO is here to indicate that we may, + // in the future, want to remove the overhead of this status request. + status, err := c.Status() + if err != nil { + return nil, errors.Wrap(err, "getting cluster status") + } + + hosts := make([]pnet.URI, len(status.Nodes)) + for i, node := range status.Nodes { + hosts[i] = node.URI.URI() + } + + keysByNode := make(map[pnet.URI][]string, len(hosts)) + for _, key := range keys { + // TODO 256 is DefaultPartitionN in Pilosa. Eventually this will likely be an index configuration parameter. + partitionID := keyPartition(index.Name(), key, 256) + uri := c.partitionOwner(partitionID, hosts) + keysByNode[uri] = append(keysByNode[uri], key) + } + + eg := errgroup.Group{} + idsByNode := make(map[pnet.URI][]uint64, len(keysByNode)) + ibnLock := &sync.Mutex{} + for uri, keys := range keysByNode { + uri := uri + keys := keys + eg.Go(func() error { + req := &internal.TranslateKeysRequest{ + Index: index.name, + Keys: keys, + } + ids, err := c.translateKeys(req, uri) + if err != nil { + return errors.Wrapf(err, "translating column keys at %v", uri) + } + ibnLock.Lock() + idsByNode[uri] = ids + ibnLock.Unlock() + return nil + }) + } + if err := eg.Wait(); err != nil { + return nil, err + } + + finalIDs := make([]uint64, len(keys)) + // put the ids back together into one slice + for uri, uriIds := range idsByNode { + uriKeys := keysByNode[uri] + kidx := 0 + for i, key := range uriKeys { + for ; keys[kidx] != key; kidx++ { + } + finalIDs[kidx] = uriIds[i] + kidx++ + } + } + return finalIDs, nil +} + +func (c *Client) partitionOwner(partitionID int, hosts []pnet.URI) pnet.URI { + nodeIndex := c.Hasher.Hash(uint64(partitionID), len(hosts)) + return hosts[nodeIndex] +} + +// Hasher represents an interface to hash integers into buckets. +type Hasher interface { + // Hashes the key into a number between [0,N). + Hash(key uint64, n int) int +} + +// jmphasher represents an implementation of jmphash. Implements Hasher. +type jmphasher struct{} + +// Hash returns the integer hash for the given key. +func (h *jmphasher) Hash(key uint64, n int) int { + b, j := int64(-1), int64(0) + for j < int64(n) { + b = j + key = key*uint64(2862933555777941757) + 1 + j = int64(float64(b+1) * (float64(int64(1)<<31) / float64((key>>33)+1))) + } + return int(b) +} + +func keyPartition(index, key string, partitionN int) int { + // Hash the bytes and mod by partition count. + h := fnv.New64a() + _, _ = h.Write([]byte(index)) + _, _ = h.Write([]byte(key)) + return int(h.Sum64() % uint64(partitionN)) +} + +func (c *Client) translateKeys(req *internal.TranslateKeysRequest, uris ...pnet.URI) ([]uint64, error) { + if len(req.Keys) == 0 { + return []uint64{}, nil + } + reqData, err := proto.Marshal(req) + if err != nil { + return nil, errors.Wrap(err, "marshalling translate keys request") + } + + var respData []byte + if len(uris) == 0 { + if _, respData, err = c.httpRequest("POST", "/internal/translate/keys", reqData, defaultProtobufHeaders(), true); err != nil { + return nil, err + } + } else { + if _, respData, err = c.doRequest(&uris[0], "POST", "/internal/translate/keys", defaultProtobufHeaders(), reqData); err != nil { + return nil, errors.Wrapf(err, "reading response body of /internal/translate/keys request to %v", uris[0]) + } + } + + idsResp := &internal.TranslateKeysResponse{} + err = proto.Unmarshal(respData, idsResp) + if err != nil { + return nil, errors.Wrap(err, "unmarshalling traslate keys response") + } + return idsResp.IDs, nil +} + +type TransactionResponse struct { + Transaction *pilosa.Transaction `json:"transaction,omitempty"` + Error string `json:"error,omitempty"` +} + +// StartTransaction tries to start a new transaction in Pilosa. It +// will continue trying until at least requestTimeout time has +// passed. If it fails due to an exclusive transaction already +// existing, it will return that transaction along with a non-nil +// error. +func (c *Client) StartTransaction(id string, timeout time.Duration, exclusive bool, requestTimeout time.Duration) (*pilosa.Transaction, error) { + return c.startTransaction(id, timeout, exclusive, time.Now().Add(requestTimeout)) +} + +func (c *Client) startTransaction(id string, timeout time.Duration, exclusive bool, deadline time.Time) (*pilosa.Transaction, error) { + trns := pilosa.Transaction{ + ID: id, + Timeout: timeout, + Exclusive: exclusive, + } + bod, err := json.Marshal(&trns) + if err != nil { + return nil, errors.Wrap(err, "marshalling transaction") + } + + status, data, err := c.httpRequest("POST", "/transaction", bod, defaultJSONHeaders(), true) + if status == http.StatusConflict && time.Now().Before(deadline) { + // if we're getting StatusConflict after all the usual timeouts/retries, keep retrying until the deadline + time.Sleep(time.Second) + return c.startTransaction(id, timeout, exclusive, deadline) + } + if err != nil { + return nil, err + } + + tr := &TransactionResponse{} + uerr := json.Unmarshal(data, &tr) + if uerr != nil { + if err != nil { + return nil, errors.Wrap(err, "unmarshal failed after") + } + return nil, errors.Wrap(uerr, "couldn't decode body") + } + + if tr.Error != "" { + err = errors.New(tr.Error) + } + + return tr.Transaction, err +} + +func (c *Client) FinishTransaction(id string) (*pilosa.Transaction, error) { + _, data, err := c.httpRequest("POST", "/transaction/"+id+"/finish", nil, defaultJSONHeaders(), true) + if err != nil && len(data) == 0 { + return nil, err + } + + tr := &TransactionResponse{} + uerr := json.Unmarshal(data, &tr) + if uerr != nil { + if err != nil { + return nil, errors.Wrap(err, "unmarshal failed after") + } + return nil, errors.Wrap(uerr, "couldn't decode body") + } + + if tr.Error != "" { + err = errors.New(tr.Error) + } + + return tr.Transaction, err +} + +func (c *Client) Transactions() (map[string]*pilosa.Transaction, error) { + _, respData, err := c.httpRequest("GET", "/transactions", nil, defaultJSONHeaders(), true) + if err != nil { + return nil, errors.Wrap(err, "getting transactions") + } + + trnsMap := make(map[string]*pilosa.Transaction) + err = json.Unmarshal(respData, &trnsMap) + if err != nil { + return nil, errors.Wrap(err, "unmarshalling transactions") + } + return trnsMap, nil +} + +func (c *Client) GetTransaction(id string) (*pilosa.Transaction, error) { + _, data, err := c.httpRequest("GET", "/transaction/"+id, nil, defaultJSONHeaders(), true) + if err != nil { + return nil, err + } + + tr := &TransactionResponse{} + uerr := json.Unmarshal(data, &tr) + if uerr != nil { + if err != nil { + return nil, errors.Wrap(err, "unmarshal failed after") + } + return nil, errors.Wrap(uerr, "couldn't decode body") + } + + if tr.Error != "" { + err = errors.New(tr.Error) + } + + return tr.Transaction, err +} + +func defaultProtobufHeaders() map[string]string { + return map[string]string{ + "Content-Type": "application/x-protobuf", + "Accept": "application/x-protobuf", + "PQL-Version": PQLVersion, + } +} + +func defaultJSONHeaders() map[string]string { + return map[string]string{ + "Content-Type": "application/json", + "Accept": "application/json", + "PQL-Version": PQLVersion, + } +} + +func buildRequest(host *pnet.URI, method, path string, headers map[string]string, data []byte) (*http.Request, error) { + request, err := http.NewRequest(method, host.Normalize()+path, bytes.NewReader(data)) + if err != nil { + return nil, err + } + + for k, v := range headers { + request.Header.Set(k, v) + } + + return request, nil +} + +func newHTTPClient(options *ClientOptions) *http.Client { + transport := &http.Transport{ + Dial: (&net.Dialer{ + Timeout: options.ConnectTimeout, + }).Dial, + TLSClientConfig: options.TLSConfig, + MaxIdleConnsPerHost: options.PoolSizePerRoute, + MaxIdleConns: options.TotalPoolSize, + } + return &http.Client{ + Transport: transport, + Timeout: options.SocketTimeout, + } +} + +func makeRequestData(query string, options *QueryOptions) ([]byte, error) { + request := &internal.QueryRequest{ + Query: query, + Shards: options.Shards, + ColumnAttrs: options.ColumnAttrs, + ExcludeRowAttrs: options.ExcludeRowAttrs, + ExcludeColumns: options.ExcludeColumns, + } + r, err := proto.Marshal(request) + if err != nil { + return nil, errors.Wrap(err, "marshaling request to protobuf") + } + return r, nil +} + +func makeRoaringImportPath(field *Field, shard uint64, params url.Values) string { + return fmt.Sprintf("/index/%s/field/%s/import-roaring/%d?%s", + field.index.name, field.name, shard, params.Encode()) +} + +type viewImports map[string]*roaring.Bitmap + +// ClientOptions control the properties of client connection to the server. +type ClientOptions struct { + SocketTimeout time.Duration + ConnectTimeout time.Duration + PoolSizePerRoute int + TotalPoolSize int + TLSConfig *tls.Config + manualServerAddress bool + tracer opentracing.Tracer + retries *int + stats stats.StatsClient + nat map[pnet.URI]pnet.URI +} + +func (co *ClientOptions) addOptions(options ...ClientOption) error { + for _, option := range options { + err := option(co) + if err != nil { + return err + } + } + return nil +} + +// ClientOption is used when creating a PilosaClient struct. +type ClientOption func(options *ClientOptions) error + +// OptClientSocketTimeout is the maximum idle socket time in nanoseconds +func OptClientSocketTimeout(timeout time.Duration) ClientOption { + return func(options *ClientOptions) error { + options.SocketTimeout = timeout + return nil + } +} + +// OptClientConnectTimeout is the maximum time to connect in nanoseconds. +func OptClientConnectTimeout(timeout time.Duration) ClientOption { + return func(options *ClientOptions) error { + options.ConnectTimeout = timeout + return nil + } +} + +// OptClientPoolSizePerRoute is the maximum number of active connections in the pool to a host. +func OptClientPoolSizePerRoute(size int) ClientOption { + return func(options *ClientOptions) error { + options.PoolSizePerRoute = size + return nil + } +} + +// OptClientTotalPoolSize is the maximum number of connections in the pool. +func OptClientTotalPoolSize(size int) ClientOption { + return func(options *ClientOptions) error { + options.TotalPoolSize = size + return nil + } +} + +// OptClientTLSConfig contains the TLS configuration. +func OptClientTLSConfig(config *tls.Config) ClientOption { + return func(options *ClientOptions) error { + options.TLSConfig = config + return nil + } +} + +// OptClientManualServerAddress forces the client use only the manual server address +func OptClientManualServerAddress(enabled bool) ClientOption { + return func(options *ClientOptions) error { + options.manualServerAddress = enabled + return nil + } +} + +// OptClientTracer sets the Open Tracing tracer +// See: https://opentracing.io +func OptClientTracer(tracer opentracing.Tracer) ClientOption { + return func(options *ClientOptions) error { + options.tracer = tracer + return nil + } +} + +// OptClientRetries sets the number of retries on HTTP request failures. +func OptClientRetries(retries int) ClientOption { + return func(options *ClientOptions) error { + if retries < 0 { + return errors.New("retries must be non-negative") + } + options.retries = &retries + return nil + } +} + +// OptClientStatsClient sets a stats client, such as Prometheus +func OptClientStatsClient(stats stats.StatsClient) ClientOption { + return func(options *ClientOptions) error { + options.stats = stats + return nil + } +} + +// OptClientNAT sets a NAT map used to translate the advertised URI to something +// else (for example, when accessing pilosa running in docker). +func OptClientNAT(nat map[string]string) ClientOption { + return func(options *ClientOptions) error { + // covert the strings to URIs + m := make(map[pnet.URI]pnet.URI) + for k, v := range nat { + if kuri, err := pnet.NewURIFromAddress(k); err != nil { + return errors.Wrapf(err, "converting string to URI: %s", k) + } else if vuri, err := pnet.NewURIFromAddress(v); err != nil { + return errors.Wrapf(err, "converting string to URI: %s", v) + } else { + m[*kuri] = *vuri + } + } + options.nat = m + return nil + } +} + +func (co *ClientOptions) withDefaults() (updated *ClientOptions) { + // copy options so the original is not updated + updated = &ClientOptions{} + *updated = *co + // impose defaults + if updated.SocketTimeout <= 0 { + updated.SocketTimeout = time.Second * 300 + } + if updated.ConnectTimeout <= 0 { + updated.ConnectTimeout = time.Second * 60 + } + if updated.PoolSizePerRoute <= 0 { + updated.PoolSizePerRoute = 50 + } + if updated.TotalPoolSize <= 0 { + updated.TotalPoolSize = 500 + } + if updated.TLSConfig == nil { + updated.TLSConfig = &tls.Config{} + } + if updated.retries == nil { + retries := 2 + updated.retries = &retries + } + return +} + +// QueryOptions contains options to customize the Query function. +type QueryOptions struct { + // Shards restricts query to a subset of shards. Queries all shards if nil. + Shards []uint64 + // ColumnAttrs enables returning columns in the query response. + ColumnAttrs bool + // ExcludeRowAttrs inhibits returning attributes + ExcludeRowAttrs bool + // ExcludeColumns inhibits returning columns + ExcludeColumns bool +} + +func (qo *QueryOptions) addOptions(options ...interface{}) error { + for i, option := range options { + switch o := option.(type) { + case nil: + if i != 0 { + return ErrInvalidQueryOption + } + continue + case *QueryOptions: + if i != 0 { + return ErrInvalidQueryOption + } + *qo = *o + case QueryOption: + err := o(qo) + if err != nil { + return err + } + default: + return ErrInvalidQueryOption + } + } + return nil +} + +// QueryOption is used when using options with a client.Query, +type QueryOption func(options *QueryOptions) error + +// OptQueryColumnAttrs enables returning column attributes in the result. +func OptQueryColumnAttrs(enable bool) QueryOption { + return func(options *QueryOptions) error { + options.ColumnAttrs = enable + return nil + } +} + +// OptQueryShards restricts the set of shards on which a query operates. +func OptQueryShards(shards ...uint64) QueryOption { + return func(options *QueryOptions) error { + options.Shards = append(options.Shards, shards...) + return nil + } +} + +// OptQueryExcludeAttrs enables discarding attributes from a result, +func OptQueryExcludeAttrs(enable bool) QueryOption { + return func(options *QueryOptions) error { + options.ExcludeRowAttrs = enable + return nil + } +} + +// OptQueryExcludeColumns enables discarding columns from a result, +func OptQueryExcludeColumns(enable bool) QueryOption { + return func(options *QueryOptions) error { + options.ExcludeColumns = enable + return nil + } +} + +// ImportOptions are the options for controlling the importer +type ImportOptions struct { + threadCount int + batchSize int + wantRoaring *bool + clear bool + skipSort bool +} + +// ImportOption is used when running imports. +type ImportOption func(options *ImportOptions) error + +// OptImportThreadCount is the number of goroutines allocated for import. +func OptImportThreadCount(count int) ImportOption { + return func(options *ImportOptions) error { + options.threadCount = count + return nil + } +} + +// OptImportBatchSize is the number of records read before importing them. +func OptImportBatchSize(batchSize int) ImportOption { + return func(options *ImportOptions) error { + options.batchSize = batchSize + return nil + } +} + +// OptImportClear sets clear import, which clears bits instead of setting them. +func OptImportClear(clear bool) ImportOption { + return func(options *ImportOptions) error { + options.clear = clear + return nil + } +} + +// OptImportRoaring enables importing using roaring bitmaps which is more performant. +func OptImportRoaring(enable bool) ImportOption { + return func(options *ImportOptions) error { + options.wantRoaring = &enable + return nil + } +} + +// OptImportSort tells the importer whether or not to sort batches of records, on +// by default. Sorting imposes some performance cost, especially on data that's +// already sorted, but dramatically improves performance in pathological +// cases. It is enabled by default because the pathological cases are awful, +// and the performance hit is comparatively small, but the performance cost can +// be significant if you know your data is sorted. +func OptImportSort(sorting bool) ImportOption { + return func(options *ImportOptions) error { + // skipSort is expressed negatively because we want to + // keep sorting enabled by default, so the zero value should + // be that default behavior. The client option expresses it + // positively because that's easier for API users. + options.skipSort = !sorting + return nil + } +} + +type fragmentNodeRoot struct { + URI fragmentNode `json:"uri"` +} + +type fragmentNode struct { + Scheme string `json:"scheme"` + Host string `json:"host"` + Port uint16 `json:"port"` +} + +func newFragmentNodeFromURI(uri *pnet.URI) fragmentNode { + return fragmentNode{ + Scheme: uri.Scheme, + Host: uri.Host, + Port: uri.Port, + } +} + +func (node fragmentNode) URI() *pnet.URI { + return &pnet.URI{ + Scheme: node.Scheme, + Host: node.Host, + Port: node.Port, + } +} + +// Info contains the configuration/host information from a Pilosa server. +type Info struct { + ShardWidth uint64 `json:"shardWidth"` // width of each shard + Memory uint64 `json:"memory"` // approximate host physical memory + CPUType string `json:"cpuType"` // "brand name string" from cpuid + CPUPhysicalCores int `json:"CPUPhysicalCores"` // physical cores (cpuid) + CPULogicalCores int `json:"CPULogicalCores"` // logical cores cpuid + CPUMHz uint64 `json:"CPUMHz"` // estimated clock speed +} + +// Status contains the status information from a Pilosa server. +type Status struct { + Nodes []StatusNode `json:"nodes"` + State string `json:"state"` + LocalID string `json:"localID"` + indexMaxShard map[string]uint64 +} + +// StatusNode contains information about a node in the cluster. +type StatusNode struct { + ID string `json:"id"` + URI StatusURI `json:"uri"` + IsPrimary bool `json:"isPrimary"` +} + +// StatusURI contains node information. +type StatusURI struct { + Scheme string `json:"scheme"` + Host string `json:"host"` + Port uint16 `json:"port"` +} + +// URI returns the StatusURI as a URI. +func (s StatusURI) URI() pnet.URI { + return pnet.URI{ + Scheme: s.Scheme, + Host: s.Host, + Port: s.Port, + } +} + +// SchemaInfo contains the indexes. +type SchemaInfo struct { + Indexes []SchemaIndex `json:"indexes"` +} + +// SchemaIndex contains index information. +type SchemaIndex struct { + Name string `json:"name"` + CreatedAt int64 `json:"createdAt,omitempty"` + Options SchemaOptions `json:"options"` + Fields []SchemaField `json:"fields"` + Shards []uint64 `json:"shards"` + ShardWidth uint64 `json:"shardWidth"` +} + +// SchemaField contains field information. +type SchemaField struct { + Name string `json:"name"` + CreatedAt int64 `json:"createdAt,omitempty"` + Options SchemaOptions `json:"options"` +} + +// SchemaOptions contains options for a field or an index. +type SchemaOptions struct { + FieldType FieldType `json:"type"` + CacheType string `json:"cacheType"` + CacheSize uint `json:"cacheSize"` + TimeQuantum string `json:"timeQuantum"` + Min pql.Decimal `json:"min"` + Max pql.Decimal `json:"max"` + Scale int64 `json:"scale"` + Keys bool `json:"keys"` + NoStandardView bool `json:"noStandardView"` + TrackExistence bool `json:"trackExistence"` +} + +func (so SchemaOptions) asIndexOptions() *IndexOptions { + return &IndexOptions{ + keys: so.Keys, + keysSet: true, + trackExistence: so.TrackExistence, + trackExistenceSet: true, + } +} + +func (so SchemaOptions) asFieldOptions() *FieldOptions { + return &FieldOptions{ + fieldType: so.FieldType, + cacheSize: int(so.CacheSize), + cacheType: CacheType(so.CacheType), + timeQuantum: TimeQuantum(so.TimeQuantum), + min: so.Min, + max: so.Max, + scale: so.Scale, + keys: so.Keys, + noStandardView: so.NoStandardView, + } +} + +type exportReader struct { + client *Client + shardURIs map[uint64]*pnet.URI + field *Field + body []byte + bodyIndex int + currentShard uint64 + shardCount uint64 +} + +func newExportReader(client *Client, shardURIs map[uint64]*pnet.URI, field *Field) *exportReader { + return &exportReader{ + client: client, + shardURIs: shardURIs, + field: field, + shardCount: uint64(len(shardURIs)), + } +} + +// Read updates the passed array with the exported CSV data and returns the number of bytes read +func (r *exportReader) Read(p []byte) (n int, err error) { + if r.currentShard >= r.shardCount { + err = io.EOF + return + } + if r.body == nil { + uri := r.shardURIs[r.currentShard] + headers := map[string]string{ + "Accept": "text/csv", + } + path := fmt.Sprintf("/export?index=%s&field=%s&shard=%d", + r.field.index.Name(), r.field.Name(), r.currentShard) + _, respData, err := r.client.doRequest(uri, "GET", path, headers, nil) + if err != nil { + return 0, errors.Wrap(err, "doing export request") + } + r.body = respData + r.bodyIndex = 0 + } + n = copy(p, r.body[r.bodyIndex:]) + r.bodyIndex += n + if n >= len(r.body) { + r.body = nil + r.currentShard++ + } + return +} diff --git a/client/client_internal_it_test.go b/client/client_internal_it_test.go new file mode 100644 index 000000000..8c88384af --- /dev/null +++ b/client/client_internal_it_test.go @@ -0,0 +1,85 @@ +//+build integration + +// Copyright 2017 Pilosa Corp. +// +// 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 client + +import ( + "reflect" + "testing" + + pnet "github.com/pilosa/pilosa/v2/net" +) + +func TestNewClientFromAddresses(t *testing.T) { + cases := []struct { + Name string + Hosts []string + ExpectErr bool + ExpectedHosts []pnet.URI + }{ + { + Name: "Cluster", + Hosts: []string{":10101", "node0.pilosa.com:10101", "node2.pilosa.com"}, + ExpectedHosts: []pnet.URI{ + {Scheme: "http", Port: 10101, Host: "localhost"}, + {Scheme: "http", Port: 10101, Host: "node0.pilosa.com"}, + {Scheme: "http", Port: 10101, Host: "node2.pilosa.com"}, + }, + }, + { + Name: "URIParseError", + Hosts: []string{"://"}, + ExpectErr: true, + }, + { + Name: "Empty", + Hosts: []string{}, + ExpectedHosts: []pnet.URI{}, + }, + { + Name: "nil", + ExpectedHosts: []pnet.URI{}, + }, + } + + for _, c := range cases { + c := c + t.Run(c.Name, func(t *testing.T) { + cli, err := NewClient(c.Hosts) + if c.ExpectErr { + if err == nil { + t.Fatalf("Did not get expected error when creating client: %v", cli.cluster.Hosts()) + } + } else { + if err != nil { + t.Fatalf("Creating client from addresses: %v", err) + } + if actualHosts := cli.cluster.Hosts(); !reflect.DeepEqual(actualHosts, c.ExpectedHosts) { + t.Fatalf("Unexpected hosts in client's cluster, got: %v, expected: %v", actualHosts, c.ExpectedHosts) + } + } + }) + } +} + +func TestDetectClusterChanges(t *testing.T) { + c := getClient() + defer c.Close() + c.shardNodes.data["blah"] = make(map[uint64][]*pnet.URI) + c.shardNodes.data["blah"][1] = []*pnet.URI{{Scheme: "zzz"}} + + c.detectClusterChanges() +} diff --git a/client/client_it_test.go b/client/client_it_test.go new file mode 100644 index 000000000..8f26bc6ed --- /dev/null +++ b/client/client_it_test.go @@ -0,0 +1,2058 @@ +//+build integration + +// Copyright 2017 Pilosa Corp. +// +// 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 client + +import ( + "bytes" + "crypto/tls" + "fmt" + "io" + "io/ioutil" + "net/http" + "net/http/httptest" + "os" + "reflect" + "strconv" + "strings" + "sync" + "testing" + "time" + + "github.com/golang/protobuf/proto" //nolint:staticcheck + "github.com/pilosa/pilosa/v2/internal" + pnet "github.com/pilosa/pilosa/v2/net" + "github.com/pkg/errors" +) + +var schema = NewSchema() +var index *Index +var indexName = "go-testindex" +var keysIndex *Index +var schemaTestIndex *Index +var testField *Field + +// AtomicRecord import test stuff +var indexAR *Index +var indexARname = "i" +var fieldAcct0 = "acct0" +var fieldAcct1 = "acct1" +var field0 *Field +var field1 *Field + +func TestMain(m *testing.M) { + Setup() + r := m.Run() + TearDown() + os.Exit(r) +} + +func Setup() { + client := getClient() + + // Make sure the existing schema is empty. + if existingSchema, err := client.Schema(); err != nil { + panic(err) + } else if indexes := existingSchema.Indexes(); len(indexes) > 0 { + TearDown() + //panic(fmt.Sprintf("Pilosa data isn't clean, found indexes: %v", indexes)) + } + + testSchema := NewSchema() + index = testSchema.Index(indexName) + keysIndex = testSchema.Index("go-testinindex-keys", OptIndexKeys(true)) + schemaTestIndex = testSchema.Index("schema-test-index", + OptIndexKeys(true), + OptIndexTrackExistence(false)) + testField = index.Field("test-field") + + indexAR = testSchema.Index(indexARname) + field0 = indexAR.Field(fieldAcct0, OptFieldTypeInt(-1000, 1000)) + field1 = indexAR.Field(fieldAcct1, OptFieldTypeInt(-1000, 1000)) + + err := client.SyncSchema(testSchema) + if err != nil { + panic(err) + } + _ = client.Close() +} + +func TearDown() { + client := getClient() + if client == nil { + return + } + defer client.Close() + if err := client.DeleteIndex(index); err != nil { + panic(err) + } + if err := client.DeleteIndex(indexAR); err != nil { + panic(err) + } + + if err := client.DeleteIndex(keysIndex); err != nil { + panic(err) + } + if err := client.DeleteIndex(schemaTestIndex); err != nil { + panic(err) + } +} + +func Reset() { + TearDown() + Setup() +} + +func TestCreateDefaultClient(t *testing.T) { + client := DefaultClient() + if client == nil { + t.Fatal() + } +} + +func TestClientReturnsResponse(t *testing.T) { + client := getClient() + defer client.Close() + response, err := client.Query(testField.Row(1)) + if err != nil { + t.Fatalf("Error querying: %s", err) + } + if response == nil { + t.Fatalf("Response should not be nil") + } +} + +func TestQueryWithShards(t *testing.T) { + Reset() + const shardWidth = 1048576 + client := getClient() + defer client.Close() + if _, err := client.Query(testField.Set(1, 100)); err != nil { + t.Fatal(err) + } + if _, err := client.Query(testField.Set(1, shardWidth)); err != nil { + t.Fatal(err) + } + if _, err := client.Query(testField.Set(1, shardWidth*3)); err != nil { + t.Fatal(err) + } + + response, err := client.Query(testField.Row(1), OptQueryShards(0, 3)) + if err != nil { + t.Fatal(err) + } + if columns := response.Result().Row().Columns; !reflect.DeepEqual(columns, []uint64{100, shardWidth * 3}) { + t.Fatalf("Unexpected results: %#v", columns) + } +} + +func TestQueryWithColumns(t *testing.T) { + Reset() + client := getClient() + defer client.Close() + targetAttrs := map[string]interface{}{ + "name": "some string", + "age": int64(95), + "registered": true, + "height": 1.83, + } + _, err := client.Query(testField.Set(1, 100)) + if err != nil { + t.Fatal(err) + } + response, err := client.Query(index.SetColumnAttrs(100, targetAttrs)) + if err != nil { + t.Fatal(err) + } + if !reflect.DeepEqual(response.Column(), ColumnItem{}) { + t.Fatalf("No columns should be returned if it wasn't explicitly requested") + } + response, err = client.Query(testField.Row(1), &QueryOptions{ColumnAttrs: true}) + if err != nil { + t.Fatal(err) + } + if len(response.ColumnAttrs()) != 1 { + t.Fatalf("Column count should be == 1") + } + columns := response.Columns() + if len(columns) != 1 { + t.Fatalf("Column count should be == 1") + } + if columns[0].ID != 100 { + t.Fatalf("Column ID should be == 100") + } + if !reflect.DeepEqual(columns[0].Attributes, targetAttrs) { + t.Fatalf("Column attrs does not match") + } + + if !reflect.DeepEqual(response.Column(), columns[0]) { + t.Fatalf("Columns() should be equivalent to first column in the response") + } +} + +func TestSetRowAttrs(t *testing.T) { + Reset() + client := getClient() + defer client.Close() + targetAttrs := map[string]interface{}{ + "name": "some string", + "age": int64(95), + "registered": true, + "height": 1.83, + } + _, err := client.Query(testField.Set(1, 100)) + if err != nil { + t.Fatal(err) + } + _, err = client.Query(testField.SetRowAttrs(1, targetAttrs)) + if err != nil { + t.Fatal(err) + } + response, err := client.Query(testField.Row(1), &QueryOptions{ColumnAttrs: true}) + if err != nil { + t.Fatal(err) + } + if !reflect.DeepEqual(targetAttrs, response.Result().Row().Attributes) { + t.Fatalf("Row attributes should be set") + } +} + +func TestOrmCount(t *testing.T) { + client := getClient() + defer client.Close() + countField := index.Field("count-test") + err := client.EnsureField(countField) + if err != nil { + t.Fatal(err) + } + qry := index.BatchQuery( + countField.Set(10, 20), + countField.Set(10, 21), + countField.Set(15, 25), + ) + _, err = client.Query(qry) + if err != nil { + t.Fatal(err) + } + response, err := client.Query(index.Count(countField.Row(10))) + if err != nil { + t.Fatal(err) + } + if response.Result().Count() != 2 { + t.Fatalf("Count should be 2") + } +} + +func TestDecimalField(t *testing.T) { + client := getClient() + defer client.Close() + decField := index.Field("a-decimal", OptFieldTypeDecimal(3)) + err := client.EnsureField(decField) + if err != nil { + t.Fatal(err) + } + + sch, err := client.Schema() + if err != nil { + t.Fatalf("getting schema: %v", err) + } + idx := sch.indexes["go-testindex"] + if opts := idx.Field("a-decimal").Options(); opts.scale != 3 { + t.Fatalf("scale should be 3, but: %v", opts) + } + +} + +func TestIntersectReturns(t *testing.T) { + client := getClient() + defer client.Close() + field := index.Field("segments") + err := client.EnsureField(field) + if err != nil { + t.Fatal(err) + } + qry1 := index.BatchQuery( + field.Set(2, 10), + field.Set(2, 15), + field.Set(3, 10), + field.Set(3, 20), + ) + _, err = client.Query(qry1) + if err != nil { + t.Fatal(err) + } + + qry2 := index.Intersect(field.Row(2), field.Row(3)) + response, err := client.Query(qry2) + if err != nil { + t.Fatal(err) + } + if len(response.Results()) != 1 { + t.Fatal("There must be 1 result") + } + if !reflect.DeepEqual(response.Result().Row().Columns, []uint64{10}) { + t.Fatal("Returned columns must be: [10]") + } +} + +func TestTopNReturns(t *testing.T) { + client := getClient() + defer client.Close() + field := index.Field("topn_test") + err := client.EnsureField(field) + if err != nil { + t.Fatal(err) + } + qry := index.BatchQuery( + field.Set(10, 5), + field.Set(10, 10), + field.Set(10, 15), + field.Set(20, 5), + field.Set(30, 5), + ) + _, err = client.Query(qry) + if err != nil { + t.Fatal(err) + } + + // XXX: The following is required to make this test pass. See: https://github.com/pilosa/pilosa/issues/625 + _, _, err = client.HTTPRequest("POST", "/recalculate-caches", nil, nil) + if err != nil { + t.Fatal(err) + } + + response, err := client.Query(field.TopN(2)) + if err != nil { + t.Fatal(err) + } + items := response.Result().CountItems() + if len(items) != 2 { + t.Fatalf("There should be 2 count items: %v", items) + } + item := items[0] + if item.ID != 10 { + t.Fatalf("Item[0] ID should be 10") + } + if item.Count != 3 { + t.Fatalf("Item[0] Count should be 3") + } + + _, err = client.Query(field.SetRowAttrs(10, map[string]interface{}{"foo": "bar"})) + if err != nil { + t.Fatal(err) + } + response, err = client.Query(field.FilterAttrTopN(5, nil, "foo", "bar")) + if err != nil { + t.Fatal(err) + } + items = response.Result().CountItems() + if len(items) != 1 { + t.Fatalf("There should be 1 count item: %v", items) + } + item = items[0] + if item.ID != 10 { + t.Fatalf("Item[0] ID should be 10") + } + if item.Count != 3 { + t.Fatalf("Item[0] Count should be 3") + } +} + +func TestMinMaxRow(t *testing.T) { + client := getClient() + defer client.Close() + field := index.Field("test-minmaxrow-field") + err := client.EnsureField(field) + if err != nil { + t.Fatal(err) + } + qry := index.BatchQuery( + field.Set(10, 5), + field.Set(10, 10), + field.Set(10, 15), + field.Set(20, 5), + field.Set(30, 5), + ) + _, err = client.Query(qry) + if err != nil { + t.Fatalf("error setting bits: %v", err) + } + + response, err := client.Query(field.MinRow()) + if err != nil { + t.Fatalf("error executing min: %v", err) + } + min := response.Result().CountItem().ID + response, err = client.Query(field.MaxRow()) + if err != nil { + t.Fatalf("error executing max: %v", err) + } + max := response.Result().CountItem().ID + + if min != 10 { + t.Fatalf("Min should be 10, got %v instead", min) + } + if max != 30 { + t.Fatalf("Max should be 30, got %v instead", max) + } +} + +func TestSetMutexField(t *testing.T) { + client := getClient() + defer client.Close() + field := index.Field("mutex-test", OptFieldTypeMutex(CacheTypeDefault, 0)) + err := client.EnsureField(field) + if err != nil { + t.Fatal(err) + } + + // can set mutex + _, err = client.Query(field.Set(1, 100)) + if err != nil { + t.Fatal(err) + } + response, err := client.Query(field.Row(1)) + if err != nil { + t.Fatal(err) + } + target := []uint64{100} + if !reflect.DeepEqual(target, response.Result().Row().Columns) { + t.Fatalf("%v != %v", target, response.Result().Row().Columns) + } + + // setting another row removes the previous + _, err = client.Query(field.Set(42, 100)) + if err != nil { + t.Fatal(err) + } + response, err = client.Query(index.BatchQuery( + field.Row(1), + field.Row(42), + )) + if err != nil { + t.Fatal(err) + } + target1 := []uint64(nil) + target42 := []uint64{100} + if !reflect.DeepEqual(target1, response.Results()[0].Row().Columns) { + t.Fatalf("%#v != %#v", target1, response.Results()[0].Row().Columns) + } + if !reflect.DeepEqual(target42, response.Results()[1].Row().Columns) { + t.Fatalf("%#v != %#v", target42, response.Results()[1].Row().Columns) + } +} + +func TestSetBoolField(t *testing.T) { + client := getClient() + defer client.Close() + field := index.Field("bool-test", OptFieldTypeBool()) + err := client.EnsureField(field) + if err != nil { + t.Fatal(err) + } + + // can set bool + _, err = client.Query(field.Set(true, 100)) + if err != nil { + t.Fatal(err) + } + response, err := client.Query(field.Row(true)) + if err != nil { + t.Fatal(err) + } + target := []uint64{100} + if !reflect.DeepEqual(target, response.Result().Row().Columns) { + t.Fatalf("%v != %v", target, response.Result().Row().Columns) + } +} + +func TestClearRowQuery(t *testing.T) { + client := getClient() + defer client.Close() + field := index.Field("clear-row-test") + err := client.EnsureField(field) + if err != nil { + t.Fatal(err) + } + + _, err = client.Query(index.BatchQuery( + field.Set(1, 100), + field.Set(1, 200), + )) + if err != nil { + t.Fatal(err) + } + response, err := client.Query(field.Row(1)) + if err != nil { + t.Fatal(err) + } + target := []uint64{100, 200} + if !reflect.DeepEqual(target, response.Result().Row().Columns) { + t.Fatalf("%v != %v", target, response.Result().Row().Columns) + } + + _, err = client.Query(field.ClearRow(1)) + if err != nil { + t.Fatal(err) + } + response, err = client.Query(field.Row(1)) + if err != nil { + t.Fatal(err) + } + target = []uint64(nil) + if !reflect.DeepEqual(target, response.Result().Row().Columns) { + t.Fatalf("%v != %v", target, response.Result().Row().Columns) + } +} + +func TestRowsQuery(t *testing.T) { + client := getClient() + defer client.Close() + field := index.Field("rows-test") + err := client.EnsureField(field) + if err != nil { + t.Fatal(err) + } + _, err = client.Query(index.BatchQuery( + field.Set(1, 100), + field.Set(1, 200), + field.Set(2, 200), + )) + if err != nil { + t.Fatal(err) + } + resp, err := client.Query(field.Rows()) + if err != nil { + t.Fatal(err) + } + target := RowIdentifiersResult{ + IDs: []uint64{1, 2}, + } + if !reflect.DeepEqual(target, resp.Result().RowIdentifiers()) { + t.Fatalf("%v != %v", target, resp.Result().RowIdentifiers()) + } +} + +func TestUnionRowsQuery(t *testing.T) { + client := getClient() + defer client.Close() + field := index.Field("rows-test") + err := client.EnsureField(field) + if err != nil { + t.Fatal(err) + } + _, err = client.Query(index.BatchQuery( + field.Set(1, 100), + field.Set(1, 200), + field.Set(2, 200), + )) + if err != nil { + t.Fatal(err) + } + resp, err := client.Query(field.Rows().Union()) + if err != nil { + t.Fatal(err) + } + target := []uint64{100, 200} + if !reflect.DeepEqual(target, resp.Result().Row().Columns) { + t.Fatalf("%v != %v", target, resp.Result().Row().Columns) + } +} + +func TestLikeQuery(t *testing.T) { + client := getClient() + defer client.Close() + field := index.Field("like-test", OptFieldKeys(true)) + err := client.EnsureField(field) + if err != nil { + t.Fatal(err) + } + _, err = client.Query(index.BatchQuery( + field.Set("a", 100), + field.Set("b", 200), + field.Set("bc", 200), + )) + if err != nil { + t.Fatal(err) + } + resp, err := client.Query(field.Like("b%")) + if err != nil { + t.Fatal(err) + } + target := RowIdentifiersResult{ + Keys: []string{"b", "bc"}, + } + if !reflect.DeepEqual(target, resp.Result().RowIdentifiers()) { + t.Fatalf("%v != %v", target, resp.Result().RowIdentifiers()) + } +} + +func TestGroupByQuery(t *testing.T) { + client := getClient() + defer client.Close() + field := index.Field("group-by-test") + err := client.EnsureField(field) + if err != nil { + t.Fatal(err) + } + _, err = client.Query(index.BatchQuery( + field.Set(1, 100), + field.Set(1, 200), + field.Set(2, 200), + )) + if err != nil { + t.Fatal(err) + } + resp, err := client.Query(index.GroupBy(field.Rows())) + if err != nil { + t.Fatal(err) + } + target := []GroupCount{ + {Groups: []FieldRow{{FieldName: "group-by-test", RowID: 1}}, Count: 2}, + {Groups: []FieldRow{{FieldName: "group-by-test", RowID: 2}}, Count: 1}, + } + + checkGroupBy(t, target, resp.Result().GroupCounts()) +} + +func TestGroupByIntQuery(t *testing.T) { + client := getClient() + defer client.Close() + field := index.Field("fint", OptFieldTypeInt(-10, 10)) + err := client.EnsureField(field) + if err != nil { + t.Fatal(err) + } + _, err = client.Query(index.RawQuery(` + Set(0, fint=1) + Set(1, fint=2) + + Set(2,fint=-2) + Set(3,fint=-1) + + Set(4,fint=4) + + Set(10, fint=0) + Set(100, fint=0) + Set(1000, fint=0) + Set(10000,fint=0) + Set(100000,fint=0) + `)) + if err != nil { + t.Fatal(err) + } + resp, err := client.Query(index.GroupBy(field.Rows())) + if err != nil { + t.Fatal(err) + } + var a, b, c, d, e, f int64 = -2, -1, 0, 1, 2, 4 + target := []GroupCount{ + {Groups: []FieldRow{{FieldName: "fint", Value: &a}}, Count: 1}, + {Groups: []FieldRow{{FieldName: "fint", Value: &b}}, Count: 1}, + {Groups: []FieldRow{{FieldName: "fint", Value: &c}}, Count: 5}, + {Groups: []FieldRow{{FieldName: "fint", Value: &d}}, Count: 1}, + {Groups: []FieldRow{{FieldName: "fint", Value: &e}}, Count: 1}, + {Groups: []FieldRow{{FieldName: "fint", Value: &f}}, Count: 1}, + } + + checkGroupBy(t, target, resp.Result().GroupCounts()) +} + +func checkGroupBy(t *testing.T, expected, results []GroupCount) { + t.Helper() + if len(results) != len(expected) { + t.Fatalf("number of groupings mismatch:\n got:%+v\nwant:%+v\n", results, expected) + } + for i, result := range results { + if !reflect.DeepEqual(expected[i], result) { + t.Fatalf("unexpected result at %d: \n got:%+v\nwant:%+v\n", i, result, expected[i]) + } + } +} + +func TestCreateDeleteIndexField(t *testing.T) { + client := getClient() + defer client.Close() + index1 := NewIndex("to-be-deleted") + field1 := index1.Field("foo") + err := client.CreateIndex(index1) + if err != nil { + t.Fatal(err) + } + err = client.CreateField(field1) + if err != nil { + t.Fatal(err) + } + err = client.DeleteField(field1) + if err != nil { + t.Fatal(err) + } + err = client.DeleteIndex(index1) + if err != nil { + t.Fatal(err) + } +} + +func TestEnsureIndexExists(t *testing.T) { + client := getClient() + defer client.Close() + err := client.EnsureIndex(index) + if err != nil { + t.Fatal(err) + } +} + +func TestEnsureFieldExists(t *testing.T) { + client := getClient() + defer client.Close() + err := client.EnsureField(testField) + if err != nil { + t.Fatal(err) + } +} + +func TestCreateFieldWithTimeQuantum(t *testing.T) { + client := getClient() + defer client.Close() + field := index.Field("field-with-timequantum", OptFieldTypeTime(TimeQuantumYear)) + err := client.CreateField(field) + if err != nil { + t.Fatal(err) + } +} + +func TestErrorCreatingIndex(t *testing.T) { + client := getClient() + defer client.Close() + err := client.CreateIndex(index) + if err == nil { + t.Fatal() + } +} + +func TestErrorCreatingField(t *testing.T) { + client := getClient() + defer client.Close() + err := client.CreateField(testField) + if err == nil { + t.Fatal() + } +} + +func TestIndexAlreadyExists(t *testing.T) { + client := getClient() + defer client.Close() + err := client.CreateIndex(index) + if err != ErrIndexExists { + t.Fatal(err) + } +} + +func TestQueryWithEmptyClusterFails(t *testing.T) { + client, _ := NewClient(DefaultCluster(), OptClientRetries(0)) + attrs := map[string]interface{}{"a": 1} + _, err := client.Query(index.SetColumnAttrs(0, attrs)) + if errors.Cause(err) != ErrEmptyCluster { + t.Fatal(err) + } +} + +func TestFailoverFail(t *testing.T) { + uri, _ := pnet.NewURIFromAddress("does-not-resolve.foo.bar") + cluster := NewClusterWithHost(uri, uri, uri, uri) + client, _ := NewClient(cluster, OptClientRetries(0)) + attrs := map[string]interface{}{"a": 1} + _, err := client.Query(index.SetColumnAttrs(0, attrs)) + if !strings.Contains(err.Error(), ErrTriedMaxHosts.Error()) { + t.Fatalf("ErrTriedMaxHosts error should be returned. Got: %v", err) + } +} + +func TestQueryFailsIfAddressNotResolved(t *testing.T) { + uri, _ := pnet.NewURIFromAddress("nonexisting.domain.pilosa.com:3456") + client, _ := NewClient(uri, OptClientRetries(0)) + _, err := client.Query(index.RawQuery("bar")) + if err == nil { + t.Fatal() + } +} + +func TestQueryFails(t *testing.T) { + client := getClient() + defer client.Close() + _, err := client.Query(index.RawQuery("Invalid query")) + if err == nil { + t.Fatal() + } +} + +func TestInvalidHttpRequest(t *testing.T) { + client := getClient() + defer client.Close() + _, _, err := client.HTTPRequest("INVALID METHOD", "/foo", nil, nil) + if err == nil { + t.Fatal() + } +} + +func TestErrorResponseNotRead(t *testing.T) { + server := getMockServer(500, []byte("Unknown error"), 512) + defer server.Close() + uri, err := pnet.NewURIFromAddress(server.URL) + if err != nil { + t.Fatal(err) + } + client, _ := NewClient(uri, OptClientRetries(0)) + response, err := client.Query(testField.Row(1)) + if err == nil { + t.Fatalf("Got response: %v", response) + } +} + +func TestResponseNotRead(t *testing.T) { + server := getMockServer(200, []byte("some content"), 512) + defer server.Close() + uri, err := pnet.NewURIFromAddress(server.URL) + if err != nil { + t.Fatal(err) + } + client, _ := NewClient(uri, OptClientRetries(0)) + response, err := client.Query(testField.Row(1)) + if err == nil { + t.Fatalf("Got response: %v", response) + } +} + +func TestSchema(t *testing.T) { + client := getClient() + defer client.Close() + schema, err := client.Schema() + if err != nil { + t.Fatal(err) + } + if len(schema.indexes) < 1 { + t.Fatalf("There should be at least 1 index in the schema") + } + f := schemaTestIndex.Field("schema-test-field", + OptFieldTypeSet(CacheTypeLRU, 9999), + OptFieldKeys(true), + ) + if f == nil { + t.Fatal("f should not be nil") + } + if err := client.EnsureField(f); err != nil { + t.Fatalf("ensuring field: %v", err) + } + err = client.SyncSchema(schema) + if err != nil { + t.Fatal(err) + } + schema, err = client.Schema() + if err != nil { + t.Fatal(err) + } + i2 := schema.indexes[schemaTestIndex.Name()] + if !reflect.DeepEqual(schemaTestIndex.options, i2.options) { + t.Fatalf("%v != %v", schemaTestIndex.options, i2.options) + } + + f2 := schema.indexes[schemaTestIndex.Name()].fields["schema-test-field"] + if f2 == nil { + t.Fatal("Field should not be nil") + } + if f2 != nil { // happy linter + opt := f2.options + if opt.cacheType != CacheTypeLRU { + t.Fatalf("cache type %s != %s", CacheTypeLRU, opt.cacheType) + } + if opt.cacheSize != 9999 { + t.Fatalf("cache size 9999 != %d", opt.cacheSize) + } + if !opt.keys { + t.Fatalf("keys true != %v", opt.keys) + } + if !reflect.DeepEqual(f.options, f2.options) { + t.Fatalf("%v != %v", f.options, f2.options) + } + } +} + +func TestSync(t *testing.T) { + client := getClient() + defer client.Close() + remoteIndex := NewIndex("remote-index-1") + err := client.EnsureIndex(remoteIndex) + if err != nil { + t.Fatal(err) + } + remoteField := remoteIndex.Field("remote-field-1") + err = client.EnsureField(remoteField) + if err != nil { + t.Fatal(err) + } + schema1 := NewSchema() + index11 := schema1.Index("diff-index1") + index11.Field("field1-1") + index11.Field("field1-2") + index12 := schema1.Index("diff-index2") + index12.Field("field2-1") + schema1.Index(remoteIndex.Name()) + + err = client.SyncSchema(schema1) + if err != nil { + t.Fatal(err) + } + err = client.DeleteIndex(remoteIndex) + if err != nil { + t.Fatal(err) + } + + err = client.DeleteIndex(index11) + if err != nil { + t.Fatal(err) + } + + err = client.DeleteIndex(index12) + if err != nil { + t.Fatal(err) + } +} + +func TestSyncFailure(t *testing.T) { + server := getMockServer(404, []byte("sorry, not found"), -1) + defer server.Close() + uri, err := pnet.NewURIFromAddress(server.URL) + if err != nil { + panic(err) + } + client, _ := NewClient(uri, OptClientRetries(0)) + err = client.SyncSchema(NewSchema()) + if err == nil { + t.Fatal("should have failed") + } +} + +func TestErrorRetrievingSchema(t *testing.T) { + server := getMockServer(404, []byte("sorry, not found"), -1) + defer server.Close() + uri, err := pnet.NewURIFromAddress(server.URL) + if err != nil { + panic(err) + } + client, _ := NewClient(uri, OptClientRetries(0)) + _, err = client.Schema() + if err == nil { + t.Fatal("should have failed") + } +} + +func TestExportReaderFailure(t *testing.T) { + server := getMockServer(404, []byte("sorry, not found"), -1) + defer server.Close() + uri, err := pnet.NewURIFromAddress(server.URL) + if err != nil { + panic(err) + } + field := index.Field("exportfield") + shardURIs := map[uint64]*pnet.URI{ + 0: uri, + } + client, _ := NewClient(uri, OptClientRetries(0)) + reader := newExportReader(client, shardURIs, field) + buf := make([]byte, 1000) + _, err = reader.Read(buf) + if err == nil { + t.Fatal("should have failed") + } +} + +func TestExportReaderReadBodyFailure(t *testing.T) { + server := getMockServer(200, []byte("not important"), 100) + defer server.Close() + uri, err := pnet.NewURIFromAddress(server.URL) + if err != nil { + t.Fatal(err) + } + field := index.Field("exportfield") + shardURIs := map[uint64]*pnet.URI{0: uri} + client, _ := NewClient(uri, OptClientRetries(0)) + reader := newExportReader(client, shardURIs, field) + buf := make([]byte, 1000) + _, err = reader.Read(buf) + if err == nil { + t.Fatal("should have failed") + } +} + +func TestFetchFragmentNodes(t *testing.T) { + client := getClient() + defer client.Close() + nodes, err := client.fetchFragmentNodes(index.Name(), 0) + if err != nil { + t.Fatal(err) + } + if len(nodes) != 1 { + t.Fatalf("1 node should be returned") + } + // running the same for coverage + nodes, err = client.fetchFragmentNodes(index.Name(), 0) + if err != nil { + t.Fatal(err) + } + if len(nodes) != 1 { + t.Fatalf("1 node should be returned") + } +} + +func TestFetchStatus(t *testing.T) { + client := getClient() + defer client.Close() + status, err := client.Status() + if err != nil { + t.Fatal(err) + } + if len(status.Nodes) == 0 { + t.Fatalf("There should be at least 1 host in the status") + } +} + +func TestFetchInfo(t *testing.T) { + client := getClient() + defer client.Close() + info, err := client.Info() + if err != nil { + t.Fatal(err) + } + if info.ShardWidth == 0 { + t.Fatalf("shard width should not be zero") + } + if info.Memory < (512 * 1024 * 1024) { + t.Fatalf("server memory [%d bytes] under 512MB seems highly improbable", info.Memory) + } + if info.CPUPhysicalCores < 1 || info.CPULogicalCores < 1 { + t.Fatalf("server did not detect any CPU cores") + } + if info.CPUType == "" { + t.Fatalf("server reported empty string for CPU type") + } + if info.CPUMHz == 0 { + t.Fatalf("server reported 0MHz processor") + } +} + +func TestRowRangeQuery(t *testing.T) { + client := getClient() + defer client.Close() + field := index.Field("test-rowrangefield", OptFieldTypeTime(TimeQuantumMonthDayHour)) + err := client.EnsureField(field) + if err != nil { + t.Fatal(err) + } + _, err = client.Query(index.BatchQuery( + field.SetTimestamp(10, 100, time.Date(2017, time.January, 1, 0, 0, 0, 0, time.UTC)), + field.SetTimestamp(10, 100, time.Date(2018, time.January, 1, 0, 0, 0, 0, time.UTC)), + field.SetTimestamp(10, 100, time.Date(2019, time.January, 1, 0, 0, 0, 0, time.UTC)), + )) + if err != nil { + t.Fatal(err) + } + start := time.Date(2017, time.January, 5, 0, 0, 0, 0, time.UTC) + end := time.Date(2018, time.January, 5, 0, 0, 0, 0, time.UTC) + resp, err := client.Query(field.RowRange(10, start, end)) + if err != nil { + t.Fatal(err) + } + target := []uint64{100} + if !reflect.DeepEqual(resp.Result().Row().Columns, target) { + t.Fatalf("%v != %v", target, resp.Result().Row().Columns) + } +} + +func TestRangeField(t *testing.T) { + client := getClient() + defer client.Close() + field := index.Field("rangefield", OptFieldTypeInt()) + field2 := index.Field("rangefield-set") + err := client.EnsureField(field) + if err != nil { + t.Fatal(err) + } + err = client.EnsureField(field2) + if err != nil { + t.Fatal(err) + } + _, err = client.Query(index.BatchQuery( + field2.Set(1, 10), + field2.Set(1, 100), + field.SetIntValue(10, 11), + field.SetIntValue(100, 15), + )) + if err != nil { + t.Fatal(err) + } + + resp, err := client.Query(field.Sum(field2.Row(1))) + if err != nil { + t.Fatal(err) + } + if resp.Result().Value() != 26 { + t.Fatalf("Sum 26 != %d", resp.Result().Value()) + } + if resp.Result().Count() != 2 { + t.Fatalf("Count 2 != %d", resp.Result().Count()) + } +} + +func TestRangeField2(t *testing.T) { + client := getClient() + defer client.Close() + field := index.Field("rangefield", OptFieldTypeInt(10, 20)) + field2 := index.Field("rangefield-set") + err := client.EnsureField(field) + if err != nil { + t.Fatal(err) + } + err = client.EnsureField(field2) + if err != nil { + t.Fatal(err) + } + + resp, err := client.Query(field.Min(field2.Row(1))) + if err != nil { + t.Fatal(err) + } + if resp.Result().Value() != 11 { + t.Fatalf("Min 11 != %d", resp.Result().Value()) + } + if resp.Result().Count() != 1 { + t.Fatalf("Count 1 != %d", resp.Result().Count()) + } + + resp, err = client.Query(field.Max(field2.Row(1))) + if err != nil { + t.Fatal(err) + } + if resp.Result().Value() != 15 { + t.Fatalf("Max 15 != %d", resp.Result().Value()) + } + if resp.Result().Count() != 1 { + t.Fatalf("Count 1 != %d", resp.Result().Count()) + } + + resp, err = client.Query(field.LT(15)) + if err != nil { + t.Fatal(err) + } + if len(resp.Result().Row().Columns) != 1 { + t.Fatalf("Count 1 != %d", len(resp.Result().Row().Columns)) + } + if resp.Result().Row().Columns[0] != 10 { + t.Fatalf("Column 10 != %d", resp.Result().Row().Columns[0]) + } +} + +func TestNotQuery(t *testing.T) { + client := getClient() + defer client.Close() + index := schema.Index("not-query-index", OptIndexTrackExistence(true)) + field := index.Field("not-field") + err := client.SyncSchema(schema) + if err != nil { + t.Fatal(err) + } + defer func() { + cerr := client.DeleteIndex(index) + if cerr != nil { + t.Errorf("failed to delete index: %v", cerr) + } + }() + + _, err = client.Query(index.BatchQuery( + field.Set(1, 10), + field.Set(1, 11), + field.Set(2, 11), + field.Set(2, 12), + field.Set(2, 13), + )) + if err != nil { + t.Fatal(err) + } + + resp, err := client.Query(index.Not(field.Row(1))) + if err != nil { + t.Fatal(err) + } + target := []uint64{12, 13} + if !reflect.DeepEqual(target, resp.Result().Row().Columns) { + t.Fatalf("%v != %v", target, resp.Result().Row().Columns) + } +} + +func TestStoreQuery(t *testing.T) { + client := getClient() + defer client.Close() + schema := NewSchema() + index := schema.Index("store-test") + fromField := index.Field("x-from-field") + toField := index.Field("x-to-field") + err := client.SyncSchema(schema) + if err != nil { + t.Fatal(err) + } + defer func() { + cerr := client.DeleteIndex(index) + if cerr != nil { + t.Errorf("failed to delete index: %v", cerr) + } + }() + + _, err = client.Query(index.BatchQuery( + fromField.Set(10, 100), + fromField.Set(10, 200), + toField.Store(fromField.Row(10), 1), + )) + if err != nil { + t.Fatal(err) + } + resp, err := client.Query(toField.Row(1)) + if err != nil { + t.Fatal(err) + } + target := []uint64{100, 200} + if !reflect.DeepEqual(target, resp.Result().Row().Columns) { + t.Fatalf("%v != %v", target, resp.Result().Row().Columns) + } +} + +func TestExcludeAttrsColumns(t *testing.T) { + client := getClient() + defer client.Close() + field := index.Field("excludecolumnsattrsfield") + err := client.EnsureField(field) + if err != nil { + t.Fatal(err) + } + attrs := map[string]interface{}{ + "foo": "bar", + } + _, err = client.Query(index.BatchQuery( + field.Set(1, 100), + field.SetRowAttrs(1, attrs), + )) + if err != nil { + t.Fatal(err) + } + + // test exclude columns. + resp, err := client.Query(field.Row(1), &QueryOptions{ExcludeColumns: true}) + if err != nil { + t.Fatal(err) + } + if len(resp.Result().Row().Columns) != 0 { + t.Fatalf("columns should be excluded") + } + if len(resp.Result().Row().Attributes) != 1 { + t.Fatalf("attributes should be included") + } + + // test exclude attributes. + resp, err = client.Query(field.Row(1), &QueryOptions{ExcludeRowAttrs: true}) + if err != nil { + t.Fatal(err) + } + if len(resp.Result().Row().Columns) != 1 { + t.Fatalf("columns should be included") + } + if len(resp.Result().Row().Attributes) != 0 { + t.Fatalf("attributes should be excluded") + } +} + +func TestMultipleClientKeyQuery(t *testing.T) { + client := getClient() + defer client.Close() + field := keysIndex.Field("multiple-client-field") + err := client.EnsureField(field) + if err != nil { + t.Fatal(err) + } + + const goroutineCount = 10 + wg := &sync.WaitGroup{} + wg.Add(goroutineCount) + for i := 0; i < goroutineCount; i++ { + go func(rowID uint64) { + if _, e := client.Query(field.Set(rowID, "col")); e != nil { + err = e + } + wg.Done() + }(uint64(i)) + } + wg.Wait() + + if err != nil { + t.Fatal(err) + } +} + +func TestDecodingFragmentNodesFails(t *testing.T) { + server := getMockServer(200, []byte("notjson"), 7) + defer server.Close() + client, _ := NewClient(server.URL, OptClientRetries(0)) + _, err := client.fetchFragmentNodes("foo", 0) + if err == nil { + t.Fatalf("fetchFragmentNodes should fail when response from /fragment/nodes cannot be decoded") + } +} + +func TestImportNodeFails(t *testing.T) { + server := getMockServer(500, []byte{}, 0) + defer server.Close() + uri, _ := pnet.NewURIFromAddress(server.URL) + client, _ := NewClient(uri, OptClientRetries(0)) + importRequest := &internal.ImportRequest{ + ColumnIDs: []uint64{}, + RowIDs: []uint64{}, + Timestamps: []int64{}, + Index: "foo", + Field: "bar", + Shard: 0, + } + data, err := proto.Marshal(importRequest) + if err != nil { + t.Fatalf("marshaling importRequest: %v", err) + } + err = client.importData(uri, "/index/foo/field/bar/import?clear=false", data) + if err == nil { + t.Fatalf("importNode should fail when posting to /import fails") + } +} + +func TestQueryUnmarshalFails(t *testing.T) { + server := getMockServer(200, []byte(`{}`), -1) + defer server.Close() + client, _ := NewClient(server.URL, OptClientRetries(0)) + field := NewSchema().Index("foo").Field("bar") + _, err := client.Query(field.Row(1)) + if err == nil { + t.Fatalf("should have failed") + } +} + +func TestResponseWithInvalidType(t *testing.T) { + qr := &internal.QueryResponse{ + Err: "", + ColumnAttrSets: []*internal.ColumnAttrSet{ + { + ID: 0, + Attrs: []*internal.Attr{ + { + Type: 9999, + StringValue: "NOVAL", + }, + }, + }, + }, + Results: []*internal.QueryResult{}, + } + data, err := proto.Marshal(qr) + if err != nil { + t.Fatal(err) + } + server := getMockServer(200, data, -1) + defer server.Close() + client, _ := NewClient(server.URL, OptClientRetries(0)) + _, err = client.Query(testField.Row(1)) + if err == nil { + t.Fatalf("Should have failed") + } +} + +func TestStatusFails(t *testing.T) { + server := getMockServer(404, nil, 0) + defer server.Close() + client, _ := NewClient(server.URL, OptClientRetries(0)) + _, err := client.Status() + if err == nil { + t.Fatalf("Should have failed") + } +} + +func TestStatusUnmarshalFails(t *testing.T) { + server := getMockServer(200, []byte("foo"), 3) + defer server.Close() + client, _ := NewClient(server.URL, OptClientRetries(0)) + _, err := client.Status() + if err == nil { + t.Fatalf("Should have failed") + } +} + +func TestStatusToNodeShardsForIndex(t *testing.T) { + client := getClient() + defer client.Close() + status := Status{ + Nodes: []StatusNode{ + { + URI: StatusURI{ + Scheme: "https", + Host: "localhost", + Port: 10101, + }, + }, + }, + indexMaxShard: map[string]uint64{ + index.Name(): 0, + }, + } + shardMap, err := client.statusToNodeShardsForIndex(status, index.Name()) + if err != nil { + t.Fatal(err) + } + if len(shardMap) != 1 { + t.Fatalf("len(shardMap) %d != %d", 1, len(shardMap)) + } + if _, ok := shardMap[0]; !ok { + t.Fatalf("shard map should have the correct shard") + } +} + +func TestHttpRequest(t *testing.T) { + client := getClient() + defer client.Close() + _, _, err := client.HTTPRequest("GET", "/status", nil, nil) + if err != nil { + t.Fatal(err) + } +} + +func TestSyncSchemaCantCreateIndex(t *testing.T) { + server := getMockServer(404, nil, 0) + defer server.Close() + client, _ := NewClient(server.URL, OptClientRetries(0)) + schema = NewSchema() + schema.Index("foo") + err := client.syncSchema(schema, NewSchema()) + if err == nil { + t.Fatalf("Should have failed") + } +} + +func TestSyncSchemaCantCreateField(t *testing.T) { + server := getMockServer(404, nil, 0) + defer server.Close() + client, _ := NewClient(server.URL, OptClientRetries(0)) + schema = NewSchema() + index := schema.Index("foo") + index.Field("foofield") + serverSchema := NewSchema() + serverSchema.Index("foo") + err := client.syncSchema(schema, serverSchema) + if err == nil { + t.Fatalf("Should have failed") + } +} + +func TestExportFieldFailure(t *testing.T) { + paths := map[string]mockResponseItem{ + "/status": { + content: []byte(`{"state":"NORMAL","nodes":[{"scheme":"http","host":"localhost","port":10101}]}`), + statusCode: 404, + contentLength: -1, + }, + "/internal/shards/max": { + content: []byte(`{"standard":{"go-testindex": 0}}`), + statusCode: 404, + contentLength: -1, + }, + } + server := getMockPathServer(paths) + defer server.Close() + client, _ := NewClient(server.URL, OptClientRetries(0)) + _, err := client.ExportField(testField) + if err == nil { + t.Fatal("should have failed") + } + statusItem := paths["/status"] + statusItem.statusCode = 200 + paths["/status"] = statusItem + _, err = client.ExportField(testField) + if err == nil { + t.Fatal("should have failed") + } + statusItem = paths["/internal/shards/max"] + statusItem.statusCode = 200 + paths["/internal/shards/max"] = statusItem + _, err = client.ExportField(testField) + if err == nil { + t.Fatal("should have failed") + } +} + +func TestShardsMaxDecodeFailure(t *testing.T) { + server := getMockServer(200, []byte(`{`), 0) + defer server.Close() + client, _ := NewClient(server.URL, OptClientRetries(0)) + _, err := client.shardsMax() + if err == nil { + t.Fatal("should have failed") + } +} + +func TestReadSchemaDecodeFailure(t *testing.T) { + server := getMockServer(200, []byte(`{`), 0) + defer server.Close() + client, _ := NewClient(server.URL, OptClientRetries(0)) + _, err := client.readSchema() + if err == nil { + t.Fatal("should have failed") + } +} + +func TestStatusToNodeShardsForIndexFailure(t *testing.T) { + server := getMockServer(200, []byte(`[]`), -1) + defer server.Close() + client, _ := NewClient(server.URL, OptClientRetries(0)) + // no shard + status := Status{ + indexMaxShard: map[string]uint64{}, + } + _, err := client.statusToNodeShardsForIndex(status, "foo") + if err == nil { + t.Fatal("should have failed") + } + + // no fragment nodes + status = Status{ + indexMaxShard: map[string]uint64{ + "foo": 0, + }, + } + _, err = client.statusToNodeShardsForIndex(status, "foo") + if err == nil { + t.Fatal("should have failed") + } +} + +func TestUserAgent(t *testing.T) { + handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + version := strings.TrimPrefix(Version, "v") + targetUserAgent := fmt.Sprintf("pilosa/client/%s", version) + if targetUserAgent != r.UserAgent() { + t.Fatalf("UserAgent %s != %s", targetUserAgent, r.UserAgent()) + } + }) + server := httptest.NewServer(handler) + defer server.Close() + client, _ := NewClient(server.URL, OptClientRetries(0)) + _, _, err := client.HTTPRequest("GET", "/version", nil, nil) + if err != nil { + t.Fatal(err) + } +} + +func TestClientRace(t *testing.T) { + uri, err := pnet.NewURIFromAddress(getPilosaBindAddress()) + if err != nil { + panic(err) + } + client, err := NewClient(uri, + OptClientTLSConfig(&tls.Config{InsecureSkipVerify: true}), + OptClientRetries(0)) + if err != nil { + panic(err) + } + f := func() { + if _, e := client.Query(testField.Row(1)); e != nil { + err = e + } + } + for i := 0; i < 10; i++ { + go f() + } + if err != nil { + panic(err) + } +} + +func TestFetchPrimaryFails(t *testing.T) { + server := getMockServer(404, []byte(`[]`), -1) + defer server.Close() + client, _ := NewClient(server.URL, OptClientRetries(0)) + _, err := client.fetchPrimaryNode() + if err == nil { + t.Fatal("should have failed") + } +} + +func TestFetchPrimaryPrimaryNotFound(t *testing.T) { + server := getMockServer(200, []byte(`{"state":"NORMAL","nodes":[{"id":"0f5c2ffc-1244-47d0-a83d-f5a25abba9bc","uri":{"scheme":"http","host":"localhost","port":10101}}],"localID":"0f5c2ffc-1244-47d0-a83d-f5a25abba9bc"}`), -1) + defer server.Close() + client, _ := NewClient(server.URL, OptClientRetries(0)) + _, err := client.fetchPrimaryNode() + if err == nil { + t.Fatal("should have failed") + } +} + +func TestServerWarning(t *testing.T) { + var herr error + defer func() { + if herr != nil { + t.Errorf("error in HTTP handler: %v", herr) + } + }() + handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + content, err := proto.Marshal(&internal.QueryResponse{}) + if err != nil { + // Cannot directly interact with t from another goroutine. + herr = err + w.WriteHeader(500) + return + } + w.Header().Set("warning", `299 pilosa/2.0 "FAKE WARNING: Deprecated PQL version: PQL v2 will remove support for SetBit() in Pilosa 2.1. Please update your client to support Set() (See https://docs.pilosa.com/pql#versioning)." "Sat, 25 Aug 2019 23:34:45 GMT"`) + w.WriteHeader(200) + _, err = io.Copy(w, bytes.NewReader(content)) + if err != nil { + // Cannot directly interact with t from another goroutine. + herr = err + return + } + }) + server := httptest.NewServer(handler) + defer server.Close() + client, _ := NewClient(server.URL, OptClientRetries(0)) + _, err := client.Query(testField.Row(1)) + if err != nil { + t.Fatal(err) + } +} + +func TestExportRowIDColumnID(t *testing.T) { + client := getClient() + defer client.Close() + field := index.Field("exportfield-rowid-colid") + err := client.EnsureField(field) + if err != nil { + t.Fatal(err) + } + _, err = client.Query(index.BatchQuery( + field.Set(1, 1), + field.Set(1, 10), + field.Set(2, 1048577), + ), nil) + if err != nil { + t.Fatal(err) + } + r, err := client.ExportField(field) + if err != nil { + t.Fatal(err) + } + s := consumeReader(t, r) + target := "1,1\n1,10\n2,1048577\n" + if target != s { + t.Fatalf("%s != %s", target, s) + } +} + +func TestExportRowIDColumnKey(t *testing.T) { + client := getClient() + defer client.Close() + field := keysIndex.Field("exportfield-rowid-colkey") + err := client.EnsureField(field) + if err != nil { + t.Fatal(err) + } + _, err = client.Query(keysIndex.BatchQuery( + field.Set(1, "one"), + field.Set(1, "ten"), + field.Set(2, "big-number"), + ), nil) + if err != nil { + t.Fatal(err) + } + r, err := client.ExportField(field) + if err != nil { + t.Fatal(err) + } + s := consumeReader(t, r) + target := "1,one\n1,ten\n2,big-number\n" + if target != s { + //t.Fatalf("%s != %s", target, s) + t.Log("TODO: these results do not necessarily come back ordered anymore!") + } +} + +func TestExportRowKeyColumnID(t *testing.T) { + client := getClient() + defer client.Close() + field := index.Field("exportfield-rowkey-colid", OptFieldKeys(true)) + err := client.EnsureField(field) + if err != nil { + t.Fatal(err) + } + _, err = client.Query(index.BatchQuery( + field.Set("one", 1), + field.Set("one", 10), + field.Set("two", 1048577), + ), nil) + if err != nil { + t.Fatal(err) + } + r, err := client.ExportField(field) + if err != nil { + t.Fatal(err) + } + s := consumeReader(t, r) + target := "one,1\none,10\ntwo,1048577\n" + if target != s { + t.Fatalf("%s != %s", target, s) + } +} + +func TestExportRowKeyColumnKey(t *testing.T) { + client := getClient() + defer client.Close() + field := keysIndex.Field("exportfield-rowkey-colkey", OptFieldKeys(true)) + err := client.EnsureField(field) + if err != nil { + t.Fatal(err) + } + _, err = client.Query(keysIndex.BatchQuery( + field.Set("one", "one"), + field.Set("one", "ten"), + field.Set("two", "big-number"), + ), nil) + if err != nil { + t.Fatal(err) + } + r, err := client.ExportField(field) + if err != nil { + t.Fatal(err) + } + s := consumeReader(t, r) + target := "one,one\none,ten\ntwo,big-number\n" + if target != s { + //t.Fatalf("%s != %s", target, s) + t.Log("TODO: these results do not necessarily come back ordered anymore!") + } +} + +func TestTranslateRowKeys(t *testing.T) { + client := getClient() + defer client.Close() + field := index.Field("translate-rowkeys", OptFieldKeys(true)) + err := client.EnsureField(field) + if err != nil { + t.Fatal(err) + } + _, err = client.Query(index.BatchQuery( + field.Set("key1", 10), + field.Set("key2", 1000), + )) + if err != nil { + t.Fatal(err) + } + rowIDs, err := client.TranslateRowKeys(field, []string{"key1", "key2"}) + if err != nil { + t.Fatal(err) + } + target := []uint64{1, 2} + if !reflect.DeepEqual(target, rowIDs) { + t.Fatalf("%v != %v", target, rowIDs) + } +} + +func TestTranslateColKeys(t *testing.T) { + client := getClient() + defer client.Close() + field := keysIndex.Field("translate-colkeys") + err := client.EnsureField(field) + if err != nil { + t.Fatal(err) + } + _, err = client.Query(keysIndex.BatchQuery( + field.Set(10, "ten"), + field.Set(1000, "one-thousand"), + )) + if err != nil { + t.Fatal(err) + } + colIDs, err := client.TranslateColumnKeys(keysIndex, []string{"ten", "one-thousand"}) + if err != nil { + t.Fatal(err) + } + target := []uint64{5242881, 82837505} + if !reflect.DeepEqual(target, colIDs) { + t.Fatalf("%v != %v", target, colIDs) + } +} + +func TestCSVExportFailure(t *testing.T) { + server := getMockServer(404, []byte("sorry, not found"), -1) + defer server.Close() + client, _ := NewClient(server.URL, OptClientRetries(0)) + field := index.Field("exportfield") + _, err := client.ExportField(field) + if err == nil { + t.Fatal("should have failed") + } +} + +func TestTransactions(t *testing.T) { + client := getClient() + defer client.Close() + + if trns, err := client.StartTransaction("blah", time.Minute, false, time.Minute); err != nil { + t.Errorf("%v", err) + } else if trns.ID != "blah" || trns.Timeout != time.Minute || !trns.Active { + t.Errorf("unexpected returned transaction: %+v", trns) + } + + if trnsMap, err := client.Transactions(); err != nil { + t.Errorf("listing transactions: %v", err) + } else if len(trnsMap) != 1 || !trnsMap["blah"].Active { + t.Errorf("unexpected trnsMap: %+v", trnsMap) + } + + if trns, err := client.GetTransaction("blah"); err != nil { + t.Errorf("%v", err) + } else if trns.ID != "blah" || trns.Timeout != time.Minute || !trns.Active { + t.Errorf("unexpected returned transaction: %+v", trns) + } + + if trns, err := client.FinishTransaction("blah"); err != nil { + t.Errorf("%v", err) + } else if trns.ID != "blah" || trns.Timeout != time.Minute || !trns.Active { + t.Errorf("unexpected returned transaction: %+v", trns) + } + +} + +func getMockServer(statusCode int, response []byte, contentLength int) *httptest.Server { + handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/x-protobuf") + if contentLength >= 0 { + w.Header().Set("Content-Length", strconv.Itoa(contentLength)) + } + w.WriteHeader(statusCode) + if response != nil { + _, _ = io.Copy(w, bytes.NewReader(response)) + } + }) + return httptest.NewServer(handler) +} + +type mockResponseItem struct { + content []byte + contentLength int + statusCode int +} + +func getMockPathServer(responses map[string]mockResponseItem) *httptest.Server { + handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/x-protobuf") + if item, ok := responses[r.RequestURI]; ok { + if item.contentLength >= 0 { + w.Header().Set("Content-Length", strconv.Itoa(item.contentLength)) + } else { + w.Header().Set("Content-Length", strconv.Itoa(len(item.content))) + } + statusCode := item.statusCode + if statusCode == 0 { + statusCode = 200 + } + w.WriteHeader(statusCode) + if item.content != nil { + _, _ = io.Copy(w, bytes.NewReader(item.content)) + } + return + } + w.WriteHeader(http.StatusNotFound) + _, _ = io.Copy(w, bytes.NewReader([]byte("not found"))) + }) + return httptest.NewServer(handler) +} + +func getClient(options ...ClientOption) *Client { + var client *Client + var err error + uri, err := pnet.NewURIFromAddress(getPilosaBindAddress()) + if err != nil { + panic(err) + } + options = append([]ClientOption{ + OptClientTLSConfig(&tls.Config{InsecureSkipVerify: true}), + OptClientRetries(0), + }, options...) + client, err = NewClient(uri, options...) + if err != nil { + panic(err) + } + return client +} + +func getPilosaBindAddress() string { + for _, kvStr := range os.Environ() { + kv := strings.SplitN(kvStr, "=", 2) + if kv[0] == "PILOSA_BIND" { + return kv[1] + } + } + return "http://:10101" +} + +func consumeReader(t *testing.T, r io.Reader) string { + b, err := ioutil.ReadAll(r) + if err != nil { + t.Fatal(err) + } + return string(b) +} + +//////////// new stuff + +func queryBalances(client *Client, acctOwnerID uint64, fieldAcct0, fieldAcct1 string) (acct0bal, acct1bal int64) { + + q := fmt.Sprintf("FieldValue(field=%v, column=%v)", fieldAcct0, acctOwnerID) + pql := NewPQLBaseQuery(q, indexAR, nil) + r, err := client.Query(pql) + panicOn(err) + acct0bal = r.ResultList[0].(*ValCountResult).Val + + q = fmt.Sprintf("FieldValue(field=%v, column=%v)", fieldAcct1, acctOwnerID) + pql = NewPQLBaseQuery(q, indexAR, nil) + r, err = client.Query(pql) + panicOn(err) + acct1bal = r.ResultList[0].(*ValCountResult).Val + + return +} + +func skipForRoaring(t *testing.T) { + src := os.Getenv("PILOSA_TXSRC") + if src == "" || strings.Contains(src, "roaring") { + t.Skip("skip if roaring pseudo-txn involved -- won't show transactional rollback/atomic commit") + } +} + +// Check the classic "bank balance transfer between accounts" +// to prevent read-anomalies when two writes are split by a read. +// +func TestImportAtomicRecord(t *testing.T) { + skipForRoaring(t) + Reset() + client := getClient() + defer client.Close() + uri := &pnet.URI{Scheme: "http", Port: 10101, Host: "localhost"} + acctOwnerID := uint64(78) // ColumnID + shard := uint64(0) + transferUSD := int64(100) + + setBal := func(bal0, bal1 int64) (data []byte, err error) { + ivr0 := &internal.ImportValueRequest{ + Index: indexARname, + Field: fieldAcct0, + Shard: shard, + ColumnIDs: []uint64{acctOwnerID}, + Values: []int64{bal0}, + } + ivr1 := &internal.ImportValueRequest{ + Index: indexARname, + Field: fieldAcct1, + Shard: shard, + ColumnIDs: []uint64{acctOwnerID}, + Values: []int64{bal1}, + } + + ar := &internal.AtomicRecord{ + Index: indexARname, + Shard: shard, + Ivr: []*internal.ImportValueRequest{ + ivr0, ivr1, + }, + } + + data, err = proto.Marshal(ar) + panicOn(err) + + return + } + + // setup 500 USD in acct1 and 700 USD in acct2. + // transfer 100 USD. + // should see 400 USD in acct, and 800 USD in acct2. + // + expectedBalStartingAcct0 := int64(500) + expectedBalStartingAcct1 := int64(700) + + data, err := setBal(expectedBalStartingAcct0, expectedBalStartingAcct1) + panicOn(err) + err = client.importData(uri, "/import-atomic-record", data) + panicOn(err) + + // start the main test, reading two balances and writing two updates. + + startingBalanceAcct0, startingBalanceAcct1 := queryBalances(client, acctOwnerID, fieldAcct0, fieldAcct1) + + //vv("starting balance: acct0=%v, acct1=%v", startingBalanceAcct0, startingBalanceAcct1) + + if startingBalanceAcct0 != expectedBalStartingAcct0 { + panic(fmt.Sprintf("expected %v, observed %v starting acct0 balance", expectedBalStartingAcct0, startingBalanceAcct0)) + } + if startingBalanceAcct1 != expectedBalStartingAcct1 { + panic(fmt.Sprintf("expected %v, observed %v starting acct1 balance", expectedBalStartingAcct1, startingBalanceAcct1)) + } + + data, err = setBal(expectedBalStartingAcct0-transferUSD, expectedBalStartingAcct1+transferUSD) + panicOn(err) + + //vv("sad path: transferUSD %v from %v -> %v, with power loss half-way through", transferUSD, fieldAcct0, fieldAcct1) + err = client.importData(uri, "/import-atomic-record?simPowerLossAfter=1", data) + if err == nil { + panic("expected to get 'update was aborted'") + } else { + if !strings.Contains(err.Error(), "update was aborted") { + panic(err) + } + } + + endingBalanceAcct0, endingBalanceAcct1 := queryBalances(client, acctOwnerID, fieldAcct0, fieldAcct1) + + // should not have been applied + if endingBalanceAcct0 != startingBalanceAcct0 || + endingBalanceAcct1 != startingBalanceAcct1 { + panic(fmt.Sprintf("problem: transaction did not abort atomically. Should have same start and end balances in both accounts, but we see: startingBalanceAcct0=%v -> endingBalanceAcct0=%v; startingBalanceAcct1=%v -> endingBalanceAcct1=%v", startingBalanceAcct0, endingBalanceAcct0, startingBalanceAcct1, endingBalanceAcct1)) + } + //vv("good: with power loss half-way, no change in account balances; acct0=%v; acct1=%v", endingBalanceAcct0, endingBalanceAcct1) + + // next part of the test, just make sure we do the update. + //vv("happy path: transferUSD %v from %v -> %v, with no interruption.", transferUSD, fieldAcct0, fieldAcct1) + + // happy path with no power failure half-way through. + + err = client.importData(uri, "/import-atomic-record?simPowerLossAfter=0", data) + panicOn(err) + endingBalanceAcct0, endingBalanceAcct1 = queryBalances(client, acctOwnerID, fieldAcct0, fieldAcct1) + + // should have been applied this time. + if endingBalanceAcct0 != startingBalanceAcct0-transferUSD || + endingBalanceAcct1 != startingBalanceAcct1+transferUSD { + panic(fmt.Sprintf("problem: transaction did not get committed/applied. transferUSD=%v, but we see: startingBalanceAcct0=%v -> endingBalanceAcct0=%v; startingBalanceAcct1=%v -> endingBalanceAcct1=%v", transferUSD, startingBalanceAcct0, endingBalanceAcct0, startingBalanceAcct1, endingBalanceAcct1)) + } + //vv("ending balance: acct0=%v, acct1=%v", endingBalanceAcct0, endingBalanceAcct1) + +} diff --git a/client/client_test.go b/client/client_test.go new file mode 100644 index 000000000..ab737c101 --- /dev/null +++ b/client/client_test.go @@ -0,0 +1,280 @@ +// Copyright 2017 Pilosa Corp. +// +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions +// are met: +// +// 1. Redistributions of source code must retain the above copyright +// notice, this list of conditions and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright +// notice, this list of conditions and the following disclaimer in the +// documentation and/or other materials provided with the distribution. +// +// 3. Neither the name of the copyright holder nor the names of its +// contributors may be used to endorse or promote products derived +// from this software without specific prior written permission. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND +// CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, +// INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF +// MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR +// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, +// BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH +// DAMAGE. + +package client + +import ( + "crypto/tls" + "errors" + "reflect" + "testing" + + "github.com/pilosa/pilosa/v2" + pnet "github.com/pilosa/pilosa/v2/net" +) + +func TestQueryWithError(t *testing.T) { + var err error + client := DefaultClient() + index := NewIndex("foo") + field := index.Field("foo") + invalid := field.FilterAttrTopN(12, field.Row(7), "$invalid$", 80, 81) + _, err = client.Query(invalid, nil) + if err == nil { + t.Fatalf("Should have failed") + } +} + +func TestClientOptions(t *testing.T) { + targets := []*ClientOptions{ + {SocketTimeout: 10}, + {ConnectTimeout: 5}, + {PoolSizePerRoute: 7}, + {TotalPoolSize: 17}, + {TLSConfig: &tls.Config{InsecureSkipVerify: true}}, + } + optionsList := [][]ClientOption{ + {OptClientSocketTimeout(10)}, + {OptClientConnectTimeout(5)}, + {OptClientPoolSizePerRoute(7)}, + {OptClientTotalPoolSize(17)}, + {OptClientTLSConfig(&tls.Config{InsecureSkipVerify: true})}, + } + + for i := 0; i < len(targets); i++ { + options := &ClientOptions{} + err := options.addOptions(optionsList[i]...) + if err != nil { + t.Fatal(err) + } + target := targets[i] + if !reflect.DeepEqual(target, options) { + t.Fatalf("%v != %v", target, options) + } + } +} + +func TestNewClientWithErrorredOption(t *testing.T) { + _, err := NewClient(":8888", ClientOptionErr(0)) + if err == nil { + t.Fatalf("Should have failed") + } +} + +func TestNewClient(t *testing.T) { + client, err := NewClient(":9999", OptClientManualServerAddress(true)) + if err != nil { + t.Fatal(err) + } + targetURI := pnet.URIFromAddress(":9999") + if !reflect.DeepEqual(targetURI, client.manualServerURI) { + t.Fatalf("%v != %v", targetURI, client.manualServerURI) + } + targetFragmentNode := &fragmentNode{ + Scheme: "http", + Host: "localhost", + Port: 9999, + } + if !reflect.DeepEqual(targetFragmentNode, client.manualFragmentNode) { + t.Fatalf("%v != %v", targetFragmentNode, client.manualFragmentNode) + } + client, err = NewClient(":9999") + if err != nil { + t.Fatal(err) + } + target := []*pnet.URI{pnet.URIFromAddress(":9999")} + if !reflect.DeepEqual(target, client.cluster.hosts) { + t.Fatalf("%v != %v", target, client.cluster.hosts) + } + client, err = NewClient([]string{":9999"}) + if err != nil { + t.Fatal(err) + } + if !reflect.DeepEqual(target, client.cluster.hosts) { + t.Fatalf("%v != %v", target, client.cluster.hosts) + } + + client, err = NewClient([]*pnet.URI{pnet.URIFromAddress(":9999"), pnet.URIFromAddress(":8888")}) + if err != nil { + t.Fatal(err) + } + target = []*pnet.URI{pnet.URIFromAddress(":9999"), pnet.URIFromAddress(":8888")} + if !reflect.DeepEqual(target, client.cluster.hosts) { + t.Fatalf("%v != %v", target, client.cluster.hosts) + } + + client, err = NewClient([]*pnet.URI{pnet.URIFromAddress(":9999")}) + if err != nil { + t.Fatal(err) + } + target = []*pnet.URI{pnet.URIFromAddress(":9999")} + if !reflect.DeepEqual(target, client.cluster.hosts) { + t.Fatalf("%v != %v", target, client.cluster.hosts) + } + + client, err = NewClient(DefaultCluster()) + if err != nil { + t.Fatal(err) + } + target = []*pnet.URI{} + if !reflect.DeepEqual(target, client.cluster.hosts) { + t.Fatalf("%v != %v", target, client.cluster.hosts) + } +} + +func TestNewClientWithInvalidAddr(t *testing.T) { + _, err := NewClient(10) + if err != ErrAddrURIClusterExpected { + t.Fatalf("%v != %v", ErrAddrURIClusterExpected, err) + } + _, err = NewClient(":invalid") + if err == nil { + t.Fatalf("should have failed: %+v", err) + } + _, err = NewClient([]string{"valid:8000", ":invalid"}) + if err != pilosa.ErrInvalidAddress { + t.Fatalf("Should have failed '%v', got '%v'", pilosa.ErrInvalidAddress, err) + } +} + +func TestNewClientManualAddressWithNoURIs(t *testing.T) { + _, err := NewClient([]string{}, OptClientManualServerAddress(true)) + if err != ErrSingleServerAddressRequired { + t.Fatalf("%v != %v", ErrSingleServerAddressRequired, err) + } + _, err = NewClient([]*pnet.URI{}, OptClientManualServerAddress(true)) + if err != ErrSingleServerAddressRequired { + t.Fatalf("%v != %v", ErrSingleServerAddressRequired, err) + } +} + +func TestNewClientManualAddressWithMultipleURIs(t *testing.T) { + _, err := NewClient([]string{":9000", ":5000"}, OptClientManualServerAddress(true)) + if err != ErrSingleServerAddressRequired { + t.Fatalf("%v != %v", ErrSingleServerAddressRequired, err) + } + _, err = NewClient([]*pnet.URI{pnet.URIFromAddress(":9000"), pnet.URIFromAddress(":5000")}, OptClientManualServerAddress(true)) + if err != ErrSingleServerAddressRequired { + t.Fatalf("%v != %v", ErrSingleServerAddressRequired, err) + } +} + +func ClientOptionErr(int) ClientOption { + return func(*ClientOptions) error { + return errors.New("Some error") + } +} + +func TestQueryOptions(t *testing.T) { + targets := []*QueryOptions{ + {ColumnAttrs: true}, + {ColumnAttrs: false}, + {ExcludeRowAttrs: true}, + {ExcludeRowAttrs: false}, + {ExcludeColumns: true}, + {ExcludeColumns: false}, + } + + optionsList := [][]interface{}{ + {OptQueryColumnAttrs(true)}, + {OptQueryColumnAttrs(false)}, + {OptQueryExcludeAttrs(true)}, + {OptQueryExcludeAttrs(false)}, + {OptQueryExcludeColumns(true)}, + {OptQueryExcludeColumns(false)}, + } + + for i := 0; i < len(targets); i++ { + options := &QueryOptions{} + err := options.addOptions(optionsList[i]...) + if err != nil { + t.Fatal(err) + } + target := targets[i] + if !reflect.DeepEqual(target, options) { + t.Fatalf("%v != %v", target, options) + } + } + + target := &QueryOptions{ + ColumnAttrs: true, + ExcludeRowAttrs: true, + ExcludeColumns: true, + } + options := &QueryOptions{} + err := options.addOptions(&QueryOptions{ + ColumnAttrs: true, + ExcludeRowAttrs: true, + ExcludeColumns: true, + }) + if err != nil { + t.Fatal(err) + } + if !reflect.DeepEqual(target, options) { + t.Fatalf("%v != %v", target, options) + } +} + +func TestQueryOptionsWithError(t *testing.T) { + options := &QueryOptions{} + err := options.addOptions(1) + if err == nil { + t.Fatalf("should have failed") + } + err = options.addOptions(OptQueryColumnAttrs(true), nil) + if err == nil { + t.Fatalf("should have failed") + } + err = options.addOptions(OptQueryColumnAttrs(true), &QueryOptions{}) + if err == nil { + t.Fatalf("should have failed") + } + err = options.addOptions(QueryOptionErr(0)) + if err == nil { + t.Fatalf("should have failed") + } +} + +func TestQueryOptionsError(t *testing.T) { + client := DefaultClient() + index := NewIndex("foo") + _, err := client.Query(index.RawQuery(""), QueryOptionErr(0)) + if err == nil { + t.Fatalf("should have failed") + } +} + +func QueryOptionErr(int) QueryOption { + return func(*QueryOptions) error { + return errors.New("Some error") + } +} diff --git a/client/cluster.go b/client/cluster.go new file mode 100644 index 000000000..9c66a23b5 --- /dev/null +++ b/client/cluster.go @@ -0,0 +1,127 @@ +// Copyright 2017 Pilosa Corp. +// +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions +// are met: +// +// 1. Redistributions of source code must retain the above copyright +// notice, this list of conditions and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright +// notice, this list of conditions and the following disclaimer in the +// documentation and/or other materials provided with the distribution. +// +// 3. Neither the name of the copyright holder nor the names of its +// contributors may be used to endorse or promote products derived +// from this software without specific prior written permission. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND +// CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, +// INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF +// MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR +// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, +// BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH +// DAMAGE. + +package client + +import ( + "sync" + + pnet "github.com/pilosa/pilosa/v2/net" +) + +// Cluster contains hosts in a Pilosa cluster. +type Cluster struct { + hosts []*pnet.URI + okList []bool + mutex *sync.RWMutex + lastHostIdx int +} + +// DefaultCluster returns the default Cluster. +func DefaultCluster() *Cluster { + return &Cluster{ + hosts: make([]*pnet.URI, 0), + okList: make([]bool, 0), + mutex: &sync.RWMutex{}, + } +} + +// NewClusterWithHost returns a cluster with the given URIs. +func NewClusterWithHost(hosts ...*pnet.URI) *Cluster { + cluster := DefaultCluster() + for _, host := range hosts { + cluster.AddHost(host) + } + return cluster +} + +// AddHost adds a host to the cluster. +func (c *Cluster) AddHost(address *pnet.URI) { + c.mutex.Lock() + defer c.mutex.Unlock() + c.hosts = append(c.hosts, address) + c.okList = append(c.okList, true) +} + +// Host returns a host in the cluster. +func (c *Cluster) Host() *pnet.URI { + c.mutex.Lock() + var host *pnet.URI + for i := range c.okList { + idx := (i + c.lastHostIdx) % len(c.okList) + ok := c.okList[idx] + if ok { + host = c.hosts[idx] + break + } + } + c.lastHostIdx++ + c.mutex.Unlock() + if host != nil { + return host + } + c.reset() + return host +} + +// RemoveHost black lists the host with the given pnet.URI from the cluster. +func (c *Cluster) RemoveHost(address *pnet.URI) { + c.mutex.Lock() + defer c.mutex.Unlock() + for i, uri := range c.hosts { + if uri.Equals(address) { + c.okList[i] = false + break + } + } +} + +// Hosts returns all available hosts in the cluster. +func (c *Cluster) Hosts() []pnet.URI { + c.mutex.RLock() + defer c.mutex.RUnlock() + hosts := make([]pnet.URI, 0, len(c.hosts)) + for i, host := range c.hosts { + if c.okList[i] { + hosts = append(hosts, *host) + } + } + return hosts +} + +func (c *Cluster) reset() { + c.mutex.Lock() + defer c.mutex.Unlock() + for i := range c.okList { + c.okList[i] = true + } +} diff --git a/client/cluster_test.go b/client/cluster_test.go new file mode 100644 index 000000000..acc8eed05 --- /dev/null +++ b/client/cluster_test.go @@ -0,0 +1,98 @@ +// Copyright 2017 Pilosa Corp. +// +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions +// are met: +// +// 1. Redistributions of source code must retain the above copyright +// notice, this list of conditions and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright +// notice, this list of conditions and the following disclaimer in the +// documentation and/or other materials provided with the distribution. +// +// 3. Neither the name of the copyright holder nor the names of its +// contributors may be used to endorse or promote products derived +// from this software without specific prior written permission. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND +// CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, +// INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF +// MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR +// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, +// BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH +// DAMAGE. + +package client + +import ( + "testing" + + pnet "github.com/pilosa/pilosa/v2/net" +) + +func TestNewClusterWithHost(t *testing.T) { + c := NewClusterWithHost(pnet.DefaultURI()) + hosts := c.Hosts() + if len(hosts) != 1 || !hosts[0].Equals(pnet.DefaultURI()) { + t.Fail() + } +} + +func TestAddHost(t *testing.T) { + const addr = "http://localhost:3000" + c := DefaultCluster() + if c.Hosts() == nil { + t.Fatalf("Hosts should not be nil") + } + uri, err := pnet.NewURIFromAddress(addr) + if err != nil { + t.Fatalf("Cannot parse address") + } + target, err := pnet.NewURIFromAddress(addr) + if err != nil { + t.Fatalf("Cannot parse address") + } + c.AddHost(uri) + hosts := c.Hosts() + if len(hosts) != 1 || !hosts[0].Equals(target) { + t.Fail() + } +} + +func TestHosts(t *testing.T) { + c := DefaultCluster() + if c.Host() != nil { + t.Fatalf("Hosts with empty cluster should return nil") + } + c = NewClusterWithHost(pnet.DefaultURI()) + if !c.Host().Equals(pnet.DefaultURI()) { + t.Fatalf("Host should return a value if there are hosts in the cluster") + } +} + +func TestRemoveHost(t *testing.T) { + uri, err := pnet.NewURIFromAddress("index1.pilosa.com:9999") + if err != nil { + t.Fatal(err) + } + c := NewClusterWithHost(uri) + if len(c.hosts) != 1 { + t.Fatalf("The cluster should contain the host") + } + uri, err = pnet.NewURIFromAddress("index1.pilosa.com:9999") + if err != nil { + t.Fatal(err) + } + c.RemoveHost(uri) + if len(c.Hosts()) != 0 { + t.Fatalf("The cluster should not contain the host") + } +} diff --git a/client/csv/csv.go b/client/csv/csv.go new file mode 100644 index 000000000..9c30a622e --- /dev/null +++ b/client/csv/csv.go @@ -0,0 +1,194 @@ +// Copyright 2017 Pilosa Corp. +// +// 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 csv + +import ( + "bufio" + "errors" + "fmt" + "io" + "strconv" + "strings" + "time" + + "github.com/pilosa/pilosa/v2/client" +) + +// Format is the format of the data in the CSV file. +type Format uint + +const ( + // RowIDColumnID formatted data is ROW_ID,COLUMN_ID. + RowIDColumnID Format = iota + // RowIDColumnKey formatted data is ROW_ID,COLUMN_KEY. + RowIDColumnKey + // RowKeyColumnID formatted data is ROW_KEY,COLUMN_ID. + RowKeyColumnID + // RowKeyColumnKey formatted data is ROW_KEY,COLUMN_ID. + RowKeyColumnKey + // ColumnID formatted data is COLUMN_ID. Valid only for value import. + ColumnID + // ColumnKey formatted data is COLUMN_KEY. Valud only for value import. + ColumnKey +) + +// ColumnUnmarshaller creates a RecordUnmarshaller for importing columns with the given format. +func ColumnUnmarshaller(format Format) RecordUnmarshaller { + return ColumnUnmarshallerWithTimestamp(format, "") +} + +// ColumnUnmarshallerWithTimestamp creates a RecordUnmarshaller for importing columns with the given format and timestamp format. +func ColumnUnmarshallerWithTimestamp(format Format, timestampFormat string) RecordUnmarshaller { + return func(text string) (client.Record, error) { + var err error + column := client.Column{} + parts := strings.Split(text, ",") + if len(parts) < 2 { + return nil, errors.New("Invalid CSV line") + } + + hasRowKey := format == RowKeyColumnID || format == RowKeyColumnKey + hasColumnKey := format == RowIDColumnKey || format == RowKeyColumnKey + + if hasRowKey { + column.RowKey = parts[0] + } else { + column.RowID, err = strconv.ParseUint(parts[0], 10, 64) + if err != nil { + return nil, errors.New("Invalid row ID") + } + } + + if hasColumnKey { + column.ColumnKey = parts[1] + } else { + column.ColumnID, err = strconv.ParseUint(parts[1], 10, 64) + if err != nil { + return nil, errors.New("Invalid column ID") + } + } + + timestamp := int64(0) + if len(parts) == 3 { + if timestampFormat == "" { + if tsInt, err := strconv.Atoi(parts[2]); err != nil { + return nil, err + } else { + timestamp = int64(tsInt) + } + } else { + t, err := time.Parse(timestampFormat, parts[2]) + if err != nil { + return nil, err + } + timestamp = t.Unix() * int64(time.Second) // Casting a duration to int64 gives the number of nanoseconds in that duration. + } + } + column.Timestamp = timestamp + + return column, nil + } +} + +// RecordUnmarshaller is a function which creates a Record from a CSV file line with column data. +type RecordUnmarshaller func(text string) (client.Record, error) + +// Iterator reads records from a Reader. +// Each line should contain a single record in the following form: +// field1,field2,... +type Iterator struct { + reader io.Reader + line int + scanner *bufio.Scanner + unmarshaller RecordUnmarshaller +} + +// NewIterator creates a CSVIterator from a Reader. +func NewIterator(reader io.Reader, unmarshaller RecordUnmarshaller) *Iterator { + return &Iterator{ + reader: reader, + line: 0, + scanner: bufio.NewScanner(reader), + unmarshaller: unmarshaller, + } +} + +// NewColumnIterator creates a new iterator for column data. +func NewColumnIterator(format Format, reader io.Reader) *Iterator { + return NewIterator(reader, ColumnUnmarshaller(format)) +} + +// NewColumnIteratorWithTimestampFormat creates a new iterator for column data with timestamp. +func NewColumnIteratorWithTimestampFormat(format Format, reader io.Reader, timestampFormat string) *Iterator { + return NewIterator(reader, ColumnUnmarshallerWithTimestamp(format, timestampFormat)) +} + +// NewValueIterator creates a new iterator for value data. +func NewValueIterator(format Format, reader io.Reader) *Iterator { + return NewIterator(reader, FieldValueUnmarshaller(format)) +} + +// NextRecord iterates on lines of a Reader. +// Returns io.EOF on end of iteration. +func (c *Iterator) NextRecord() (client.Record, error) { + if ok := c.scanner.Scan(); ok { + c.line++ + text := strings.TrimSpace(c.scanner.Text()) + if text != "" { + rc, err := c.unmarshaller(text) + if err != nil { + return nil, fmt.Errorf("%s at line: %d", err.Error(), c.line) + } + return rc, nil + } + } + err := c.scanner.Err() + if err != nil { + return nil, err + } + return nil, io.EOF +} + +// FieldValueUnmarshaller is a function which creates a Record from a CSV file line with value data. +func FieldValueUnmarshaller(format Format) RecordUnmarshaller { + return func(text string) (client.Record, error) { + parts := strings.Split(text, ",") + if len(parts) < 2 { + return nil, errors.New("Invalid CSV") + } + value, err := strconv.ParseInt(parts[1], 10, 64) + if err != nil { + return nil, errors.New("Invalid value") + } + switch format { + case ColumnID: + columnID, err := strconv.ParseUint(parts[0], 10, 64) + if err != nil { + return nil, errors.New("Invalid column ID at line: %d") + } + return client.FieldValue{ + ColumnID: uint64(columnID), + Value: value, + }, nil + case ColumnKey: + return client.FieldValue{ + ColumnKey: parts[0], + Value: value, + }, nil + default: + return nil, fmt.Errorf("Invalid format: %d", format) + } + } +} diff --git a/client/csv/csv_it_test.go b/client/csv/csv_it_test.go new file mode 100644 index 000000000..3df6badc9 --- /dev/null +++ b/client/csv/csv_it_test.go @@ -0,0 +1,78 @@ +//+build integration + +// Copyright 2017 Pilosa Corp. +// +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions +// are met: +// +// 1. Redistributions of source code must retain the above copyright +// notice, this list of conditions and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright +// notice, this list of conditions and the following disclaimer in the +// documentation and/or other materials provided with the distribution. +// +// 3. Neither the name of the copyright holder nor the names of its +// contributors may be used to endorse or promote products derived +// from this software without specific prior written permission. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND +// CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, +// INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF +// MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR +// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, +// BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH +// DAMAGE. + +package csv_test + +import ( + "io" + "reflect" + "strings" + "testing" + + "github.com/pilosa/pilosa/v2/client" + "github.com/pilosa/pilosa/v2/client/csv" +) + +func TestCSVIterate(t *testing.T) { + text := `10,7 + 10,5 + 2,3 + 7,1` + iterator := csv.NewColumnIterator(csv.RowIDColumnID, strings.NewReader(text)) + recs := consumeIterator(t, iterator) + target := []client.Record{ + client.Column{RowID: 10, ColumnID: 7}, + client.Column{RowID: 10, ColumnID: 5}, + client.Column{RowID: 2, ColumnID: 3}, + client.Column{RowID: 7, ColumnID: 1}, + } + if !reflect.DeepEqual(target, recs) { + t.Fatalf("%v != %v", target, recs) + } +} + +func consumeIterator(t *testing.T, it *csv.Iterator) []client.Record { + recs := []client.Record{} + for { + r, err := it.NextRecord() + if err == io.EOF { + break + } + if err != nil { + t.Fatal(err) + } + recs = append(recs, r) + } + return recs +} diff --git a/client/csv/csv_test.go b/client/csv/csv_test.go new file mode 100644 index 000000000..4ef162b8e --- /dev/null +++ b/client/csv/csv_test.go @@ -0,0 +1,267 @@ +// Copyright 2017 Pilosa Corp. +// +// 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 csv_test + +import ( + "errors" + "io" + "reflect" + "strings" + "testing" + + "github.com/pilosa/pilosa/v2" + "github.com/pilosa/pilosa/v2/client" + "github.com/pilosa/pilosa/v2/client/csv" +) + +func TestCSVColumnIterator(t *testing.T) { + reader := strings.NewReader(`1,10,683793200 + 5,20,683793300 + 3,41,683793385`) + iterator := csv.NewColumnIterator(csv.RowIDColumnID, reader) + columns := []client.Record{} + for { + column, err := iterator.NextRecord() + if err == io.EOF { + break + } + if err != nil { + t.Fatal(err) + } + columns = append(columns, column) + } + if len(columns) != 3 { + t.Fatalf("There should be 3 columns") + } + target := []client.Column{ + {RowID: 1, ColumnID: 10, Timestamp: 683793200}, + {RowID: 5, ColumnID: 20, Timestamp: 683793300}, + {RowID: 3, ColumnID: 41, Timestamp: 683793385}, + } + for i := range target { + if !reflect.DeepEqual(target[i], columns[i]) { + t.Fatalf("%v != %v", target[i], columns[i]) + } + } +} + +func TestCSVColumnIteratorWithTimestampFormatRowIDColumnID(t *testing.T) { + format := "2006-01-02T03:04" + reader := strings.NewReader(`1,10,1991-09-02T09:33 + 5,20,1991-09-02T09:35 + 3,41,1991-09-02T09:36`) + iterator := csv.NewColumnIteratorWithTimestampFormat(csv.RowIDColumnID, reader, format) + records := []client.Record{} + for { + record, err := iterator.NextRecord() + if err == io.EOF { + break + } + if err != nil { + t.Fatal(err) + } + records = append(records, record) + } + target := []client.Column{ + {RowID: 1, ColumnID: 10, Timestamp: 683803980000000000}, + {RowID: 5, ColumnID: 20, Timestamp: 683804100000000000}, + {RowID: 3, ColumnID: 41, Timestamp: 683804160000000000}, + } + if len(records) != len(target) { + t.Fatalf("There should be %d columns", len(target)) + } + for i := range target { + if !reflect.DeepEqual(target[i], records[i]) { + t.Fatalf("%v != %v", target[i], records[i]) + } + } +} + +func TestCSVColumnIteratorWithTimestampFormatRowKeyColumnKey(t *testing.T) { + format := "2006-01-02T03:04" + reader := strings.NewReader(`one,ten,1991-09-02T09:33 + five,twenty,1991-09-02T09:35 + three,forty-one,1991-09-02T09:36`) + iterator := csv.NewColumnIteratorWithTimestampFormat(csv.RowKeyColumnKey, reader, format) + records := []client.Record{} + for { + record, err := iterator.NextRecord() + if err == io.EOF { + break + } + if err != nil { + t.Fatal(err) + } + records = append(records, record) + } + target := []client.Column{ + {RowKey: "one", ColumnKey: "ten", Timestamp: 683803980000000000}, + {RowKey: "five", ColumnKey: "twenty", Timestamp: 683804100000000000}, + {RowKey: "three", ColumnKey: "forty-one", Timestamp: 683804160000000000}, + } + if len(records) != len(target) { + t.Fatalf("There should be %d columns", len(target)) + } + for i := range target { + if !reflect.DeepEqual(target[i], records[i]) { + t.Fatalf("%v != %v", target[i], records[i]) + } + } +} + +func TestCSVColumnIteratorWithTimestampFormatFail(t *testing.T) { + format := "2014-07-16" + reader := strings.NewReader(`1,10,X`) + iterator := csv.NewColumnIteratorWithTimestampFormat(csv.RowIDColumnID, reader, format) + _, err := iterator.NextRecord() + if err == nil { + t.Fatalf("Should have failed") + } +} + +func TestCSVValueIteratorWithColumnID(t *testing.T) { + reader := strings.NewReader(`1,10 + 5,-20 + 3,41 + `) + iterator := csv.NewValueIterator(csv.ColumnID, reader) + values := []client.Record{} + for { + value, err := iterator.NextRecord() + if err == io.EOF { + break + } + if err != nil { + t.Fatal(err) + } + values = append(values, value) + } + target := []pilosa.FieldValue{ + {ColumnID: 1, Value: 10}, + {ColumnID: 5, Value: -20}, + {ColumnID: 3, Value: 41}, + } + if len(values) != len(target) { + t.Fatalf("There should be %d values, got %d", len(target), len(values)) + } + for i := range target { + v := values[i].(client.FieldValue) + if !reflect.DeepEqual(pilosa.FieldValue(v), target[i]) { + t.Fatalf("'%+v' != '%+v'", target[i], values[i]) + } + } +} + +func TestCSVValueIteratorWithColumnKey(t *testing.T) { + reader := strings.NewReader(`one,10 + five,-20 + three,41 + `) + iterator := csv.NewValueIterator(csv.ColumnKey, reader) + values := []client.Record{} + for { + value, err := iterator.NextRecord() + if err == io.EOF { + break + } + if err != nil { + t.Fatal(err) + } + values = append(values, value) + } + target := []pilosa.FieldValue{ + {ColumnKey: "one", Value: 10}, + {ColumnKey: "five", Value: -20}, + {ColumnKey: "three", Value: 41}, + } + if len(values) != len(target) { + t.Fatalf("There should be %d values, got %d", len(target), len(values)) + } + for i := range target { + v := values[i].(client.FieldValue) + if !reflect.DeepEqual(pilosa.FieldValue(v), target[i]) { + t.Fatalf("%v != %v", target[i], values[i]) + } + } +} + +func TestCSValueIteratorWithInvalidFormat(t *testing.T) { + reader := strings.NewReader("1,2") + iterator := csv.NewValueIterator(csv.RowIDColumnID, reader) + _, err := iterator.NextRecord() + if err == nil { + t.Fatalf("should have failed") + } +} + +func TestCSVColumnIteratorInvalidInput(t *testing.T) { + invalidInputs := []string{ + // less than 2 columns + "155", + // invalid row ID + "a5,155", + // invalid column ID + "155,a5", + // invalid timestamp + "155,255,a5", + } + for _, text := range invalidInputs { + iterator := csv.NewColumnIterator(csv.RowIDColumnID, strings.NewReader(text)) + _, err := iterator.NextRecord() + if err == nil { + t.Fatalf("CSVColumnIterator input: %s should fail", text) + } + } +} + +func TestCSVValueIteratorInvalidInput(t *testing.T) { + invalidInputs := []string{ + // less than 2 columns + "155", + // invalid column ID + "a5,155", + // invalid value + "155,a5", + } + for _, text := range invalidInputs { + iterator := csv.NewValueIterator(csv.ColumnID, strings.NewReader(text)) + _, err := iterator.NextRecord() + if err == nil { + t.Fatalf("CSVValueIterator input: %s should fail", text) + } + } +} + +func TestCSVColumnIteratorError(t *testing.T) { + iterator := csv.NewColumnIterator(csv.RowIDColumnID, &BrokenReader{}) + _, err := iterator.NextRecord() + if err == nil { + t.Fatal("CSVColumnIterator should fail with error") + } +} + +func TestCSVValueIteratorError(t *testing.T) { + iterator := csv.NewValueIterator(csv.ColumnID, &BrokenReader{}) + _, err := iterator.NextRecord() + if err == nil { + t.Fatal("CSVValueIterator should fail with error") + } +} + +type BrokenReader struct{} + +func (r BrokenReader) Read(p []byte) (n int, err error) { + return 0, errors.New("broken reader") +} diff --git a/client/doc.go b/client/doc.go new file mode 100644 index 000000000..212642ff1 --- /dev/null +++ b/client/doc.go @@ -0,0 +1,83 @@ +// Copyright 2017 Pilosa Corp. +// +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions +// are met: +// +// 1. Redistributions of source code must retain the above copyright +// notice, this list of conditions and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright +// notice, this list of conditions and the following disclaimer in the +// documentation and/or other materials provided with the distribution. +// +// 3. Neither the name of the copyright holder nor the names of its +// contributors may be used to endorse or promote products derived +// from this software without specific prior written permission. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND +// CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, +// INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF +// MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR +// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, +// BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH +// DAMAGE. + +/* +Package pilosa enables querying a Pilosa server. + +This client uses Pilosa's http+protobuf API. + +Usage: + + import ( + "fmt" + "github.com/pilosa/pilosa/v2/client" + ) + + // Create a Client instance + client := client.DefaultClient() + + // Create a Schema instance + schema, err := client.Schema() + if err != nil { + panic(err) + } + + // Create an Index instance + index, err := schema.Index("repository") + if err != nil { + panic(err) + } + + // Create a Field instance + stargazer, err := index.Field("stargazer") + if err != nil { + panic(err) + } + + // Sync the schema with the server-side, so non-existing indexes/fields are created on the server-side. + err = client.SyncSchema(schema) + if err != nil { + panic(err) + } + + // Execute a query + response, err := client.Query(stargazer.Row(5)) + if err != nil { + panic(err) + } + + // Act on the result + fmt.Println(response.Result()) + +See also https://www.pilosa.com/docs/api-reference/ and https://www.pilosa.com/docs/query-language/. +*/ +package client diff --git a/client/egpool/egpool.go b/client/egpool/egpool.go new file mode 100644 index 000000000..92449cf7c --- /dev/null +++ b/client/egpool/egpool.go @@ -0,0 +1,123 @@ +// Copyright 2017 Pilosa Corp. +// +// 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 egpool + +import ( + "errors" + "fmt" + "sync" +) + +type Group struct { + PoolSize int + + jobs chan func() error + + sema chan struct{} + errMu sync.Mutex + firstErr error + errs []error +} + +func (eg *Group) Go(f func() error) { + if eg.PoolSize <= 0 { + eg.PoolSize = 1 + } + + if eg.jobs == nil { + eg.jobs = make(chan func() error) + eg.sema = make(chan struct{}, eg.PoolSize) + } + + // Start the job in an idle worker if possible. + select { + case eg.jobs <- f: + return + default: + } + + // Start a new worker if necessary. + select { + case eg.jobs <- f: + // A worker finished its previous job and took this one over. + return + case eg.sema <- struct{}{}: + // Start a new worker. + go eg.processJobs() + eg.jobs <- f + } +} + +func (eg *Group) err(err error) { + eg.errMu.Lock() + defer eg.errMu.Unlock() + + if eg.firstErr == nil { + eg.firstErr = err + } + eg.errs = append(eg.errs, err) +} + +type ErrPanic struct { + Value interface{} +} + +func (p ErrPanic) Error() string { + return fmt.Sprintf("panic: %v", p.Value) +} + +var ErrGoexit = errors.New("runtime.Goexit used in job function") + +func (eg *Group) processJobs() { + // Notify pool of shutdown. + defer func() { <-eg.sema }() + + // Handle panic and Goexit. + var finished bool + defer func() { + if !finished { + if p := recover(); p != nil { + eg.err(ErrPanic{p}) + } else { + eg.err(ErrGoexit) + } + } + }() + + // Run jobs from queue. + for jobFn := range eg.jobs { + err := jobFn() + if err != nil { + eg.err(err) + } + } + + finished = true +} + +func (eg *Group) Wait() error { + if eg.jobs == nil { + return nil + } + close(eg.jobs) + for i := 0; i < eg.PoolSize; i++ { + eg.sema <- struct{}{} + } + return eg.firstErr +} + +func (eg *Group) Errors() []error { + return eg.errs +} diff --git a/client/egpool/egpool_test.go b/client/egpool/egpool_test.go new file mode 100644 index 000000000..288c9fcb6 --- /dev/null +++ b/client/egpool/egpool_test.go @@ -0,0 +1,50 @@ +// Copyright 2017 Pilosa Corp. +// +// 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 egpool_test + +import ( + "errors" + "testing" + + "github.com/pilosa/pilosa/v2/client/egpool" +) + +func TestEGPool(t *testing.T) { + eg := egpool.Group{} + + a := make([]int, 10) + + for i := 0; i < 10; i++ { + i := i + eg.Go(func() error { + a[i] = i + if i == 7 { + return errors.New("blah") + } + return nil + }) + } + + err := eg.Wait() + if err == nil || err.Error() != "blah" { + t.Errorf("expected err blah, got: %v", err) + } + + for i := 0; i < 10; i++ { + if a[i] != i { + t.Errorf("expected a[%d] to be %d, but is %d", i, i, a[i]) + } + } +} diff --git a/client/error.go b/client/error.go new file mode 100644 index 000000000..ec9bc7912 --- /dev/null +++ b/client/error.go @@ -0,0 +1,56 @@ +// Copyright 2017 Pilosa Corp. +// +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions +// are met: +// +// 1. Redistributions of source code must retain the above copyright +// notice, this list of conditions and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright +// notice, this list of conditions and the following disclaimer in the +// documentation and/or other materials provided with the distribution. +// +// 3. Neither the name of the copyright holder nor the names of its +// contributors may be used to endorse or promote products derived +// from this software without specific prior written permission. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND +// CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, +// INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF +// MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR +// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, +// BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH +// DAMAGE. + +package client + +import "github.com/pkg/errors" + +// Predefined Pilosa errors. +var ( + ErrEmptyCluster = errors.New("No usable addresses in the cluster") + ErrIndexExists = errors.New("Index exists") + ErrFieldExists = errors.New("Field exists") + ErrInvalidIndexName = errors.New("Invalid index name") + ErrInvalidFieldName = errors.New("Invalid field name") + ErrInvalidLabel = errors.New("Invalid label") + ErrInvalidKey = errors.New("Invalid key") + ErrTriedMaxHosts = errors.New("Tried max hosts, still failing") + ErrAddrURIClusterExpected = errors.New("Addresses, URIs or a cluster is expected") + ErrInvalidQueryOption = errors.New("Invalid query option") + ErrInvalidIndexOption = errors.New("Invalid index option") + ErrInvalidFieldOption = errors.New("Invalid field option") + ErrNoFragmentNodes = errors.New("No fragment nodes") + ErrNoShard = errors.New("Index has no shards") + ErrUnknownType = errors.New("Unknown type") + ErrSingleServerAddressRequired = errors.New("OptClientManualServerAddress requires a single URI or address") + ErrPreconditionFailed = errors.New("Precondition failed") +) diff --git a/client/logimport.go b/client/logimport.go new file mode 100644 index 000000000..56202b3a0 --- /dev/null +++ b/client/logimport.go @@ -0,0 +1,45 @@ +// Copyright 2017 Pilosa Corp. +// +// 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 client + +import ( + "encoding/gob" + "io" +) + +type importLog struct { + Index string + Path string + Shard uint64 + IsRoaring bool + Timestamp int64 // Unix Nanoseconds + Data []byte +} + +type encoder interface { + Encode(thing interface{}) error +} + +func newImportLogEncoder(w io.Writer) encoder { + return gob.NewEncoder(w) +} + +type decoder interface { + Decode(thing interface{}) error +} + +func newImportLogDecoder(r io.Reader) decoder { + return gob.NewDecoder(r) +} diff --git a/client/logimport_test.go b/client/logimport_test.go new file mode 100644 index 000000000..85ceab535 --- /dev/null +++ b/client/logimport_test.go @@ -0,0 +1,155 @@ +// Copyright 2017 Pilosa Corp. +// +// 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 client + +import ( + "bytes" + "fmt" + "io/ioutil" + "os" + "reflect" + "testing" +) + +func TestEncodeDecode(t *testing.T) { + tests := []importLog{ + { + Index: "go-testindex", + Path: "/index/go-testindex/field/importfield-batchsize/import?clear=false", + Shard: 0, + Data: make([]byte, 3918), + }, + { + Index: "go-testindex", + Path: "/index/go-testindex/field/importfield-batchsize/import?clear=false", + Shard: 0, + Data: make([]byte, 3918), + }, + { + Index: "eheh", + Path: "blah", + Shard: 9, + Data: []byte("something"), + }, + { + Index: "", + Path: "", + Shard: 0, + Data: nil, + }, + { + Index: "eheh", + Path: "blah", + Shard: 10, + Data: []byte("blahaslkdjfeoiwujf"), + }, + { + Index: "eheh", + Path: "blah", + Shard: 10, + Data: make([]byte, 10000), + }, + { + Index: "zoop", + Path: "blah", + Shard: 8923734, + Data: []byte("blahaslkdjfeoiwujf"), + }, + } + + for i, test := range tests { + t.Run(fmt.Sprintf("%d", i), func(t *testing.T) { + nl := importLog{ + Index: test.Index, + Path: test.Path, + Shard: test.Shard, + Data: make([]byte, len(test.Data)), + } + copy(nl.Data, test.Data) + buf := &bytes.Buffer{} + enc := newImportLogEncoder(buf) + err := enc.Encode(nl) + if err != nil { + t.Fatalf("writing to buf: %v", err) + } + dec := newImportLogDecoder(buf) + l2 := &importLog{} + err = dec.Decode(l2) + if err != nil { + t.Fatalf("reading from buf: %v", err) + } + if l2.Index != test.Index { + t.Errorf("indexes not equal:\n%s\n%s", test.Index, l2.Index) + } + if l2.Path != test.Path { + t.Errorf("paths not equal:\n%s\n%s", test.Path, l2.Path) + } + if l2.Shard != test.Shard { + t.Errorf("shards not equal exp: %d got %d", test.Shard, l2.Shard) + } + if !reflect.DeepEqual(test.Data, l2.Data) { + t.Errorf("data not equal \n%v\n%v", test.Data, l2.Data) + } + + }) + } + buf, err := ioutil.TempFile("", "") + if err != nil { + t.Fatalf("getting temp file: %v", err) + } + enc := newImportLogEncoder(buf) + for _, test := range tests { + a := &test + err := enc.Encode(a) + if err != nil { + t.Errorf("encoding to buf: %v", err) + } + } + + name := buf.Name() + err = buf.Close() + if err != nil { + t.Fatalf("closing temp file: %v", err) + } + + buf, err = os.Open(name) + if err != nil { + t.Fatalf("reopening: %v", err) + } + + dec := newImportLogDecoder(buf) + for i, test := range tests { + t.Run(fmt.Sprintf("%d", i), func(t *testing.T) { + l := &importLog{} + err := dec.Decode(l) + // err := l.ReadFrom(buf) + if err != nil { + t.Errorf("reading from buf: %v", err) + } + if l.Index != test.Index { + t.Errorf("indexes not equal:\n%s\n%s", test.Index, l.Index) + } + if l.Path != test.Path { + t.Errorf("paths not equal:\n%s\n%s", test.Path, l.Path) + } + if l.Shard != test.Shard { + t.Errorf("shards not equal exp: %d got %d", test.Shard, l.Shard) + } + if !reflect.DeepEqual(test.Data, l.Data) { + t.Errorf("data not equal \n%v\n%v", test.Data, l.Data) + } + }) + } +} diff --git a/client/metrics.go b/client/metrics.go new file mode 100644 index 000000000..aaa14060b --- /dev/null +++ b/client/metrics.go @@ -0,0 +1,29 @@ +// Copyright 2017 Pilosa Corp. +// +// 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 client + +const ( + // MetricBatchImportDurationSeconds records the full time of the + // RecordBatch.Import call. This includes starting and finishing a + // transaction, doing key translation, building fragments locally, + // importing all data, and resetting internal structures. + MetricBatchImportDurationSeconds = "batch_import_duration_seconds" + + // MetricBatchFlushDurationSeconds records the full time for + // RecordBatch.Flush (if splitBatchMode is in use). This includes + // starting and finishing a transaction, importing all data, and + // resetting internal structures. + MetricBatchFlushDurationSeconds = "batch_flush_duration_seconds" +) diff --git a/client/orm.go b/client/orm.go new file mode 100644 index 000000000..71e44bf58 --- /dev/null +++ b/client/orm.go @@ -0,0 +1,1628 @@ +// Copyright 2017 Pilosa Corp. +// +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions +// are met: +// +// 1. Redistributions of source code must retain the above copyright +// notice, this list of conditions and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright +// notice, this list of conditions and the following disclaimer in the +// documentation and/or other materials provided with the distribution. +// +// 3. Neither the name of the copyright holder nor the names of its +// contributors may be used to endorse or promote products derived +// from this software without specific prior written permission. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND +// CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, +// INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF +// MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR +// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, +// BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH +// DAMAGE. + +package client + +import ( + "encoding/json" + "fmt" + "math" + "sort" + "strconv" + "strings" + "sync" + "time" + + "github.com/pilosa/pilosa/v2/pql" + "github.com/pkg/errors" +) + +const timeFormat = "2006-01-02T15:04" + +// Schema contains the index properties +type Schema struct { + mu sync.RWMutex + indexes map[string]*Index +} + +func (s *Schema) String() string { + s.mu.RLock() + defer s.mu.RUnlock() + return fmt.Sprintf("%s", s.indexes) +} + +// NewSchema creates a new Schema +func NewSchema() *Schema { + return &Schema{ + indexes: make(map[string]*Index), + } +} + +// Index returns an index with a name. +func (s *Schema) Index(name string, options ...IndexOption) *Index { + s.mu.Lock() + defer s.mu.Unlock() + if index, ok := s.indexes[name]; ok { + return index + } + indexOptions := &IndexOptions{} + indexOptions.addOptions(options...) + return s.indexWithOptions(name, 0, 0, indexOptions) +} + +func (s *Schema) indexWithOptions(name string, createdAt int64, shardWidth uint64, options *IndexOptions) *Index { + index := NewIndex(name) + if createdAt != 0 { + index.createdAt = createdAt + } + + index.options = options.withDefaults() + index.shardWidth = shardWidth + if index.Opts().TrackExistence() { + index.Field("_exists") + } + s.indexes[name] = index + return index +} + +// Indexes return a copy of the indexes in this schema +func (s *Schema) Indexes() map[string]*Index { + s.mu.RLock() + defer s.mu.RUnlock() + result := make(map[string]*Index) + for k, v := range s.indexes { + result[k] = v.copy() + } + return result +} + +// HasIndex returns true if the given index is in the schema. +func (s *Schema) HasIndex(indexName string) bool { + s.mu.RLock() + defer s.mu.RUnlock() + _, ok := s.indexes[indexName] + return ok +} + +func (s *Schema) diff(other *Schema) *Schema { + result := NewSchema() + for indexName, index := range s.indexes { + if otherIndex, ok := other.indexes[indexName]; !ok { + // if the index doesn't exist in the other schema, simply copy it + result.indexes[indexName] = index.copy() + } else { + // the index exists in the other schema; check the fields + resultIndex := NewIndex(indexName) + for fieldName, field := range index.fields { + if _, ok := otherIndex.fields[fieldName]; !ok { + // the field doesn't exist in the other schema, copy it + resultIndex.fields[fieldName] = field.copy() + } + } + // check whether we modified result index + if len(resultIndex.fields) > 0 { + // if so, move it to the result + result.indexes[indexName] = resultIndex + } + } + } + return result +} + +type SerializedQuery interface { + String() string + HasWriteKeys() bool +} + +type serializedQuery struct { + query string + hasWriteKeys bool +} + +func newSerializedQuery(query string, hasWriteKeys bool) serializedQuery { + return serializedQuery{ + query: query, + hasWriteKeys: hasWriteKeys, + } +} + +func (s serializedQuery) String() string { + return s.query +} + +func (s serializedQuery) HasWriteKeys() bool { + return s.hasWriteKeys +} + +// PQLQuery is an interface for PQL queries. +type PQLQuery interface { + Index() *Index + Serialize() SerializedQuery + Error() error +} + +// PQLBaseQuery is the base implementation for PQLQuery. +type PQLBaseQuery struct { + index *Index + pql string + err error + hasKeys bool +} + +// NewPQLBaseQuery creates a new PQLQuery with the given PQL and index. +func NewPQLBaseQuery(pql string, index *Index, err error) *PQLBaseQuery { + var hasKeys bool + if index != nil { + hasKeys = index.options.keys + } + return &PQLBaseQuery{ + index: index, + pql: pql, + err: err, + hasKeys: hasKeys, + } +} + +// Index returns the index for this query +func (q *PQLBaseQuery) Index() *Index { + return q.index +} + +func (q *PQLBaseQuery) Serialize() SerializedQuery { + return newSerializedQuery(q.pql, q.hasKeys) +} + +// Error returns the error or nil for this query. +func (q PQLBaseQuery) Error() error { + return q.err +} + +// PQLRowQuery is the return type for row queries. +type PQLRowQuery struct { + index *Index + pql string + err error + hasKeys bool +} + +// Index returns the index for this query/ +func (q *PQLRowQuery) Index() *Index { + return q.index +} + +func (q *PQLRowQuery) Serialize() SerializedQuery { + return q.serialize() +} + +func (q *PQLRowQuery) serialize() SerializedQuery { + return newSerializedQuery(q.pql, q.hasKeys) +} + +// Error returns the error or nil for this query. +func (q PQLRowQuery) Error() error { + return q.err +} + +// PQLBatchQuery contains a batch of PQL queries. +// Use Index.BatchQuery function to create an instance. +// +// Usage: +// +// repo, err := NewIndex("repository") +// stargazer, err := repo.Field("stargazer") +// query := repo.BatchQuery( +// stargazer.Row(5), +// stargazer.Row(15), +// repo.Union(stargazer.Row(20), stargazer.Row(25))) +type PQLBatchQuery struct { + index *Index + queries []string + err error + hasKeys bool +} + +// Index returns the index for this query. +func (q *PQLBatchQuery) Index() *Index { + return q.index +} + +func (q *PQLBatchQuery) Serialize() SerializedQuery { + query := strings.Join(q.queries, "") + return newSerializedQuery(query, q.hasKeys) +} + +func (q *PQLBatchQuery) Error() error { + return q.err +} + +// Add adds a query to the batch. +func (q *PQLBatchQuery) Add(query PQLQuery) { + err := query.Error() + if err != nil { + q.err = err + } + serializedQuery := query.Serialize() + q.hasKeys = q.hasKeys || serializedQuery.HasWriteKeys() + q.queries = append(q.queries, serializedQuery.String()) +} + +// NewPQLRowQuery creates a new PqlRowQuery. +func NewPQLRowQuery(pql string, index *Index, err error) *PQLRowQuery { + return &PQLRowQuery{ + index: index, + pql: pql, + err: err, + hasKeys: index.options.keys, + } +} + +// IndexOptions contains options to customize Index objects. +type IndexOptions struct { + keys bool + keysSet bool + trackExistence bool + trackExistenceSet bool +} + +func (io *IndexOptions) withDefaults() (updated *IndexOptions) { + // copy options so the original is not updated + updated = &IndexOptions{} + *updated = *io + if !updated.keysSet { + updated.keys = false + } + if !updated.trackExistenceSet { + updated.trackExistence = true + } + return +} + +// Keys return true if this index has keys. +func (io IndexOptions) Keys() bool { + return io.keys +} + +// TrackExistence returns true if existence is tracked for this index. +func (io IndexOptions) TrackExistence() bool { + return io.trackExistence +} + +// String serializes this index to a JSON string. +func (io IndexOptions) String() string { + mopt := map[string]interface{}{} + if io.keysSet { + mopt["keys"] = io.keys + } + if io.trackExistenceSet { + mopt["trackExistence"] = io.trackExistence + } + return fmt.Sprintf(`{"options":%s}`, encodeMap(mopt)) +} + +func (io *IndexOptions) addOptions(options ...IndexOption) { + for _, option := range options { + if option == nil { + continue + } + option(io) + } +} + +// IndexOption is used to pass an option to Index function. +type IndexOption func(options *IndexOptions) + +// OptIndexKeys sets whether index uses string keys. +func OptIndexKeys(keys bool) IndexOption { + return func(options *IndexOptions) { + options.keys = keys + options.keysSet = true + } +} + +// OptIndexTrackExistence enables keeping track of existence of columns. +func OptIndexTrackExistence(trackExistence bool) IndexOption { + return func(options *IndexOptions) { + options.trackExistence = trackExistence + options.trackExistenceSet = true + } +} + +// OptionsOptions is used to pass an option to Option call. +type OptionsOptions struct { + columnAttrs bool + excludeColumns bool + excludeRowAttrs bool + shards []uint64 +} + +func (oo OptionsOptions) marshal() string { + part1 := fmt.Sprintf("columnAttrs=%s,excludeColumns=%s,excludeRowAttrs=%s", + strconv.FormatBool(oo.columnAttrs), + strconv.FormatBool(oo.excludeColumns), + strconv.FormatBool(oo.excludeRowAttrs)) + if oo.shards != nil { + shardsStr := make([]string, len(oo.shards)) + for i, shard := range oo.shards { + shardsStr[i] = strconv.FormatUint(shard, 10) + } + return fmt.Sprintf("%s,shards=[%s]", part1, strings.Join(shardsStr, ",")) + } + return part1 +} + +// OptionsOption is an option for Index.Options call. +type OptionsOption func(options *OptionsOptions) + +// OptOptionsColumnAttrs enables returning column attributes. +func OptOptionsColumnAttrs(enable bool) OptionsOption { + return func(options *OptionsOptions) { + options.columnAttrs = enable + } +} + +// OptOptionsExcludeColumns enables preventing returning columns. +func OptOptionsExcludeColumns(enable bool) OptionsOption { + return func(options *OptionsOptions) { + options.excludeColumns = enable + } +} + +// OptOptionsExcludeRowAttrs enables preventing returning row attributes. +func OptOptionsExcludeRowAttrs(enable bool) OptionsOption { + return func(options *OptionsOptions) { + options.excludeRowAttrs = enable + } +} + +// OptOptionsShards run the query using only the data from the given shards. +// By default, the entire data set (i.e. data from all shards) is used. +func OptOptionsShards(shards ...uint64) OptionsOption { + return func(options *OptionsOptions) { + options.shards = shards + } +} + +// Index is a Pilosa index. The purpose of the Index is to represent a data namespace. +// You cannot perform cross-index queries. Column-level attributes are global to the Index. +type Index struct { + mu sync.RWMutex + name string + createdAt int64 + options *IndexOptions + fields map[string]*Field + shardWidth uint64 +} + +func (idx *Index) String() string { + return fmt.Sprintf(`{name: "%s", options: "%s", fields: %s, shardWidth: %d}`, idx.name, idx.options, idx.fields, idx.shardWidth) +} + +// NewIndex creates an index with a name. +func NewIndex(name string) *Index { + options := &IndexOptions{} + return &Index{ + name: name, + options: options.withDefaults(), + fields: map[string]*Field{}, + } +} + +func (idx *Index) ShardWidth() uint64 { + return idx.shardWidth +} + +// Fields return a copy of the fields in this index +func (idx *Index) Fields() map[string]*Field { + idx.mu.Lock() + defer idx.mu.Unlock() + result := make(map[string]*Field) + for k, v := range idx.fields { + result[k] = v.copy() + } + return result +} + +// HasFields returns true if the given field exists in the index. +func (idx *Index) HasField(fieldName string) bool { + idx.mu.Lock() + defer idx.mu.Unlock() + _, ok := idx.fields[fieldName] + return ok +} + +func (idx *Index) copy() *Index { + idx.mu.Lock() + defer idx.mu.Unlock() + fields := make(map[string]*Field) + for name, f := range idx.fields { + fields[name] = f.copy() + } + index := &Index{ + name: idx.name, + createdAt: idx.createdAt, + options: &IndexOptions{}, + fields: fields, + shardWidth: idx.shardWidth, + } + *index.options = *idx.options + return index +} + +// Name returns the name of this index. +func (idx *Index) Name() string { + return idx.name +} + +func (idx *Index) CreatedAt() int64 { + idx.mu.RLock() + defer idx.mu.RUnlock() + return idx.createdAt +} + +// Opts returns the options of this index. +func (idx *Index) Opts() IndexOptions { + return *idx.options +} + +// Field creates a Field struct with the specified name and defaults. +func (idx *Index) Field(name string, options ...FieldOption) *Field { + idx.mu.Lock() + defer idx.mu.Unlock() + if field, ok := idx.fields[name]; ok { + return field + } + fieldOptions := &FieldOptions{} + fieldOptions = fieldOptions.withDefaults() + fieldOptions.addOptions(options...) + return idx.fieldWithOptions(name, 0, fieldOptions) +} + +func (idx *Index) fieldWithOptions(name string, createdAt int64, fieldOptions *FieldOptions) *Field { + field := newField(name, idx) + if createdAt != 0 { + field.createdAt = createdAt + } + fieldOptions = fieldOptions.withDefaults() + field.options = fieldOptions + idx.fields[name] = field + return field +} + +// BatchQuery creates a batch query with the given queries. +func (idx *Index) BatchQuery(queries ...PQLQuery) *PQLBatchQuery { + stringQueries := make([]string, 0, len(queries)) + hasKeys := false + for _, query := range queries { + serializedQuery := query.Serialize() + hasKeys = hasKeys || serializedQuery.HasWriteKeys() + stringQueries = append(stringQueries, serializedQuery.String()) + } + return &PQLBatchQuery{ + index: idx, + queries: stringQueries, + hasKeys: hasKeys, + } +} + +// RawQuery creates a query with the given string. +// Note that the query is not validated before sending to the server. +func (idx *Index) RawQuery(query string) *PQLBaseQuery { + q := NewPQLBaseQuery(query, idx, nil) + // NOTE: raw queries always assumed to have keys set + q.hasKeys = true + return q +} + +// Union creates a Union query. +// Union performs a logical OR on the results of each ROW_CALL query passed to it. +func (idx *Index) Union(rows ...*PQLRowQuery) *PQLRowQuery { + return idx.rowOperation("Union", rows...) +} + +// Intersect creates an Intersect query. +// Intersect performs a logical AND on the results of each ROW_CALL query passed to it. +func (idx *Index) Intersect(rows ...*PQLRowQuery) *PQLRowQuery { + if len(rows) < 1 { + return NewPQLRowQuery("", idx, errors.New("Intersect operation requires at least 1 row")) + } + return idx.rowOperation("Intersect", rows...) +} + +// Difference creates an Intersect query. +// Difference returns all of the columns from the first ROW_CALL argument passed to it, without the columns from each subsequent ROW_CALL. +func (idx *Index) Difference(rows ...*PQLRowQuery) *PQLRowQuery { + if len(rows) < 1 { + return NewPQLRowQuery("", idx, errors.New("Difference operation requires at least 1 row")) + } + return idx.rowOperation("Difference", rows...) +} + +// Xor creates an Xor query. +func (idx *Index) Xor(rows ...*PQLRowQuery) *PQLRowQuery { + if len(rows) < 2 { + return NewPQLRowQuery("", idx, errors.New("Xor operation requires at least 2 rows")) + } + return idx.rowOperation("Xor", rows...) +} + +// Not creates a Not query. +func (idx *Index) Not(row *PQLRowQuery) *PQLRowQuery { + return NewPQLRowQuery(fmt.Sprintf("Not(%s)", row.serialize()), idx, row.Error()) +} + +// Count creates a Count query. +// Returns the number of set columns in the ROW_CALL passed in. +func (idx *Index) Count(row *PQLRowQuery) *PQLBaseQuery { + serializedQuery := row.serialize() + q := NewPQLBaseQuery(fmt.Sprintf("Count(%s)", serializedQuery.String()), idx, nil) + q.hasKeys = q.hasKeys || serializedQuery.HasWriteKeys() + return q +} + +// All creates an All query. +// Returns the set columns with existence true. +func (idx *Index) All() *PQLRowQuery { + q := NewPQLRowQuery("All()", idx, nil) + return q +} + +// TODO: impelement AllLimit(limit, offset uint64) *PQLRowQuery + +// SetColumnAttrs creates a SetColumnAttrs query. +// SetColumnAttrs associates arbitrary key/value pairs with a column in an index. +// Following types are accepted: integer, float, string and boolean types. +func (idx *Index) SetColumnAttrs(colIDOrKey interface{}, attrs map[string]interface{}) *PQLBaseQuery { + colStr, err := formatIDKey(colIDOrKey) + if err != nil { + return NewPQLBaseQuery("", idx, err) + } + attrsString, err := createAttributesString(attrs) + if err != nil { + return NewPQLBaseQuery("", idx, err) + } + q := fmt.Sprintf("SetColumnAttrs(%s,%s)", colStr, attrsString) + return NewPQLBaseQuery(q, idx, nil) +} + +// Options creates an Options query. +func (idx *Index) Options(row *PQLRowQuery, opts ...OptionsOption) *PQLBaseQuery { + oo := &OptionsOptions{} + for _, opt := range opts { + opt(oo) + } + text := fmt.Sprintf("Options(%s,%s)", row.serialize(), oo.marshal()) + return NewPQLBaseQuery(text, idx, nil) +} + +type groupByBuilder struct { + rows []*PQLRowsQuery + limit int64 + filter *PQLRowQuery + aggregate *PQLBaseQuery + having *PQLBaseQuery +} + +// GroupByBuilderOption is a functional option type for index.GroupBy +type GroupByBuilderOption func(g *groupByBuilder) error + +// OptGroupByBuilderRows is a functional option on groupByBuilder +// used to set the rows. +func OptGroupByBuilderRows(rows ...*PQLRowsQuery) GroupByBuilderOption { + return func(g *groupByBuilder) error { + g.rows = rows + return nil + } +} + +// OptGroupByBuilderLimit is a functional option on groupByBuilder +// used to set the limit. +func OptGroupByBuilderLimit(l int64) GroupByBuilderOption { + return func(g *groupByBuilder) error { + g.limit = l + return nil + } +} + +// OptGroupByBuilderFilter is a functional option on groupByBuilder +// used to set the filter. +func OptGroupByBuilderFilter(q *PQLRowQuery) GroupByBuilderOption { + return func(g *groupByBuilder) error { + g.filter = q + return nil + } +} + +// OptGroupByBuilderAggregate is a functional option on groupByBuilder +// used to set the aggregate. +func OptGroupByBuilderAggregate(agg *PQLBaseQuery) GroupByBuilderOption { + return func(g *groupByBuilder) error { + g.aggregate = agg + return nil + } +} + +// OptGroupByBuilderHaving is a functional option on groupByBuilder +// used to set the having clause. +func OptGroupByBuilderHaving(having *PQLBaseQuery) GroupByBuilderOption { + return func(g *groupByBuilder) error { + g.having = having + return nil + } +} + +// GroupByBase creates a GroupBy query with the given functional options. +func (idx *Index) GroupByBase(opts ...GroupByBuilderOption) *PQLBaseQuery { + bldr := &groupByBuilder{} + for _, opt := range opts { + err := opt(bldr) + if err != nil { + return NewPQLBaseQuery("", idx, errors.Wrap(err, "applying option")) + } + } + + if len(bldr.rows) < 1 { + return NewPQLBaseQuery("", idx, errors.New("there should be at least one rows query")) + } + if bldr.limit < 0 { + return NewPQLBaseQuery("", idx, errors.New("limit must be non-negative")) + } + + // rows + text := fmt.Sprintf("GroupBy(%s", strings.Join(serializeGroupBy(bldr.rows...), ",")) + + // limit + if bldr.limit > 0 { + text += fmt.Sprintf(",limit=%d", bldr.limit) + } + + // filter + if bldr.filter != nil { + filterText := bldr.filter.serialize().String() + text += fmt.Sprintf(",filter=%s", filterText) + } + + // aggregate + if bldr.aggregate != nil { + aggregateText := bldr.aggregate.Serialize().String() + text += fmt.Sprintf(",aggregate=%s", aggregateText) + } + + // having + if bldr.having != nil { + havingText := bldr.having.Serialize().String() + text += fmt.Sprintf(",having=%s", havingText) + } + + text += ")" + return NewPQLBaseQuery(text, idx, nil) +} + +// GroupBy creates a GroupBy query with the given Rows queries +func (idx *Index) GroupBy(rowsQueries ...*PQLRowsQuery) *PQLBaseQuery { + if len(rowsQueries) < 1 { + return NewPQLBaseQuery("", idx, errors.New("there should be at least one rows query")) + } + text := fmt.Sprintf("GroupBy(%s)", strings.Join(serializeGroupBy(rowsQueries...), ",")) + return NewPQLBaseQuery(text, idx, nil) +} + +// GroupByLimit creates a GroupBy query with the given limit and Rows queries +func (idx *Index) GroupByLimit(limit int64, rowsQueries ...*PQLRowsQuery) *PQLBaseQuery { + if len(rowsQueries) < 1 { + return NewPQLBaseQuery("", idx, errors.New("there should be at least one rows query")) + } + if limit < 0 { + return NewPQLBaseQuery("", idx, errors.New("limit must be non-negative")) + } + text := fmt.Sprintf("GroupBy(%s,limit=%d)", strings.Join(serializeGroupBy(rowsQueries...), ","), limit) + return NewPQLBaseQuery(text, idx, nil) +} + +// GroupByFilter creates a GroupBy query with the given filter and Rows queries +func (idx *Index) GroupByFilter(filterQuery *PQLRowQuery, rowsQueries ...*PQLRowsQuery) *PQLBaseQuery { + if len(rowsQueries) < 1 { + return NewPQLBaseQuery("", idx, errors.New("there should be at least one rows query")) + } + filterText := filterQuery.serialize().String() + text := fmt.Sprintf("GroupBy(%s,filter=%s)", strings.Join(serializeGroupBy(rowsQueries...), ","), filterText) + return NewPQLBaseQuery(text, idx, nil) +} + +// GroupByLimitFilter creates a GroupBy query with the given filter and Rows queries +func (idx *Index) GroupByLimitFilter(limit int64, filterQuery *PQLRowQuery, rowsQueries ...*PQLRowsQuery) *PQLBaseQuery { + if len(rowsQueries) < 1 { + return NewPQLBaseQuery("", idx, errors.New("there should be at least one rows query")) + } + if limit < 0 { + return NewPQLBaseQuery("", idx, errors.New("limit must be non-negative")) + } + filterText := filterQuery.serialize().String() + text := fmt.Sprintf("GroupBy(%s,limit=%d,filter=%s)", strings.Join(serializeGroupBy(rowsQueries...), ","), limit, filterText) + return NewPQLBaseQuery(text, idx, nil) +} + +func (idx *Index) rowOperation(name string, rows ...*PQLRowQuery) *PQLRowQuery { + var err error + args := make([]string, 0, len(rows)) + for _, row := range rows { + if err = row.Error(); err != nil { + return NewPQLRowQuery("", idx, err) + } + args = append(args, row.serialize().String()) + } + query := NewPQLRowQuery(fmt.Sprintf("%s(%s)", name, strings.Join(args, ",")), idx, nil) + return query +} + +func serializeGroupBy(rowsQueries ...*PQLRowsQuery) []string { + qs := make([]string, 0, len(rowsQueries)) + for _, qry := range rowsQueries { + qs = append(qs, qry.serialize().String()) + } + return qs +} + +// FieldInfo represents schema information for a field. +type FieldInfo struct { + Name string `json:"name"` +} + +// FieldOptions contains options to customize Field objects and field queries. +type FieldOptions struct { + fieldType FieldType + timeQuantum TimeQuantum + cacheType CacheType + cacheSize int + min pql.Decimal + max pql.Decimal + scale int64 + keys bool + noStandardView bool + foreignIndex string +} + +// Type returns the type of the field. Currently "set", "int", or "time". +func (fo FieldOptions) Type() FieldType { + return fo.fieldType +} + +// TimeQuantum returns the configured time quantum for a time field. Empty +// string otherwise. +func (fo FieldOptions) TimeQuantum() TimeQuantum { + return fo.timeQuantum +} + +// CacheType returns the configured cache type for a "set" field. Empty string +// otherwise. +func (fo FieldOptions) CacheType() CacheType { + return fo.cacheType +} + +// CacheSize returns the cache size for a set field. Zero otherwise. +func (fo FieldOptions) CacheSize() int { + return fo.cacheSize +} + +// Min returns the minimum accepted value for an integer field. Zero otherwise. +func (fo FieldOptions) Min() pql.Decimal { + return fo.min +} + +// Max returns the maximum accepted value for an integer field. Zero otherwise. +func (fo FieldOptions) Max() pql.Decimal { + return fo.max +} + +// Scale returns the scale for a decimal field. +func (fo FieldOptions) Scale() int64 { + return fo.scale +} + +// Keys returns whether this field uses keys instead of IDs +func (fo FieldOptions) Keys() bool { + return fo.keys +} + +func (fo FieldOptions) ForeignIndex() string { + return fo.foreignIndex +} + +// NoStandardView suppresses creating the standard view for supported field types (currently, time) +func (fo FieldOptions) NoStandardView() bool { + return fo.noStandardView +} + +func (fo *FieldOptions) withDefaults() (updated *FieldOptions) { + // copy options so the original is not updated + updated = &FieldOptions{} + *updated = *fo + if updated.fieldType == "" { + updated.fieldType = FieldTypeSet + } + return +} + +func (fo FieldOptions) String() string { + mopt := map[string]interface{}{} + + switch fo.fieldType { + case FieldTypeSet, FieldTypeMutex: + if fo.cacheType != CacheTypeDefault { + mopt["cacheType"] = string(fo.cacheType) + } + if fo.cacheSize > 0 { + mopt["cacheSize"] = fo.cacheSize + } + case FieldTypeInt: + mopt["min"] = fo.min + mopt["max"] = fo.max + case FieldTypeDecimal: + mopt["min"] = fo.min + mopt["max"] = fo.max + mopt["scale"] = fo.scale + case FieldTypeTime: + mopt["timeQuantum"] = string(fo.timeQuantum) + mopt["noStandardView"] = fo.noStandardView + } + + if fo.fieldType != FieldTypeDefault { + mopt["type"] = string(fo.fieldType) + } + if fo.keys { + mopt["keys"] = fo.keys + } + if fo.foreignIndex != "" { + mopt["foreignIndex"] = fo.foreignIndex + } + return fmt.Sprintf(`{"options":%s}`, encodeMap(mopt)) +} + +func (fo *FieldOptions) addOptions(options ...FieldOption) { + for _, option := range options { + if option == nil { + continue + } + option(fo) + } +} + +// FieldOption is used to pass an option to index.Field function. +type FieldOption func(options *FieldOptions) + +// OptFieldTypeSet adds a set field. +// Specify CacheTypeDefault for the default cache type. +// Specify CacheSizeDefault for the default cache size. +func OptFieldTypeSet(cacheType CacheType, cacheSize int) FieldOption { + return func(options *FieldOptions) { + options.fieldType = FieldTypeSet + options.cacheType = cacheType + options.cacheSize = cacheSize + } +} + +// OptFieldTypeInt adds an integer field. +// No arguments: min = min_int, max = max_int +// 1 argument: min = limit[0], max = max_int +// 2 or more arguments: min = limit[0], max = limit[1] +func OptFieldTypeInt(limits ...int64) FieldOption { + min := pql.NewDecimal(math.MinInt64, 0) + max := pql.NewDecimal(math.MaxInt64, 0) + + if len(limits) > 2 { + panic("error: OptFieldTypeInt accepts at most 2 arguments") + } + if len(limits) > 0 { + min = pql.NewDecimal(limits[0], 0) + } + if len(limits) > 1 { + max = pql.NewDecimal(limits[1], 0) + } + + return func(options *FieldOptions) { + options.fieldType = FieldTypeInt + options.min = min + options.max = max + } +} + +// OptFieldTypeTime adds a time field. +func OptFieldTypeTime(quantum TimeQuantum, opts ...bool) FieldOption { + return func(options *FieldOptions) { + options.fieldType = FieldTypeTime + options.timeQuantum = quantum + if len(opts) > 0 && opts[0] { + options.noStandardView = true + } + } +} + +// OptFieldTypeMutex adds a mutex field. +func OptFieldTypeMutex(cacheType CacheType, cacheSize int) FieldOption { + return func(options *FieldOptions) { + options.fieldType = FieldTypeMutex + options.cacheType = cacheType + options.cacheSize = cacheSize + } +} + +// OptFieldTypeBool adds a bool field. +func OptFieldTypeBool() FieldOption { + return func(options *FieldOptions) { + options.fieldType = FieldTypeBool + } +} + +func OptFieldTypeDecimal(scale int64, minmax ...pql.Decimal) FieldOption { + min, max := pql.MinMax(scale) + if len(minmax) > 2 { + panic("error: OptFieldTypeDecimal accepts at most 2 arguments") + } + if len(minmax) > 0 { + min = minmax[0] + } + if len(minmax) > 1 { + max = minmax[1] + } + return func(options *FieldOptions) { + options.fieldType = FieldTypeDecimal + options.scale = scale + options.min = min + options.max = max + } +} + +// OptFieldKeys sets whether field uses string keys. +func OptFieldKeys(keys bool) FieldOption { + return func(options *FieldOptions) { + options.keys = keys + } +} + +func OptFieldForeignIndex(index string) FieldOption { + return func(options *FieldOptions) { + options.foreignIndex = index + } +} + +// Field structs are used to segment and define different functional characteristics within your entire index. +// You can think of a Field as a table-like data partition within your Index. +// Row-level attributes are namespaced at the Field level. +type Field struct { + name string + createdAt int64 + index *Index + options *FieldOptions +} + +func (f *Field) String() string { + return fmt.Sprintf(`{name: "%s", index: "%s", options: "%s"}`, f.name, f.index.name, f.options) +} + +func newField(name string, index *Index) *Field { + return &Field{ + name: name, + index: index, + options: &FieldOptions{}, + } +} + +// Name returns the name of the field +func (f *Field) Name() string { + return f.name +} + +func (f *Field) CreatedAt() int64 { + return f.createdAt +} + +// Opts returns the options of the field +func (f *Field) Opts() FieldOptions { + return *f.options +} + +func (f *Field) copy() *Field { + field := newField(f.name, f.index) + field.createdAt = f.createdAt + *field.options = *f.options + return field +} + +// Row creates a Row query. +// Row retrieves the indices of all the set columns in a row. +// It also retrieves any attributes set on that row or column. +func (f *Field) Row(rowIDOrKey interface{}) *PQLRowQuery { + rowStr, err := formatIDKeyBool(rowIDOrKey) + if err != nil { + return NewPQLRowQuery("", f.index, err) + } + text := fmt.Sprintf("Row(%s=%s)", f.name, rowStr) + q := NewPQLRowQuery(text, f.index, nil) + return q +} + +// Set creates a Set query. +// Set, assigns a value of 1 to a bit in the binary matrix, thus associating the given row in the given field with the given column. +func (f *Field) Set(rowIDOrKey, colIDOrKey interface{}) *PQLBaseQuery { + rowStr, colStr, err := formatRowColIDKey(rowIDOrKey, colIDOrKey) + if err != nil { + return NewPQLBaseQuery("", f.index, err) + } + text := fmt.Sprintf("Set(%s,%s=%s)", colStr, f.name, rowStr) + q := NewPQLBaseQuery(text, f.index, nil) + q.hasKeys = f.options.keys || f.index.options.keys + return q +} + +// SetTimestamp creates a Set query with timestamp. +// Set, assigns a value of 1 to a column in the binary matrix, +// thus associating the given row in the given field with the given column. +func (f *Field) SetTimestamp(rowIDOrKey, colIDOrKey interface{}, timestamp time.Time) *PQLBaseQuery { + rowStr, colStr, err := formatRowColIDKey(rowIDOrKey, colIDOrKey) + if err != nil { + return NewPQLBaseQuery("", f.index, err) + } + text := fmt.Sprintf("Set(%s,%s=%s,%s)", colStr, f.name, rowStr, timestamp.Format(timeFormat)) + q := NewPQLBaseQuery(text, f.index, nil) + q.hasKeys = f.options.keys || f.index.options.keys + return q +} + +// Clear creates a Clear query. +// Clear, assigns a value of 0 to a bit in the binary matrix, thus disassociating the given row in the given field from the given column. +func (f *Field) Clear(rowIDOrKey, colIDOrKey interface{}) *PQLBaseQuery { + rowStr, colStr, err := formatRowColIDKey(rowIDOrKey, colIDOrKey) + if err != nil { + return NewPQLBaseQuery("", f.index, err) + } + text := fmt.Sprintf("Clear(%s,%s=%s)", colStr, f.name, rowStr) + q := NewPQLBaseQuery(text, f.index, nil) + q.hasKeys = f.options.keys || f.index.options.keys + return q +} + +// ClearRow creates a ClearRow query. +// ClearRow sets all bits to 0 in a given row of the binary matrix, thus disassociating the given row in the given field from all columns. +func (f *Field) ClearRow(rowIDOrKey interface{}) *PQLBaseQuery { + rowStr, err := formatIDKeyBool(rowIDOrKey) + if err != nil { + return NewPQLBaseQuery("", f.index, err) + } + text := fmt.Sprintf("ClearRow(%s=%s)", f.name, rowStr) + q := NewPQLBaseQuery(text, f.index, nil) + return q +} + +// TopN creates a TopN query with the given item count. +// Returns the id and count of the top n rows (by count of columns) in the field. +func (f *Field) TopN(n uint64) *PQLRowQuery { + q := NewPQLRowQuery(fmt.Sprintf("TopN(%s,n=%d)", f.name, n), f.index, nil) + return q +} + +// RowTopN creates a TopN query with the given item count and row. +// This variant supports customizing the row query. +func (f *Field) RowTopN(n uint64, row *PQLRowQuery) *PQLRowQuery { + q := NewPQLRowQuery(fmt.Sprintf("TopN(%s,%s,n=%d)", + f.name, row.serialize(), n), f.index, nil) + return q +} + +// FilterAttrTopN creates a TopN query with the given item count, row, attribute name and filter values for that field +// The attrName and attrValues arguments work together to only return Rows which have the attribute specified by attrName with one of the values specified in attrValues. +func (f *Field) FilterAttrTopN(n uint64, row *PQLRowQuery, attrName string, attrValues ...interface{}) *PQLRowQuery { + return f.filterAttrTopN(n, row, attrName, attrValues...) +} + +func (f *Field) filterAttrTopN(n uint64, row *PQLRowQuery, field string, values ...interface{}) *PQLRowQuery { + if err := validateLabel(field); err != nil { + return NewPQLRowQuery("", f.index, err) + } + b, err := json.Marshal(values) + if err != nil { + return NewPQLRowQuery("", f.index, err) + } + var q *PQLRowQuery + if row == nil { + q = NewPQLRowQuery(fmt.Sprintf("TopN(%s,n=%d,attrName='%s',attrValues=%s)", + f.name, n, field, string(b)), f.index, nil) + } else { + serializedRow := row.serialize() + q = NewPQLRowQuery(fmt.Sprintf("TopN(%s,%s,n=%d,attrName='%s',attrValues=%s)", + f.name, serializedRow.String(), n, field, string(b)), f.index, nil) + } + return q +} + +// Range creates a Range query. +// Similar to Row, but only returns columns which were set with timestamps between the given start and end timestamps. +// *Deprecated at Pilosa 1.3* +func (f *Field) Range(rowIDOrKey interface{}, start time.Time, end time.Time) *PQLRowQuery { + rowStr, err := formatIDKeyBool(rowIDOrKey) + if err != nil { + return NewPQLRowQuery("", f.index, err) + } + text := fmt.Sprintf("Range(%s=%s,%s,%s)", f.name, rowStr, start.Format(timeFormat), end.Format(timeFormat)) + q := NewPQLRowQuery(text, f.index, nil) + return q +} + +// RowRange creates a Row query with timestamps. +// Similar to Row, but only returns columns which were set with timestamps between the given start and end timestamps. +// *Introduced at Pilosa 1.3* +func (f *Field) RowRange(rowIDOrKey interface{}, start time.Time, end time.Time) *PQLRowQuery { + rowStr, err := formatIDKeyBool(rowIDOrKey) + if err != nil { + return NewPQLRowQuery("", f.index, err) + } + text := fmt.Sprintf("Row(%s=%s,from='%s',to='%s')", f.name, rowStr, start.Format(timeFormat), end.Format(timeFormat)) + q := NewPQLRowQuery(text, f.index, nil) + return q +} + +// SetRowAttrs creates a SetRowAttrs query. +// SetRowAttrs associates arbitrary key/value pairs with a row in a field. +// Following types are accepted: integer, float, string and boolean types. +func (f *Field) SetRowAttrs(rowIDOrKey interface{}, attrs map[string]interface{}) *PQLBaseQuery { + rowStr, err := formatIDKeyBool(rowIDOrKey) + if err != nil { + return NewPQLBaseQuery("", f.index, err) + } + attrsString, err := createAttributesString(attrs) + if err != nil { + return NewPQLBaseQuery("", f.index, err) + } + text := fmt.Sprintf("SetRowAttrs(%s,%s,%s)", f.name, rowStr, attrsString) + q := NewPQLBaseQuery(text, f.index, nil) + q.hasKeys = f.options.keys || f.index.options.keys + return q +} + +// Store creates a Store call. +// Store writes the result of the row query to the specified row. If the row already exists, it will be replaced. The destination field must be of field type set. +func (f *Field) Store(row *PQLRowQuery, rowIDOrKey interface{}) *PQLBaseQuery { + rowStr, err := formatIDKeyBool(rowIDOrKey) + if err != nil { + return NewPQLBaseQuery("", f.index, err) + } + return NewPQLBaseQuery(fmt.Sprintf("Store(%s,%s=%s)", row.serialize().String(), f.name, rowStr), f.index, nil) +} + +func createAttributesString(attrs map[string]interface{}) (string, error) { + attrsList := make([]string, 0, len(attrs)) + for k, v := range attrs { + // TODO: validate the type of v is one of string, int64, float64, bool + if err := validateLabel(k); err != nil { + return "", err + } + if vs, ok := v.(string); ok { + attrsList = append(attrsList, fmt.Sprintf("%s=%s", k, strconv.Quote(vs))) + } else { + attrsList = append(attrsList, fmt.Sprintf("%s=%v", k, v)) + } + } + sort.Strings(attrsList) + return strings.Join(attrsList, ","), nil +} + +func formatIDKey(idKey interface{}) (string, error) { + switch v := idKey.(type) { + case uint: + return strconv.FormatUint(uint64(v), 10), nil + case uint32: + return strconv.FormatUint(uint64(v), 10), nil + case uint64: + return strconv.FormatUint(v, 10), nil + case int: + return strconv.FormatInt(int64(v), 10), nil + case int32: + return strconv.FormatInt(int64(v), 10), nil + case int64: + return strconv.FormatInt(v, 10), nil + case string: + v = strings.ReplaceAll(v, `\`, `\\`) + return fmt.Sprintf(`'%s'`, strings.ReplaceAll(v, `'`, `\'`)), nil + default: + return "", errors.Errorf("id/key is not a string or integer type: %#v", idKey) + } +} + +func formatIDKeyBool(idKeyBool interface{}) (string, error) { + if b, ok := idKeyBool.(bool); ok { + return strconv.FormatBool(b), nil + } + if flt, ok := idKeyBool.(float64); ok { + return fmt.Sprintf("%f", flt), nil + } + return formatIDKey(idKeyBool) +} + +func formatRowColIDKey(rowIDOrKey, colIDOrKey interface{}) (string, string, error) { + rowStr, err := formatIDKeyBool(rowIDOrKey) + if err != nil { + return "", "", errors.Wrap(err, "formatting row") + } + colStr, err := formatIDKey(colIDOrKey) + if err != nil { + return "", "", errors.Wrap(err, "formatting column") + } + return rowStr, colStr, err +} + +// FieldType is the type of a field. +// See: https://www.pilosa.com/docs/latest/data-model/#field-type +type FieldType string + +const ( + // FieldTypeDefault is the default field type. + FieldTypeDefault FieldType = "" + // FieldTypeSet is the set field type. + // See: https://www.pilosa.com/docs/latest/data-model/#set + FieldTypeSet FieldType = "set" + // FieldTypeInt is the int field type. + // See: https://www.pilosa.com/docs/latest/data-model/#int + FieldTypeInt FieldType = "int" + // FieldTypeTime is the time field type. + // See: https://www.pilosa.com/docs/latest/data-model/#time + FieldTypeTime FieldType = "time" + // FieldTypeMutex is the mutex field type. + // See: https://www.pilosa.com/docs/latest/data-model/#mutex + FieldTypeMutex FieldType = "mutex" + // FieldTypeBool is the boolean field type. + // See: https://www.pilosa.com/docs/latest/data-model/#boolean + FieldTypeBool FieldType = "bool" + // FieldTypeDecimal can store floating point numbers as integers + // with a scale factor. This field type is only available in + // Molecula's Pilosa with enterprise extensions. + FieldTypeDecimal FieldType = "decimal" +) + +// TimeQuantum type represents valid time quantum values time fields. +type TimeQuantum string + +// TimeQuantum constants +const ( + TimeQuantumNone TimeQuantum = "" + TimeQuantumYear TimeQuantum = "Y" + TimeQuantumMonth TimeQuantum = "M" + TimeQuantumDay TimeQuantum = "D" + TimeQuantumHour TimeQuantum = "H" + TimeQuantumYearMonth TimeQuantum = "YM" + TimeQuantumMonthDay TimeQuantum = "MD" + TimeQuantumDayHour TimeQuantum = "DH" + TimeQuantumYearMonthDay TimeQuantum = "YMD" + TimeQuantumMonthDayHour TimeQuantum = "MDH" + TimeQuantumYearMonthDayHour TimeQuantum = "YMDH" +) + +// CacheType represents cache type for a field +type CacheType string + +// CacheType constants +const ( + CacheTypeDefault CacheType = "" + CacheTypeLRU CacheType = "lru" + CacheTypeRanked CacheType = "ranked" + CacheTypeNone CacheType = "none" +) + +// CacheSizeDefault is the default cache size +const CacheSizeDefault = 0 + +// Options returns the options set for the field. Which fields of the +// FieldOptions struct are actually being used depends on the field's type. +// *DEPRECATED* +func (f *Field) Options() *FieldOptions { + return f.options +} + +type IntOrFloat interface{} + +type intOrFloatVal struct { + IntOrFloat +} + +func (i intOrFloatVal) String() string { + switch i.IntOrFloat.(type) { + case float64: + // In order to test expected values, we set the precision + // to 8. TODO: It's likely we'll need to address this + // at some point. + return fmt.Sprintf("%.8f", i.IntOrFloat) + default: + return fmt.Sprintf("%d", i.IntOrFloat) + } +} + +// LT creates a less than query. +func (f *Field) LT(n IntOrFloat) *PQLRowQuery { + return f.binaryOperation("<", n) +} + +// LTE creates a less than or equal query. +func (f *Field) LTE(n IntOrFloat) *PQLRowQuery { + return f.binaryOperation("<=", n) +} + +// GT creates a greater than query. +func (f *Field) GT(n IntOrFloat) *PQLRowQuery { + return f.binaryOperation(">", n) +} + +// GTE creates a greater than or equal query. +func (f *Field) GTE(n IntOrFloat) *PQLRowQuery { + return f.binaryOperation(">=", n) +} + +// Equals creates an equals query. +func (f *Field) Equals(n IntOrFloat) *PQLRowQuery { + return f.binaryOperation("==", n) +} + +// NotEquals creates a not equals query. +func (f *Field) NotEquals(n IntOrFloat) *PQLRowQuery { + return f.binaryOperation("!=", n) +} + +// NotNull creates a not equal to null query. +func (f *Field) NotNull() *PQLRowQuery { + text := fmt.Sprintf("Row(%s != null)", f.name) + q := NewPQLRowQuery(text, f.index, nil) + q.hasKeys = f.options.keys || f.index.options.keys + return q +} + +// Between creates a between query. +func (f *Field) Between(a IntOrFloat, b IntOrFloat) *PQLRowQuery { + text := fmt.Sprintf("Row(%s >< [%s,%s])", f.name, intOrFloatVal{a}, intOrFloatVal{b}) + q := NewPQLRowQuery(text, f.index, nil) + q.hasKeys = f.options.keys || f.index.options.keys + return q +} + +// Sum creates a sum query. +func (f *Field) Sum(row *PQLRowQuery) *PQLBaseQuery { + return f.valQuery("Sum", row) +} + +// Min creates a min query. +func (f *Field) Min(row *PQLRowQuery) *PQLBaseQuery { + return f.valQuery("Min", row) +} + +// Max creates a max query. +func (f *Field) Max(row *PQLRowQuery) *PQLBaseQuery { + return f.valQuery("Max", row) +} + +// MinRow creates a min row query. +func (f *Field) MinRow() *PQLBaseQuery { + q := fmt.Sprintf("MinRow(field='%s')", f.name) + return NewPQLBaseQuery(q, f.index, nil) +} + +// MaxRow creates a max row query. +func (f *Field) MaxRow() *PQLBaseQuery { + q := fmt.Sprintf("MaxRow(field='%s')", f.name) + return NewPQLBaseQuery(q, f.index, nil) +} + +// SetIntValue creates a Set query. +func (f *Field) SetIntValue(colIDOrKey interface{}, value int) *PQLBaseQuery { + colStr, err := formatIDKey(colIDOrKey) + if err != nil { + return NewPQLBaseQuery("", f.index, err) + } + q := fmt.Sprintf("Set(%s, %s=%d)", colStr, f.name, value) + return NewPQLBaseQuery(q, f.index, nil) +} + +// PQLRowsQuery is the return type for Rows calls. +type PQLRowsQuery struct { + index *Index + pql string + err error +} + +// NewPQLRowsQuery creates a new PQLRowsQuery. +func NewPQLRowsQuery(pql string, index *Index, err error) *PQLRowsQuery { + return &PQLRowsQuery{ + index: index, + pql: pql, + err: err, + } +} + +// Index returns the index for this query/ +func (q *PQLRowsQuery) Index() *Index { + return q.index +} + +func (q *PQLRowsQuery) Serialize() SerializedQuery { + return q.serialize() +} + +func (q *PQLRowsQuery) serialize() SerializedQuery { + return newSerializedQuery(q.pql, false) +} + +// Error returns the error or nil for this query. +func (q PQLRowsQuery) Error() error { + return q.err +} + +// Union returns the union of all matched rows. +func (q *PQLRowsQuery) Union() *PQLRowQuery { + return NewPQLRowQuery(fmt.Sprintf("UnionRows(%s)", q.serialize().String()), q.index, nil) +} + +// Rows creates a Rows query with defaults +func (f *Field) Rows() *PQLRowsQuery { + text := fmt.Sprintf("Rows(field='%s')", f.name) + return NewPQLRowsQuery(text, f.index, nil) +} + +// Like creates a Rows query filtered by a pattern. +// An underscore ('_') can be used as a placeholder for a single UTF-8 codepoint or a percent sign ('%') can be used as a placeholder for 0 or more codepoints. +// All other codepoints in the pattern are matched exactly. +func (f *Field) Like(pattern string) *PQLRowsQuery { + pattern = strings.ReplaceAll(pattern, `\`, `\\`) + pattern = strings.ReplaceAll(pattern, `'`, `\'`) + text := fmt.Sprintf("Rows(field='%s',like='%s')", f.name, pattern) + return NewPQLRowsQuery(text, f.index, nil) +} + +// RowsPrevious creates a Rows query with the given previous row ID/key +func (f *Field) RowsPrevious(rowIDOrKey interface{}) *PQLRowsQuery { + idKey, err := formatIDKey(rowIDOrKey) + if err != nil { + return NewPQLRowsQuery("", f.index, err) + } + text := fmt.Sprintf("Rows(field='%s',previous=%s)", f.name, idKey) + return NewPQLRowsQuery(text, f.index, nil) +} + +// RowsLimit creates a Rows query with the given limit +func (f *Field) RowsLimit(limit int64) *PQLRowsQuery { + if limit < 0 { + return NewPQLRowsQuery("", f.index, errors.New("rows limit must be non-negative")) + } + text := fmt.Sprintf("Rows(field='%s',limit=%d)", f.name, limit) + return NewPQLRowsQuery(text, f.index, nil) +} + +// RowsColumn creates a Rows query with the given column ID/key +func (f *Field) RowsColumn(columnIDOrKey interface{}) *PQLRowsQuery { + idKey, err := formatIDKey(columnIDOrKey) + if err != nil { + return NewPQLRowsQuery("", f.index, err) + } + text := fmt.Sprintf("Rows(field='%s',column=%s)", f.name, idKey) + return NewPQLRowsQuery(text, f.index, nil) +} + +// RowsPreviousLimit creates a Rows query with the given previous row ID/key and limit +func (f *Field) RowsPreviousLimit(rowIDOrKey interface{}, limit int64) *PQLRowsQuery { + idKey, err := formatIDKey(rowIDOrKey) + if err != nil { + return NewPQLRowsQuery("", f.index, err) + } + if limit < 0 { + return NewPQLRowsQuery("", f.index, errors.New("rows limit must be non-negative")) + } + text := fmt.Sprintf("Rows(field='%s',previous=%s,limit=%d)", f.name, idKey, limit) + return NewPQLRowsQuery(text, f.index, nil) +} + +// RowsPreviousColumn creates a Rows query with the given previous row ID/key and column ID/key +func (f *Field) RowsPreviousColumn(rowIDOrKey interface{}, columnIDOrKey interface{}) *PQLRowsQuery { + rowIDKey, err := formatIDKey(rowIDOrKey) + if err != nil { + return NewPQLRowsQuery("", f.index, err) + } + columnIDKey, err := formatIDKey(columnIDOrKey) + if err != nil { + return NewPQLRowsQuery("", f.index, err) + } + text := fmt.Sprintf("Rows(field='%s',previous=%s,column=%s)", f.name, rowIDKey, columnIDKey) + return NewPQLRowsQuery(text, f.index, nil) +} + +// RowsLimitColumn creates a Row query with the given limit and column ID/key +func (f *Field) RowsLimitColumn(limit int64, columnIDOrKey interface{}) *PQLRowsQuery { + if limit < 0 { + return NewPQLRowsQuery("", f.index, errors.New("rows limit must be non-negative")) + } + columnIDKey, err := formatIDKey(columnIDOrKey) + if err != nil { + return NewPQLRowsQuery("", f.index, err) + } + text := fmt.Sprintf("Rows(field='%s',limit=%d,column=%s)", f.name, limit, columnIDKey) + return NewPQLRowsQuery(text, f.index, nil) +} + +// RowsPreviousLimitColumn creates a Row query with the given previous row ID/key, limit and column ID/key +func (f *Field) RowsPreviousLimitColumn(rowIDOrKey interface{}, limit int64, columnIDOrKey interface{}) *PQLRowsQuery { + rowIDKey, err := formatIDKey(rowIDOrKey) + if err != nil { + return NewPQLRowsQuery("", f.index, err) + } + if limit < 0 { + return NewPQLRowsQuery("", f.index, errors.New("rows limit must be non-negative")) + } + columnIDKey, err := formatIDKey(columnIDOrKey) + if err != nil { + return NewPQLRowsQuery("", f.index, err) + } + text := fmt.Sprintf("Rows(field='%s',previous=%s,limit=%d,column=%s)", f.name, rowIDKey, limit, columnIDKey) + return NewPQLRowsQuery(text, f.index, nil) +} + +// Distinct creates a Distinct query. +func (f *Field) Distinct() *PQLRowQuery { + text := fmt.Sprintf("Distinct(Row(%s!=null),index='%s',field='%s')", f.name, f.index.Name(), f.name) + return NewPQLRowQuery(text, f.index, nil) +} + +// RowDistinct creates a Distinct query with the given row filter. +func (f *Field) RowDistinct(row *PQLRowQuery) *PQLRowQuery { + text := fmt.Sprintf("Distinct(%s,index='%s',field='%s')", row.serialize(), f.index.Name(), f.name) + return NewPQLRowQuery(text, f.index, nil) +} + +func (f *Field) binaryOperation(op string, n IntOrFloat) *PQLRowQuery { + text := fmt.Sprintf("Row(%s %s %s)", f.name, op, intOrFloatVal{n}) + q := NewPQLRowQuery(text, f.index, nil) + q.hasKeys = f.options.keys || f.index.options.keys + return q +} + +func (f *Field) valQuery(op string, row *PQLRowQuery) *PQLBaseQuery { + rowStr := "" + hasKeys := f.options.keys || f.index.options.keys + if row != nil { + serializedRow := row.serialize() + hasKeys = hasKeys || serializedRow.HasWriteKeys() + rowStr = fmt.Sprintf("%s,", serializedRow.String()) + } + text := fmt.Sprintf("%s(%sfield='%s')", op, rowStr, f.name) + q := NewPQLBaseQuery(text, f.index, nil) + q.hasKeys = hasKeys + return q +} + +func encodeMap(m map[string]interface{}) string { + result, err := json.Marshal(m) + if err != nil { + panic(err) + } + return string(result) +} diff --git a/client/orm_test.go b/client/orm_test.go new file mode 100644 index 000000000..61e2c71d7 --- /dev/null +++ b/client/orm_test.go @@ -0,0 +1,1257 @@ +// Copyright 2017 Pilosa Corp. +// +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions +// are met: +// +// 1. Redistributions of source code must retain the above copyright +// notice, this list of conditions and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright +// notice, this list of conditions and the following disclaimer in the +// documentation and/or other materials provided with the distribution. +// +// 3. Neither the name of the copyright holder nor the names of its +// contributors may be used to endorse or promote products derived +// from this software without specific prior written permission. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND +// CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, +// INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF +// MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR +// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, +// BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH +// DAMAGE. + +package client + +import ( + "fmt" + "math" + "reflect" + "sort" + "strings" + "testing" + "time" + + "github.com/pilosa/pilosa/v2/pql" + "github.com/pkg/errors" +) + +func TestORM(t *testing.T) { + var schema = NewSchema() + var sampleIndex = schema.Index("sample-index") + var sampleField = sampleIndex.Field("sample-field") + var projectIndex = schema.Index("project-index") + var collabField = projectIndex.Field("collaboration") + var b1 = sampleField.Row(10) + var b2 = sampleField.Row(20) + var b3 = sampleField.Row(42) + var b4 = collabField.Row(2) + + t.Run("SchemaDiff", func(t *testing.T) { + schema1 := NewSchema() + index11 := schema1.Index("diff-index1") + index11.Field("field1-1") + index11.Field("field1-2") + index12 := schema1.Index("diff-index2", OptIndexKeys(true), OptIndexTrackExistence(false)) + index12.Field("field2-1") + + schema2 := NewSchema() + index21 := schema2.Index("diff-index1") + index21.Field("another-field") + + targetDiff12 := NewSchema() + targetIndex1 := targetDiff12.Index("diff-index1", OptIndexTrackExistence(false)) + targetIndex1.Field("field1-1") + targetIndex1.Field("field1-2") + targetIndex2 := targetDiff12.Index("diff-index2", OptIndexKeys(true), OptIndexTrackExistence(false)) + targetIndex2.Field("field2-1") + targetIndex1.options = &IndexOptions{} + + diff12 := schema1.diff(schema2) + + strTargetDiff12 := fmt.Sprintf("%+v", targetDiff12.indexes) + strDiff12 := fmt.Sprintf("%+v", diff12.indexes) + if strDiff12 != strTargetDiff12 { + t.Fatalf("The diff must be correctly calculated, but exp/got\n%s\n%s", strTargetDiff12, strDiff12) + } + }) + + t.Run("SchemaIndexes", func(t *testing.T) { + schema1 := NewSchema() + index11 := schema1.Index("diff-index1") + index12 := schema1.Index("diff-index2") + indexes := schema1.Indexes() + target := map[string]*Index{ + "diff-index1": index11, + "diff-index2": index12, + } + if !reflect.DeepEqual(target, indexes) { + t.Fatalf("calling schema.Indexes should return indexes") + } + }) + + t.Run("SchemaToString", func(t *testing.T) { + schema1 := NewSchema() + _ = schema1.Index("test-index") + target := `map[test-index:{name: "test-index", options: "{"options":{}}", fields: map[_exists:{name: "_exists", index: "test-index", options: "{"options":{"type":"set"}}"}], shardWidth: 0}]` + if target != schema1.String() { + t.Fatalf("%s != %s", target, schema1.String()) + } + }) + + t.Run("NewIndex", func(t *testing.T) { + index1 := schema.Index("index-name") + if index1.Name() != "index-name" { + t.Fatalf("index name was not set") + } + // calling schema.Index again should return the same index + index2 := schema.Index("index-name") + if index1 != index2 { + t.Fatalf("calling schema.Index again should return the same index") + } + if !schema.HasIndex("index-name") { + t.Fatalf("HasIndex should return true") + } + if schema.HasIndex("index-x") { + t.Fatalf("HasIndex should return false") + } + }) + + t.Run("NewIndexCopy", func(t *testing.T) { + index := schema.Index("my-index-4copy", OptIndexKeys(true)) + index.Field("my-field-4copy", OptFieldTypeTime(TimeQuantumDayHour)) + copiedIndex := index.copy() + if !reflect.DeepEqual(index, copiedIndex) { + t.Fatalf("copied index should be equivalent") + } + }) + + t.Run("NewIndexOptions", func(t *testing.T) { + schema := NewSchema() + // test the defaults + index := schema.Index("index-default-options") + target := `{"options":{}}` + if target != index.options.String() { + t.Fatalf("%s != %s", target, index.options.String()) + } + + index = schema.Index("index-keys", OptIndexKeys(true)) + if true != index.Opts().Keys() { + t.Fatalf("index keys %v != %v", true, index.Opts().Keys()) + } + target = `{"options":{"keys":true}}` + if target != index.options.String() { + t.Fatalf("%s != %s", target, index.options.String()) + } + + index = schema.Index("index-trackexistence", OptIndexTrackExistence(false)) + if false != index.Opts().TrackExistence() { + t.Fatalf("index trackExistene %v != %v", true, index.Opts().TrackExistence()) + } + target = `{"options":{"trackExistence":false}}` + if target != index.options.String() { + t.Fatalf("%s != %s", target, index.options.String()) + } + }) + + t.Run("NilIndexOption", func(t *testing.T) { + schema.Index("index-with-nil-option", nil) + }) + + t.Run("IndexFields", func(t *testing.T) { + schema1 := NewSchema() + index11 := schema1.Index("diff-index1", OptIndexTrackExistence(false)) + field11 := index11.Field("field1-1") + field12 := index11.Field("field1-2") + fields := index11.Fields() + target := map[string]*Field{ + "field1-1": field11, + "field1-2": field12, + } + if !reflect.DeepEqual(target, fields) { + t.Fatalf("calling index.Fields should return fields") + } + if !index11.HasField("field1-1") { + t.Fatalf("HasField should return true") + } + if index11.HasField("field-x") { + t.Fatalf("HasField should return false") + } + }) + + t.Run("IndexToString", func(t *testing.T) { + schema1 := NewSchema() + index := schema1.Index("test-index") + target := `{name: "test-index", options: "{"options":{}}", fields: map[_exists:{name: "_exists", index: "test-index", options: "{"options":{"type":"set"}}"}], shardWidth: 0}` + if target != index.String() { + t.Fatalf("indexes not equal exp/got:\n%s\n%s", target, index.String()) + } + }) + + t.Run("Field", func(t *testing.T) { + field1 := sampleIndex.Field("nonexistent-field") + field2 := sampleIndex.Field("nonexistent-field") + if field1 != field2 { + t.Fatalf("calling index.Field again should return the same field") + } + if field1.Name() != "nonexistent-field" { + t.Fatalf("calling field.Name should return field's name") + } + }) + + t.Run("FieldCopy", func(t *testing.T) { + field := sampleIndex.Field("my-field-4copy", OptFieldTypeSet(CacheTypeRanked, 123456)) + copiedField := field.copy() + if !reflect.DeepEqual(field, copiedField) { + t.Fatalf("copied field should be equivalent") + } + }) + + t.Run("FieldToString", func(t *testing.T) { + schema1 := NewSchema() + index := schema1.Index("test-index") + field := index.Field("test-field") + target := `{name: "test-field", index: "test-index", options: "{"options":{"type":"set"}}"}` + if target != field.String() { + t.Fatalf("%s != %s", target, field.String()) + } + }) + + t.Run("NilFieldOption", func(t *testing.T) { + schema1 := NewSchema() + index := schema1.Index("test-index") + index.Field("test-field-with-nil-option", nil) + }) + + t.Run("FieldSetType", func(t *testing.T) { + schema1 := NewSchema() + index := schema1.Index("test-index") + field := index.Field("test-set-field", OptFieldTypeSet(CacheTypeLRU, 1000), OptFieldKeys(true)) + target := `{"options":{"type":"set","cacheType":"lru","cacheSize":1000,"keys":true}}` + if sortedString(target) != sortedString(field.options.String()) { + t.Fatalf("%s != %s", target, field.options.String()) + } + + field = index.Field("test-set-field2", OptFieldTypeSet(CacheTypeLRU, -10), OptFieldKeys(true)) + target = `{"options":{"type":"set","cacheType":"lru","keys":true}}` + if sortedString(target) != sortedString(field.options.String()) { + t.Fatalf("%s != %s", target, field.options.String()) + } + }) + + t.Run("Row", func(t *testing.T) { + comparePQL(t, + "Row(collaboration=5)", + collabField.Row(5)) + + comparePQL(t, + "Row(collaboration='b7feb014-8ea7-49a8-9cd8-19709161ab63')", + collabField.Row("b7feb014-8ea7-49a8-9cd8-19709161ab63")) + + q := collabField.Row(nil) + if q.err == nil { + t.Fatalf("should have failed") + } + }) + + t.Run("Set", func(t *testing.T) { + comparePQL(t, + "Set(10,collaboration=5)", + collabField.Set(5, 10)) + + comparePQL(t, + `Set('some_id',collaboration='b7feb014-8ea7-49a8-9cd8-19709161ab63')`, + collabField.Set("b7feb014-8ea7-49a8-9cd8-19709161ab63", "some_id")) + + q := collabField.Set(nil, 10) + if q.err == nil { + t.Fatalf("should have failed") + } + q = collabField.Set(5, false) + if q.err == nil { + t.Fatalf("should have failed") + } + }) + + t.Run("Timestamp", func(t *testing.T) { + timestamp := time.Date(2017, time.April, 24, 12, 14, 0, 0, time.UTC) + comparePQL(t, + "Set(20,collaboration=10,2017-04-24T12:14)", + collabField.SetTimestamp(10, 20, timestamp)) + + comparePQL(t, + "Set('mycol',collaboration='myrow',2017-04-24T12:14)", + collabField.SetTimestamp("myrow", "mycol", timestamp)) + + q := collabField.SetTimestamp(nil, 20, timestamp) + if q.err == nil { + t.Fatalf("should have failed") + } + }) + + t.Run("Clear", func(t *testing.T) { + comparePQL(t, + "Clear(10,collaboration=5)", + collabField.Clear(5, 10)) + + comparePQL(t, + "Clear('some_id',collaboration='b7feb014-8ea7-49a8-9cd8-19709161ab63')", + collabField.Clear("b7feb014-8ea7-49a8-9cd8-19709161ab63", "some_id")) + comparePQL(t, + `Clear('bill\'s',collaboration='will\'s')`, + collabField.Clear("will's", "bill's")) + + q := collabField.Clear(nil, 10) + if q.err == nil { + t.Fatalf("should have failed") + } + q = collabField.Clear(5, false) + if q.err == nil { + t.Fatalf("should have failed") + } + }) + + t.Run("ClearRow", func(t *testing.T) { + comparePQL(t, + "ClearRow(collaboration=5)", + collabField.ClearRow(5)) + + comparePQL(t, + "ClearRow(collaboration='five')", + collabField.ClearRow("five")) + + comparePQL(t, + "ClearRow(collaboration=true)", + collabField.ClearRow(true)) + + q := collabField.ClearRow(nil) + if q.err == nil { + t.Fatalf("should have failed") + } + }) + + t.Run("Union", func(t *testing.T) { + comparePQL(t, + "Union(Row(sample-field=10),Row(sample-field=20))", + sampleIndex.Union(b1, b2)) + comparePQL(t, + "Union(Row(sample-field=10),Row(sample-field=20),Row(sample-field=42))", + sampleIndex.Union(b1, b2, b3)) + comparePQL(t, + "Union(Row(sample-field=10),Row(collaboration=2))", + sampleIndex.Union(b1, b4)) + comparePQL(t, + "Union(Row(sample-field=10))", + sampleIndex.Union(b1)) + comparePQL(t, + "Union()", + sampleIndex.Union()) + }) + + t.Run("Intersect", func(t *testing.T) { + comparePQL(t, + "Intersect(Row(sample-field=10),Row(sample-field=20))", + sampleIndex.Intersect(b1, b2)) + comparePQL(t, + "Intersect(Row(sample-field=10),Row(sample-field=20),Row(sample-field=42))", + sampleIndex.Intersect(b1, b2, b3)) + comparePQL(t, + "Intersect(Row(sample-field=10),Row(collaboration=2))", + sampleIndex.Intersect(b1, b4)) + comparePQL(t, + "Intersect(Row(sample-field=10))", + sampleIndex.Intersect(b1)) + }) + + t.Run("Difference", func(t *testing.T) { + comparePQL(t, + "Difference(Row(sample-field=10),Row(sample-field=20))", + sampleIndex.Difference(b1, b2)) + comparePQL(t, + "Difference(Row(sample-field=10),Row(sample-field=20),Row(sample-field=42))", + sampleIndex.Difference(b1, b2, b3)) + comparePQL(t, + "Difference(Row(sample-field=10),Row(collaboration=2))", + sampleIndex.Difference(b1, b4)) + comparePQL(t, + "Difference(Row(sample-field=10))", + sampleIndex.Difference(b1)) + }) + + t.Run("Xor", func(t *testing.T) { + comparePQL(t, + "Xor(Row(sample-field=10),Row(sample-field=20))", + sampleIndex.Xor(b1, b2)) + comparePQL(t, + "Xor(Row(sample-field=10),Row(sample-field=20),Row(sample-field=42))", + sampleIndex.Xor(b1, b2, b3)) + comparePQL(t, + "Xor(Row(sample-field=10),Row(collaboration=2))", + sampleIndex.Xor(b1, b4)) + }) + + t.Run("Not", func(t *testing.T) { + comparePQL(t, + "Not(Row(sample-field=10))", + sampleIndex.Not(sampleField.Row(10))) + }) + + t.Run("TopN", func(t *testing.T) { + comparePQL(t, + "TopN(collaboration,n=27)", + collabField.TopN(27)) + comparePQL(t, + "TopN(collaboration,Row(collaboration=3),n=10)", + collabField.RowTopN(10, collabField.Row(3))) + comparePQL(t, + "TopN(sample-field,Row(collaboration=7),n=12,attrName='category',attrValues=[80,81])", + sampleField.FilterAttrTopN(12, collabField.Row(7), "category", 80, 81)) + comparePQL(t, + "TopN(sample-field,n=12,attrName='category',attrValues=[80,81])", + sampleField.FilterAttrTopN(12, nil, "category", 80, 81)) + }) + + t.Run("FieldLT", func(t *testing.T) { + comparePQL(t, + "Row(collaboration < 10)", + collabField.LT(10)) + comparePQL(t, + "Row(collaboration < 10.12300000)", + collabField.LT(10.123)) + }) + + t.Run("FieldLTE", func(t *testing.T) { + comparePQL(t, + "Row(collaboration <= 10)", + collabField.LTE(10)) + comparePQL(t, + "Row(collaboration <= 10.12300000)", + collabField.LTE(10.123)) + }) + + t.Run("FieldGT", func(t *testing.T) { + comparePQL(t, + "Row(collaboration > 10)", + collabField.GT(10)) + comparePQL(t, + "Row(collaboration > 10.12300000)", + collabField.GT(10.123)) + }) + + t.Run("FieldGTE", func(t *testing.T) { + comparePQL(t, + "Row(collaboration >= 10)", + collabField.GTE(10)) + comparePQL(t, + "Row(collaboration >= 10.12300000)", + collabField.GTE(10.123)) + }) + + t.Run("FieldEQ", func(t *testing.T) { + comparePQL(t, + "Row(collaboration == 10)", + collabField.Equals(10)) + comparePQL(t, + "Row(collaboration == 10.12300000)", + collabField.Equals(10.123)) + }) + + t.Run("FieldNEQ", func(t *testing.T) { + comparePQL(t, + "Row(collaboration != 10)", + collabField.NotEquals(10)) + comparePQL(t, + "Row(collaboration != 10.12300000)", + collabField.NotEquals(10.123)) + }) + + t.Run("FieldNotNull", func(t *testing.T) { + comparePQL(t, + "Row(collaboration != null)", + collabField.NotNull()) + }) + + t.Run("FieldBetween", func(t *testing.T) { + comparePQL(t, + "Row(collaboration >< [10,20])", + collabField.Between(10, 20)) + comparePQL(t, + "Row(collaboration >< [10.12300000,20.45600000])", + collabField.Between(10.123, 20.456)) + }) + + t.Run("FieldSum", func(t *testing.T) { + comparePQL(t, + "Sum(Row(collaboration=10),field='collaboration')", + collabField.Sum(collabField.Row(10))) + comparePQL(t, + "Sum(field='collaboration')", + collabField.Sum(nil)) + }) + + t.Run("FieldMinRow", func(t *testing.T) { + comparePQL(t, + "MinRow(field='sample-field')", + sampleField.MinRow()) + }) + + t.Run("FieldMaxRow", func(t *testing.T) { + comparePQL(t, + "MaxRow(field='sample-field')", + sampleField.MaxRow()) + }) + + t.Run("FieldSetValue", func(t *testing.T) { + comparePQL(t, + "Set(50, collaboration=15)", + collabField.SetIntValue(50, 15)) + + comparePQL(t, + "Set('mycol', sample-field=22)", + sampleField.SetIntValue("mycol", 22)) + + q := sampleField.SetIntValue(false, 22) + if q.err == nil { + t.Fatalf("should have failed") + } + }) + + t.Run("FilterFieldTopNInvalidField", func(t *testing.T) { + q := sampleField.FilterAttrTopN(12, collabField.Row(7), "$invalid$", 80, 81) + if q.Error() == nil { + t.Fatalf("should have failed") + } + }) + + t.Run("FilterFieldTopNInvalidValue", func(t *testing.T) { + q := sampleField.FilterAttrTopN(12, collabField.Row(7), "category", 80, func() {}) + if q.Error() == nil { + t.Fatalf("should have failed") + } + }) + + t.Run("RowOperationInvalidArg", func(t *testing.T) { + invalid := sampleField.FilterAttrTopN(12, collabField.Row(7), "$invalid$", 80, 81) + // invalid argument in pos 1 + q := sampleIndex.Union(invalid, b1) + if q.Error() == nil { + t.Fatalf("should have failed") + } + // invalid argument in pos 2 + q = sampleIndex.Intersect(b1, invalid) + if q.Error() == nil { + t.Fatalf("should have failed") + } + // invalid argument in pos 3 + q = sampleIndex.Intersect(b1, b2, invalid) + if q.Error() == nil { + t.Fatalf("should have failed") + } + // not enough rows supplied + q = sampleIndex.Difference() + if q.Error() == nil { + t.Fatalf("should have failed") + } + // not enough rows supplied + q = sampleIndex.Intersect() + if q.Error() == nil { + t.Fatalf("should have failed") + } + + // not enough rows supplied + q = sampleIndex.Xor(b1) + if q.Error() == nil { + t.Fatalf("should have failed") + } + }) + + t.Run("SetColumnAttrs", func(t *testing.T) { + attrs := map[string]interface{}{ + "quote": "\"Don't worry, be happy\"", + "happy": true, + } + comparePQL(t, + "SetColumnAttrs(5,happy=true,quote=\"\\\"Don't worry, be happy\\\"\")", + projectIndex.SetColumnAttrs(5, attrs)) + + q := projectIndex.SetColumnAttrs(false, attrs) + if q.err == nil { + t.Fatalf("should have failed") + } + }) + + t.Run("SetColumnAttrsInvalidAttr", func(t *testing.T) { + attrs := map[string]interface{}{ + "color": "blue", + "$invalid$": true, + } + if projectIndex.SetColumnAttrs(5, attrs).Error() == nil { + t.Fatalf("Should have failed") + } + }) + + t.Run("SetRowAttrs", func(t *testing.T) { + attrs := map[string]interface{}{ + "quote": "\"Don't worry, be happy\"", + "active": true, + } + comparePQL(t, + `SetRowAttrs(collaboration,5,active=true,quote="\"Don't worry, be happy\"")`, + collabField.SetRowAttrs(5, attrs)) + + comparePQL(t, + "SetRowAttrs(collaboration,'foo',active=true,quote=\"\\\"Don't worry, be happy\\\"\")", + collabField.SetRowAttrs("foo", attrs)) + + q := collabField.SetRowAttrs(nil, attrs) + if q.err == nil { + t.Fatalf("should have failed") + } + }) + + t.Run("SetRowAttrsInvalidAttr", func(t *testing.T) { + attrs := map[string]interface{}{ + "color": "blue", + "$invalid$": true, + } + if collabField.SetRowAttrs(5, attrs).Error() == nil { + t.Fatalf("Should have failed") + } + + if collabField.SetRowAttrs("foo", attrs).Error() == nil { + t.Fatalf("Should have failed") + } + }) + + t.Run("Store", func(t *testing.T) { + comparePQL(t, + "Store(Row(collaboration=5),sample-field=10)", + sampleField.Store(collabField.Row(5), 10)) + q := sampleField.Store(collabField.Row(5), nil) + if q.Error() == nil { + t.Fatalf("query error should be not nil") + } + }) + + t.Run("Options", func(t *testing.T) { + comparePQL(t, + "Options(Row(collaboration=5),columnAttrs=true,excludeColumns=true,excludeRowAttrs=true,shards=[1,3])", + sampleIndex.Options(collabField.Row(5), + OptOptionsColumnAttrs(true), + OptOptionsExcludeColumns(true), + OptOptionsExcludeRowAttrs(true), + OptOptionsShards(1, 3), + )) + comparePQL(t, + "Options(Row(collaboration=5),columnAttrs=true,excludeColumns=false,excludeRowAttrs=false)", + sampleIndex.Options(collabField.Row(5), + OptOptionsColumnAttrs(true), + )) + }) + + t.Run("BatchQuery", func(t *testing.T) { + q := sampleIndex.BatchQuery() + if q.Index() != sampleIndex { + t.Fatalf("The correct index should be assigned") + } + q.Add(sampleField.Row(44)) + q.Add(sampleField.Row(10101)) + if q.Error() != nil { + t.Fatalf("Error should be nil") + } + comparePQL(t, "Row(sample-field=44)Row(sample-field=10101)", q) + + q2 := sampleField.Row(nil) + if q2.err == nil { + t.Fatalf("should have failed") + } + }) + + t.Run("BatchQueryWithError", func(t *testing.T) { + q := sampleIndex.BatchQuery() + q.Add(sampleField.FilterAttrTopN(12, collabField.Row(7), "$invalid$", 80, 81)) + if q.Error() == nil { + t.Fatalf("The error must be set") + } + }) + + t.Run("Count", func(t *testing.T) { + q := projectIndex.Count(collabField.Row(42)) + comparePQL(t, "Count(Row(collaboration=42))", q) + }) + + t.Run("Range", func(t *testing.T) { + start := time.Date(1970, time.January, 1, 0, 0, 0, 0, time.UTC) + end := time.Date(2000, time.February, 2, 3, 4, 0, 0, time.UTC) + comparePQL(t, + "Range(collaboration=10,1970-01-01T00:00,2000-02-02T03:04)", + collabField.Range(10, start, end)) + + comparePQL(t, + "Range(collaboration='foo',1970-01-01T00:00,2000-02-02T03:04)", + collabField.Range("foo", start, end)) + + q := collabField.Range(nil, start, end) + if q.err == nil { + t.Fatalf("should have failed") + } + }) + + t.Run("RowRange", func(t *testing.T) { + start := time.Date(1970, time.January, 1, 0, 0, 0, 0, time.UTC) + end := time.Date(2000, time.February, 2, 3, 4, 0, 0, time.UTC) + comparePQL(t, + "Row(collaboration=10,from='1970-01-01T00:00',to='2000-02-02T03:04')", + collabField.RowRange(10, start, end)) + + comparePQL(t, + "Row(collaboration='foo',from='1970-01-01T00:00',to='2000-02-02T03:04')", + collabField.RowRange("foo", start, end)) + comparePQL(t, + `Row(collaboration='bill\'s',from='1970-01-01T00:00',to='2000-02-02T03:04')`, + collabField.RowRange("bill's", start, end)) + + q := collabField.RowRange(nil, start, end) + if q.err == nil { + t.Fatalf("should have failed") + } + }) + + t.Run("Rows", func(t *testing.T) { + comparePQL(t, + "Rows(field='collaboration')", + collabField.Rows()) + }) + + t.Run("UnionRows", func(t *testing.T) { + comparePQL(t, + "UnionRows(Rows(field='collaboration'))", + collabField.Rows().Union()) + }) + + t.Run("Like", func(t *testing.T) { + comparePQL(t, + "Rows(field='collaboration',like='_')", + collabField.Like("_")) + comparePQL(t, + `Rows(field='collaboration',like='_\\')`, + collabField.Like(`_\`)) + comparePQL(t, + `Rows(field='collaboration',like='_\'')`, + collabField.Like(`_'`)) + }) + + t.Run("RowPrevious", func(t *testing.T) { + comparePQL(t, + "Rows(field='collaboration',previous=42)", + collabField.RowsPrevious(42)) + comparePQL(t, + "Rows(field='collaboration',previous='forty-two')", + collabField.RowsPrevious("forty-two")) + comparePQL(t, + `Rows(field='collaboration',previous='bill\'s')`, + collabField.RowsPrevious("bill's")) + q := collabField.RowsPrevious(1.2) + if q.Error() == nil { + t.Fatalf("should have failed") + } + }) + + t.Run("RowLimit", func(t *testing.T) { + comparePQL(t, + "Rows(field='collaboration',limit=10)", + collabField.RowsLimit(10)) + q := collabField.RowsLimit(-1) + if q.Error() == nil { + t.Fatalf("should have failed") + } + }) + + t.Run("RowsColumn", func(t *testing.T) { + comparePQL(t, + "Rows(field='collaboration',column=1000)", + collabField.RowsColumn(1000)) + comparePQL(t, + "Rows(field='collaboration',column='one-thousand')", + collabField.RowsColumn("one-thousand")) + q := collabField.RowsColumn(1.2) + if q.Error() == nil { + t.Fatalf("should have failed") + } + }) + + t.Run("RowsPreviousLimit", func(t *testing.T) { + comparePQL(t, + "Rows(field='collaboration',previous=42,limit=10)", + collabField.RowsPreviousLimit(42, 10)) + comparePQL(t, + "Rows(field='collaboration',previous='forty-two',limit=10)", + collabField.RowsPreviousLimit("forty-two", 10)) + q := collabField.RowsPreviousLimit(1.2, 10) + if q.Error() == nil { + t.Fatalf("should have failed") + } + q = collabField.RowsPreviousLimit("forty-two", -1) + if q.Error() == nil { + t.Fatalf("should have failed") + } + }) + + t.Run("RowsPreviousColumn", func(t *testing.T) { + comparePQL(t, + "Rows(field='collaboration',previous=42,column=1000)", + collabField.RowsPreviousColumn(42, 1000)) + comparePQL(t, + "Rows(field='collaboration',previous='forty-two',column='one-thousand')", + collabField.RowsPreviousColumn("forty-two", "one-thousand")) + q := collabField.RowsPreviousColumn(1.2, 1000) + if q.Error() == nil { + t.Fatalf("should have failed") + } + q = collabField.RowsPreviousColumn("forty-two", 1.2) + if q.Error() == nil { + t.Fatalf("should have failed") + } + }) + + t.Run("All", func(t *testing.T) { + comparePQL(t, + "All()", + projectIndex.All()) + }) + + t.Run("Distinct", func(t *testing.T) { + comparePQL(t, + "Distinct(Row(collaboration!=null),index='project-index',field='collaboration')", + collabField.Distinct()) + }) + + t.Run("RowDistinct", func(t *testing.T) { + comparePQL(t, + "Distinct(Row(sample-field=44),index='project-index',field='collaboration')", + collabField.RowDistinct(sampleField.Row(44))) + }) + + t.Run("RowLimitColumn", func(t *testing.T) { + comparePQL(t, + "Rows(field='collaboration',limit=10,column=1000)", + collabField.RowsLimitColumn(10, 1000)) + comparePQL(t, + "Rows(field='collaboration',limit=10,column='one-thousand')", + collabField.RowsLimitColumn(10, "one-thousand")) + q := collabField.RowsLimitColumn(10, 1.2) + if q.Error() == nil { + t.Fatalf("should have failed") + } + q = collabField.RowsLimitColumn(-1, 1000) + if q.Error() == nil { + t.Fatalf("should have failed") + } + }) + + t.Run("RowsPreviousLimitColumn", func(t *testing.T) { + comparePQL(t, + "Rows(field='collaboration',previous=42,limit=10,column=1000)", + collabField.RowsPreviousLimitColumn(42, 10, 1000)) + comparePQL(t, + "Rows(field='collaboration',previous='forty-two',limit=10,column='one-thousand')", + collabField.RowsPreviousLimitColumn("forty-two", 10, "one-thousand")) + q := collabField.RowsPreviousLimitColumn(1.2, 10, 1000) + if q.Error() == nil { + t.Fatalf("should have failed") + } + q = collabField.RowsPreviousLimitColumn(42, -1, 1000) + if q.Error() == nil { + t.Fatalf("should have failed") + } + q = collabField.RowsPreviousLimitColumn(42, 10, 1.2) + if q.Error() == nil { + t.Fatalf("should have failed") + } + }) + + t.Run("GroupBy", func(t *testing.T) { + field := sampleIndex.Field("test") + comparePQL(t, + "GroupBy(Rows(field='collaboration'))", + sampleIndex.GroupBy(collabField.Rows())) + comparePQL(t, + "GroupBy(Rows(field='collaboration'),Rows(field='test'))", + sampleIndex.GroupBy(collabField.Rows(), field.Rows())) + q := sampleIndex.GroupBy() + if q.Error() == nil { + t.Fatalf("should have failed") + } + }) + + t.Run("GroupByLimit", func(t *testing.T) { + field := sampleIndex.Field("test") + comparePQL(t, + "GroupBy(Rows(field='collaboration'),limit=10)", + sampleIndex.GroupByLimit(10, collabField.Rows())) + comparePQL(t, + "GroupBy(Rows(field='collaboration'),Rows(field='test'),limit=10)", + sampleIndex.GroupByLimit(10, collabField.Rows(), field.Rows())) + q := sampleIndex.GroupByLimit(10) + if q.Error() == nil { + t.Fatalf("should have failed") + } + q = sampleIndex.GroupByLimit(-1, collabField.Rows()) + if q.Error() == nil { + t.Fatalf("should have failed") + } + }) + + t.Run("GroupByFilter", func(t *testing.T) { + field := sampleIndex.Field("test") + comparePQL(t, + "GroupBy(Rows(field='collaboration'),filter=Row(test=5))", + sampleIndex.GroupByFilter(field.Row(5), collabField.Rows())) + comparePQL(t, + "GroupBy(Rows(field='collaboration'),Rows(field='test'),filter=Row(test=5))", + sampleIndex.GroupByFilter(field.Row(5), collabField.Rows(), field.Rows())) + q := sampleIndex.GroupByFilter(field.Row(5)) + if q.Error() == nil { + t.Fatalf("should have failed") + } + }) + + t.Run("GroupByLimitFilter", func(t *testing.T) { + field := sampleIndex.Field("test") + comparePQL(t, + "GroupBy(Rows(field='collaboration'),limit=10,filter=Row(test=5))", + sampleIndex.GroupByLimitFilter(10, field.Row(5), collabField.Rows())) + comparePQL(t, + "GroupBy(Rows(field='collaboration'),Rows(field='test'),limit=10,filter=Row(test=5))", + sampleIndex.GroupByLimitFilter(10, field.Row(5), collabField.Rows(), field.Rows())) + q := sampleIndex.GroupByLimitFilter(10, field.Row(5)) + if q.Error() == nil { + t.Fatalf("should have failed") + } + q = sampleIndex.GroupByLimitFilter(-1, field.Row(5), collabField.Rows()) + if q.Error() == nil { + t.Fatalf("should have failed") + } + }) + + t.Run("GroupByBase", func(t *testing.T) { + field := sampleIndex.Field("test") + comparePQL(t, + "GroupBy(Rows(field='collaboration'))", + sampleIndex.GroupByBase( + OptGroupByBuilderRows(collabField.Rows()), + ), + ) + comparePQL(t, + "GroupBy(Rows(field='collaboration'),Rows(field='test'))", + sampleIndex.GroupByBase( + OptGroupByBuilderRows(collabField.Rows(), field.Rows()), + ), + ) + + comparePQL(t, + "GroupBy(Rows(field='collaboration'),limit=10)", + sampleIndex.GroupByBase( + OptGroupByBuilderLimit(10), + OptGroupByBuilderRows(collabField.Rows()), + ), + ) + comparePQL(t, + "GroupBy(Rows(field='collaboration'),Rows(field='test'),limit=10)", + sampleIndex.GroupByBase( + OptGroupByBuilderLimit(10), + OptGroupByBuilderRows(collabField.Rows(), field.Rows()), + ), + ) + + comparePQL(t, + "GroupBy(Rows(field='collaboration'),filter=Row(test=5))", + sampleIndex.GroupByBase( + OptGroupByBuilderFilter(field.Row(5)), + OptGroupByBuilderRows(collabField.Rows()), + ), + ) + comparePQL(t, + "GroupBy(Rows(field='collaboration'),Rows(field='test'),filter=Row(test=5))", + sampleIndex.GroupByBase( + OptGroupByBuilderFilter(field.Row(5)), + OptGroupByBuilderRows(collabField.Rows(), field.Rows()), + ), + ) + + comparePQL(t, + "GroupBy(Rows(field='collaboration'),limit=10,filter=Row(test=5))", + sampleIndex.GroupByBase( + OptGroupByBuilderLimit(10), + OptGroupByBuilderFilter(field.Row(5)), + OptGroupByBuilderRows(collabField.Rows()), + ), + ) + comparePQL(t, + "GroupBy(Rows(field='collaboration'),Rows(field='test'),limit=10,filter=Row(test=5))", + sampleIndex.GroupByBase( + OptGroupByBuilderLimit(10), + OptGroupByBuilderFilter(field.Row(5)), + OptGroupByBuilderRows(collabField.Rows(), field.Rows()), + ), + ) + + field2 := sampleIndex.Field("age") + comparePQL(t, + "GroupBy(Rows(field='collaboration'),Rows(field='test'),aggregate=Sum(Row(age=20),field='age'))", + sampleIndex.GroupByBase( + OptGroupByBuilderRows(collabField.Rows(), field.Rows()), + OptGroupByBuilderAggregate(field2.Sum(field2.Row(20))), + ), + ) + }) + + t.Run("FieldOptions", func(t *testing.T) { + field := sampleIndex.Field("foo", OptFieldKeys(true)) + if true != field.Opts().Keys() { + t.Fatalf("field keys: %v != %v", true, field.Opts().Keys()) + } + }) + + t.Run("SetFieldOptions", func(t *testing.T) { + field := sampleIndex.Field("set-field", OptFieldTypeSet(CacheTypeRanked, 9999)) + jsonString := field.options.String() + targetString := `{"options":{"type":"set","cacheType":"ranked","cacheSize":9999}}` + if sortedString(targetString) != sortedString(jsonString) { + t.Fatalf("`%s` != `%s`", targetString, jsonString) + } + compareFieldOptions(t, + field.Options(), + FieldTypeSet, + TimeQuantumNone, + CacheTypeRanked, + 9999, + pql.NewDecimal(0, 0), + pql.NewDecimal(0, 0), + "") + }) + + t.Run("IntFieldOptions", func(t *testing.T) { + field := sampleIndex.Field("int-field", OptFieldTypeInt(-10, 100)) + jsonString := field.options.String() + targetString := `{"options":{"type":"int","min":-10,"max":100}}` + if sortedString(targetString) != sortedString(jsonString) { + t.Fatalf("`%s` != `%s`", targetString, jsonString) + } + compareFieldOptions(t, + field.Options(), + FieldTypeInt, + TimeQuantumNone, + CacheTypeDefault, + 0, + pql.NewDecimal(-10, 0), + pql.NewDecimal(100, 0), + "") + + field = sampleIndex.Field("int-field2", OptFieldTypeInt(-10)) + jsonString = field.options.String() + targetString = fmt.Sprintf(`{"options":{"type":"int","min":-10,"max":%d}}`, math.MaxInt64) + if sortedString(targetString) != sortedString(jsonString) { + t.Fatalf("`%s` != `%s`", targetString, jsonString) + } + + compareFieldOptions(t, + field.Options(), + FieldTypeInt, + TimeQuantumNone, + CacheTypeDefault, + 0, + pql.NewDecimal(-10, 0), + pql.NewDecimal(math.MaxInt64, 0), + "") + field = sampleIndex.Field("int-field3", OptFieldTypeInt()) + jsonString = field.options.String() + targetString = fmt.Sprintf(`{"options":{"type":"int","min":%d,"max":%d}}`, math.MinInt64, math.MaxInt64) + if sortedString(targetString) != sortedString(jsonString) { + t.Fatalf("`%s` != `%s`", targetString, jsonString) + } + compareFieldOptions(t, + field.Options(), + FieldTypeInt, + TimeQuantumNone, + CacheTypeDefault, + 0, + pql.NewDecimal(math.MinInt64, 0), + pql.NewDecimal(math.MaxInt64, 0), + "") + + field = sampleIndex.Field("int-field4", OptFieldTypeInt(), OptFieldForeignIndex("blerg")) + jsonString = field.options.String() + targetString = fmt.Sprintf(`{"options":{"type":"int","min":%d,"max":%d,"foreignIndex":"blerg"}}`, math.MinInt64, math.MaxInt64) + if sortedString(targetString) != sortedString(jsonString) { + t.Fatalf("`%s` != `%s`", targetString, jsonString) + } + compareFieldOptions(t, + field.Options(), + FieldTypeInt, + TimeQuantumNone, + CacheTypeDefault, + 0, + pql.NewDecimal(math.MinInt64, 0), + pql.NewDecimal(math.MaxInt64, 0), + "blerg") + }) + + t.Run("TimeFieldOptions", func(t *testing.T) { + field := sampleIndex.Field("time-field", OptFieldTypeTime(TimeQuantumDayHour, true)) + if true != field.Opts().NoStandardView() { + t.Fatalf("field noStandardView %v != %v", true, field.Opts().NoStandardView()) + } + jsonString := field.options.String() + targetString := `{"options":{"noStandardView":true,"type":"time","timeQuantum":"DH"}}` + if sortedString(targetString) != sortedString(jsonString) { + t.Fatalf("`%s` != `%s`", targetString, jsonString) + } + compareFieldOptions(t, + field.Options(), + FieldTypeTime, + TimeQuantumDayHour, + CacheTypeDefault, + 0, + pql.NewDecimal(0, 0), + pql.NewDecimal(0, 0), + "") + }) + + t.Run("MutexFieldOptions", func(t *testing.T) { + field := sampleIndex.Field("mutex-field", OptFieldTypeMutex(CacheTypeRanked, 9999)) + jsonString := field.options.String() + targetString := `{"options":{"type":"mutex","cacheType":"ranked","cacheSize":9999}}` + if sortedString(targetString) != sortedString(jsonString) { + t.Fatalf("`%s` != `%s`", targetString, jsonString) + } + compareFieldOptions(t, + field.Options(), + FieldTypeMutex, + TimeQuantumNone, + CacheTypeRanked, + 9999, + pql.NewDecimal(0, 0), + pql.NewDecimal(0, 0), + "") + }) + + t.Run("BoolFieldOptions", func(t *testing.T) { + field := sampleIndex.Field("bool-field", OptFieldTypeBool()) + jsonString := field.options.String() + targetString := `{"options":{"type":"bool"}}` + if sortedString(targetString) != sortedString(jsonString) { + t.Fatalf("`%s` != `%s`", targetString, jsonString) + } + compareFieldOptions(t, + field.Options(), + FieldTypeBool, + TimeQuantumNone, + CacheTypeDefault, + 0, + pql.NewDecimal(0, 0), + pql.NewDecimal(0, 0), + "") + }) + + t.Run("DecimalFieldOptions", func(t *testing.T) { + field := sampleIndex.Field("decimal-field", OptFieldTypeDecimal(3, pql.NewDecimal(7, 3), pql.NewDecimal(999, 3))) + jsonString := field.options.String() + targetString := `{"options":{"type":"decimal","scale":3,"max":0.999,"min":0.007}}` + if sortedString(targetString) != sortedString(jsonString) { + t.Fatalf("`%s` != `%s`", targetString, jsonString) + } + compareFieldOptions(t, + field.Options(), + FieldTypeDecimal, + TimeQuantumNone, + CacheTypeDefault, + 0, + pql.NewDecimal(7, 3), + pql.NewDecimal(999, 3), + "") + }) + + t.Run("EncodeMapPanicsOnMarshalFailure", func(t *testing.T) { + defer func() { + _ = recover() + }() + m := map[string]interface{}{ + "foo": func() {}, + } + encodeMap(m) + t.Fatal("Should have panicked") + }) + + t.Run("FormatIDKey", func(t *testing.T) { + testCase := [][]interface{}{ + {uint(42), "42", nil}, + {uint32(42), "42", nil}, + {uint64(42), "42", nil}, + {42, "42", nil}, + {int32(42), "42", nil}, + {int64(42), "42", nil}, + {"foo", `'foo'`, nil}, + {false, "", errors.New("error")}, + } + for i, item := range testCase { + s, err := formatIDKey(item[0]) + if item[2] != nil { + if err == nil { + t.Fatalf("Should have failed: %d", i) + } + continue + } + if item[1] != s { + t.Fatalf("%s != %s", item[1], s) + } + } + }) +} + +func comparePQL(t *testing.T, target string, q PQLQuery) { + t.Helper() + pql := q.Serialize().String() + if target != pql { + t.Fatalf("%s != %s", target, pql) + } +} + +func compareFieldOptions(t *testing.T, opts *FieldOptions, fieldType FieldType, timeQuantum TimeQuantum, cacheType CacheType, cacheSize int, min pql.Decimal, max pql.Decimal, foreignIndex string) { + if fieldType != opts.Type() { + t.Fatalf("%s != %s", fieldType, opts.Type()) + } + if timeQuantum != opts.TimeQuantum() { + t.Fatalf("%s != %s", timeQuantum, opts.TimeQuantum()) + } + if cacheType != opts.CacheType() { + t.Fatalf("%s != %s", cacheType, opts.CacheType()) + } + if cacheSize != opts.CacheSize() { + t.Fatalf("%d != %d", cacheSize, opts.CacheSize()) + } + if min != opts.Min() { + t.Fatalf("%d != %d", min, opts.Min()) + } + if max != opts.Max() { + t.Fatalf("%d != %d", max, opts.Max()) + } + if foreignIndex != opts.ForeignIndex() { + t.Fatalf("%s != %s", foreignIndex, opts.ForeignIndex()) + } +} + +func sortedString(s string) string { + arr := strings.Split(s, "") + sort.Strings(arr) + return strings.Join(arr, "") +} diff --git a/client/record.go b/client/record.go new file mode 100644 index 000000000..d12ec6a92 --- /dev/null +++ b/client/record.go @@ -0,0 +1,90 @@ +// Copyright 2017 Pilosa Corp. +// +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions +// are met: +// +// 1. Redistributions of source code must retain the above copyright +// notice, this list of conditions and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright +// notice, this list of conditions and the following disclaimer in the +// documentation and/or other materials provided with the distribution. +// +// 3. Neither the name of the copyright holder nor the names of its +// contributors may be used to endorse or promote products derived +// from this software without specific prior written permission. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND +// CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, +// INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF +// MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR +// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, +// BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH +// DAMAGE. + +package client + +// Record is a Column or a FieldValue. +type Record interface { + Shard(shardWidth uint64) uint64 + Less(other Record) bool +} + +// RecordIterator is an iterator for a record. +type RecordIterator interface { + NextRecord() (Record, error) +} + +// Column defines a single Pilosa column. +type Column struct { + RowID uint64 + ColumnID uint64 + RowKey string + ColumnKey string + Timestamp int64 +} + +// Shard returns the shard for this column. +func (b Column) Shard(shardWidth uint64) uint64 { + return b.ColumnID / shardWidth +} + +// Less returns true if this column sorts before the given Record. +func (b Column) Less(other Record) bool { + if ob, ok := other.(Column); ok { + if b.RowID == ob.RowID { + return b.ColumnID < ob.ColumnID + } + return b.RowID < ob.RowID + } + return false +} + +// FieldValue represents the value for a column within a +// range-encoded field. +type FieldValue struct { + ColumnID uint64 + ColumnKey string + Value int64 +} + +// Shard returns the shard for this field value. +func (v FieldValue) Shard(shardWidth uint64) uint64 { + return v.ColumnID / shardWidth +} + +// Less returns true if this field value sorts before the given Record. +func (v FieldValue) Less(other Record) bool { + if ov, ok := other.(FieldValue); ok { + return v.ColumnID < ov.ColumnID + } + return false +} diff --git a/client/record_test.go b/client/record_test.go new file mode 100644 index 000000000..8e4d72db4 --- /dev/null +++ b/client/record_test.go @@ -0,0 +1,98 @@ +// Copyright 2017 Pilosa Corp. +// +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions +// are met: +// +// 1. Redistributions of source code must retain the above copyright +// notice, this list of conditions and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright +// notice, this list of conditions and the following disclaimer in the +// documentation and/or other materials provided with the distribution. +// +// 3. Neither the name of the copyright holder nor the names of its +// contributors may be used to endorse or promote products derived +// from this software without specific prior written permission. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND +// CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, +// INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF +// MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR +// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, +// BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH +// DAMAGE. + +package client_test + +import ( + "testing" + + "github.com/pilosa/pilosa/v2/client" +) + +func TestColumnShard(t *testing.T) { + a := client.Column{RowID: 15, ColumnID: 55, Timestamp: 100101} + target := uint64(0) + if a.Shard(100) != target { + t.Fatalf("shard %d != %d", target, a.Shard(100)) + } + target = 5 + if a.Shard(10) != target { + t.Fatalf("shard %d != %d", target, a.Shard(10)) + } +} + +func TestColumnLess(t *testing.T) { + a := client.Column{RowID: 10, ColumnID: 200} + a2 := client.Column{RowID: 10, ColumnID: 1000} + b := client.Column{RowID: 200, ColumnID: 10} + c := client.FieldValue{ColumnID: 1} + if !a.Less(a2) { + t.Fatalf("%v should be less than %v", a, a2) + } + if !a.Less(b) { + t.Fatalf("%v should be less than %v", a, b) + } + if b.Less(a) { + t.Fatalf("%v should not be less than %v", b, a) + } + if c.Less(a) { + t.Fatalf("%v should not be less than %v", c, a) + } +} + +func TestFieldValueShard(t *testing.T) { + a := client.FieldValue{ColumnID: 55, Value: 125} + target := uint64(0) + if a.Shard(100) != target { + t.Fatalf("shard %d != %d", target, a.Shard(100)) + } + target = 5 + if a.Shard(10) != target { + t.Fatalf("shard %d != %d", target, a.Shard(10)) + } + +} + +func TestFieldValueLess(t *testing.T) { + a := client.FieldValue{ColumnID: 55, Value: 125} + b := client.FieldValue{ColumnID: 100, Value: 125} + c := client.Column{ColumnID: 1, RowID: 2} + if !a.Less(b) { + t.Fatalf("%v should be less than %v", a, b) + } + if b.Less(a) { + t.Fatalf("%v should not be less than %v", b, a) + } + if c.Less(a) { + t.Fatalf("%v should not be less than %v", c, a) + } +} diff --git a/client/response.go b/client/response.go new file mode 100644 index 000000000..111286d2a --- /dev/null +++ b/client/response.go @@ -0,0 +1,610 @@ +// Copyright 2017 Pilosa Corp. +// +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions +// are met: +// +// 1. Redistributions of source code must retain the above copyright +// notice, this list of conditions and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright +// notice, this list of conditions and the following disclaimer in the +// documentation and/or other materials provided with the distribution. +// +// 3. Neither the name of the copyright holder nor the names of its +// contributors may be used to endorse or promote products derived +// from this software without specific prior written permission. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND +// CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, +// INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF +// MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR +// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, +// BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH +// DAMAGE. + +package client + +import ( + "encoding/json" + "errors" + "fmt" + + "github.com/pilosa/pilosa/v2/internal" +) + +// QueryResponse types. +const ( + QueryResultTypeNil uint32 = iota + QueryResultTypeRow + QueryResultTypePairs + QueryResultTypePairsField + QueryResultTypeValCount + QueryResultTypeUint64 + QueryResultTypeBool + QueryResultTypeRowIDs // this is not used by the client + QueryResultTypeGroupCounts + QueryResultTypeRowIdentifiers + QueryResultTypePair + QueryResultTypePairField + QueryResultTypeSignedRow +) + +// QueryResponse represents the response from a Pilosa query. +type QueryResponse struct { + ResultList []QueryResult `json:"results,omitempty"` + ColumnList []ColumnItem `json:"columns,omitempty"` + ErrorMessage string `json:"error-message,omitempty"` + Success bool `json:"success,omitempty"` +} + +func newQueryResponseFromInternal(response *internal.QueryResponse) (*QueryResponse, error) { + if response.Err != "" { + return &QueryResponse{ + ErrorMessage: response.Err, + Success: false, + }, nil + } + results := make([]QueryResult, 0, len(response.Results)) + for _, r := range response.Results { + result, err := newQueryResultFromInternal(r) + if err != nil { + return nil, err + } + results = append(results, result) + } + columns := make([]ColumnItem, 0, len(response.ColumnAttrSets)) + for _, p := range response.ColumnAttrSets { + columnItem, err := newColumnItemFromInternal(p) + if err != nil { + return nil, err + } + columns = append(columns, columnItem) + } + + return &QueryResponse{ + ResultList: results, + ColumnList: columns, + Success: true, + }, nil +} + +// Results returns all results in the response. +func (qr *QueryResponse) Results() []QueryResult { + return qr.ResultList +} + +// Result returns the first result or nil. +func (qr *QueryResponse) Result() QueryResult { + if len(qr.ResultList) == 0 { + return nil + } + return qr.ResultList[0] +} + +// Columns returns all column attributes in the response. +// *DEPRECATED* +func (qr *QueryResponse) Columns() []ColumnItem { + return qr.ColumnList +} + +// Column returns the attributes for first column. +// *DEPRECATED* +func (qr *QueryResponse) Column() ColumnItem { + if len(qr.ColumnList) == 0 { + return ColumnItem{} + } + return qr.ColumnList[0] +} + +// ColumnAttrs returns all column attributes in the response. +func (qr *QueryResponse) ColumnAttrs() []ColumnItem { + return qr.ColumnList +} + +// QueryResult represents one of the results in the response. +type QueryResult interface { + Type() uint32 + Row() RowResult + CountItems() []CountResultItem + CountItem() CountResultItem + Count() int64 + Value() int64 + Changed() bool + GroupCounts() []GroupCount + RowIdentifiers() RowIdentifiersResult +} + +func newQueryResultFromInternal(result *internal.QueryResult) (QueryResult, error) { + switch result.Type { + case QueryResultTypeNil: + return NilResult{}, nil + case QueryResultTypeRow: + return newRowResultFromInternal(result.Row) + case QueryResultTypePairs: + return countItemsFromInternal(result.Pairs), nil + case QueryResultTypePairsField: + return countItemsFromInternal(result.PairsField.Pairs), nil + case QueryResultTypeValCount: + return &ValCountResult{ + Val: result.ValCount.Val, + Cnt: result.ValCount.Count, + }, nil + case QueryResultTypeUint64: + return IntResult(result.N), nil + case QueryResultTypeBool: + return BoolResult(result.Changed), nil + case QueryResultTypeRowIdentifiers: + return &RowIdentifiersResult{ + IDs: result.RowIdentifiers.Rows, + Keys: result.RowIdentifiers.Keys, + }, nil + case QueryResultTypeGroupCounts: + return groupCountsFromInternal(result.GroupCounts), nil + case QueryResultTypePair: + return CountItem{CountResultItem: countItemFromInternal(result.Pairs[0])}, nil + case QueryResultTypePairField: + return CountItem{CountResultItem: countItemFromInternal(result.PairField.Pair)}, nil + } + + return nil, ErrUnknownType +} + +// CountResultItem represents a result from TopN call. +type CountResultItem struct { + ID uint64 `json:"id"` + Key string `json:"key,omitempty"` + Count uint64 `json:"count"` +} + +func (c *CountResultItem) String() string { + if c.Key != "" { + return fmt.Sprintf("%s:%d", c.Key, c.Count) + } + return fmt.Sprintf("%d:%d", c.ID, c.Count) +} + +type CountItem struct { + CountResultItem +} + +// Type is the type of this result. +func (CountItem) Type() uint32 { return QueryResultTypePairField } + +// Row returns a RowResult. +func (CountItem) Row() RowResult { return RowResult{} } + +// CountItems returns a CountResultItem slice. +func (t CountItem) CountItems() []CountResultItem { return []CountResultItem{t.CountResultItem} } + +// CountItem returns a CountResultItem +func (t CountItem) CountItem() CountResultItem { return t.CountResultItem } + +// Count returns the result of a Count call. +func (CountItem) Count() int64 { return 0 } + +// Value returns the result of a Min, Max or Sum call. +func (CountItem) Value() int64 { return 0 } + +// Changed returns whether the corresponding Set or Clear call changed the value of a bit. +func (CountItem) Changed() bool { return false } + +// GroupCounts returns the result of a GroupBy call. +func (CountItem) GroupCounts() []GroupCount { return nil } + +// RowIdentifiers returns the result of a Rows call. +func (CountItem) RowIdentifiers() RowIdentifiersResult { return RowIdentifiersResult{} } + +func countItemFromInternal(item *internal.Pair) CountResultItem { + return CountResultItem{ID: item.ID, Key: item.Key, Count: item.Count} +} + +func countItemsFromInternal(items []*internal.Pair) TopNResult { + result := make([]CountResultItem, 0, len(items)) + for _, v := range items { + result = append(result, countItemFromInternal(v)) + } + return TopNResult(result) +} + +// TopNResult is returned from TopN call. +type TopNResult []CountResultItem + +// Type is the type of this result. +func (TopNResult) Type() uint32 { return QueryResultTypePairsField } + +// Row returns a RowResult. +func (TopNResult) Row() RowResult { return RowResult{} } + +// CountItems returns a CountResultItem slice. +func (t TopNResult) CountItems() []CountResultItem { return t } + +// CountItem returns a CountResultItem +func (t TopNResult) CountItem() CountResultItem { + if len(t) >= 1 { + return t[0] + } + return CountResultItem{} +} + +// Count returns the result of a Count call. +func (TopNResult) Count() int64 { return 0 } + +// Value returns the result of a Min, Max or Sum call. +func (TopNResult) Value() int64 { return 0 } + +// Changed returns whether the corresponding Set or Clear call changed the value of a bit. +func (TopNResult) Changed() bool { return false } + +// GroupCounts returns the result of a GroupBy call. +func (TopNResult) GroupCounts() []GroupCount { return nil } + +// RowIdentifiers returns the result of a Rows call. +func (TopNResult) RowIdentifiers() RowIdentifiersResult { return RowIdentifiersResult{} } + +// RowResult represents a result from Row, Union, Intersect, Difference and Range PQL calls. +type RowResult struct { + Attributes map[string]interface{} `json:"attrs"` + Columns []uint64 `json:"columns"` + Keys []string `json:"keys"` +} + +func newRowResultFromInternal(row *internal.Row) (*RowResult, error) { + attrs, err := convertInternalAttrsToMap(row.Attrs) + if err != nil { + return nil, err + } + result := &RowResult{ + Attributes: attrs, + Columns: row.Columns, + Keys: row.Keys, + } + return result, nil +} + +// Type is the type of this result. +func (RowResult) Type() uint32 { return QueryResultTypeRow } + +// Row returns a RowResult. +func (b RowResult) Row() RowResult { return b } + +// CountItems returns a CountResultItem slice. +func (RowResult) CountItems() []CountResultItem { return nil } + +// CountItem returns a CountResultItem +func (RowResult) CountItem() CountResultItem { return CountResultItem{} } + +// Count returns the result of a Count call. +func (RowResult) Count() int64 { return 0 } + +// Value returns the result of a Min, Max or Sum call. +func (RowResult) Value() int64 { return 0 } + +// Changed returns whether the corresponding Set or Clear call changed the value of a bit. +func (RowResult) Changed() bool { return false } + +// GroupCounts returns the result of a GroupBy call. +func (RowResult) GroupCounts() []GroupCount { return nil } + +// RowIdentifiers returns the result of a Rows call. +func (RowResult) RowIdentifiers() RowIdentifiersResult { return RowIdentifiersResult{} } + +// MarshalJSON serializes this row result. +func (b RowResult) MarshalJSON() ([]byte, error) { + columns := b.Columns + if columns == nil { + columns = []uint64{} + } + keys := b.Keys + if keys == nil { + keys = []string{} + } + return json.Marshal(struct { + Attributes map[string]interface{} `json:"attrs"` + Columns []uint64 `json:"columns"` + Keys []string `json:"keys"` + }{ + Attributes: b.Attributes, + Columns: columns, + Keys: keys, + }) +} + +// ValCountResult is returned from Min, Max and Sum calls. +type ValCountResult struct { + Val int64 `json:"val"` + Cnt int64 `json:"count"` +} + +// Type is the type of this result. +func (ValCountResult) Type() uint32 { return QueryResultTypeValCount } + +// Row returns a RowResult. +func (ValCountResult) Row() RowResult { return RowResult{} } + +// CountItems returns a CountResultItem slice. +func (ValCountResult) CountItems() []CountResultItem { return nil } + +// CountItem returns a CountResultItem +func (ValCountResult) CountItem() CountResultItem { return CountResultItem{} } + +// Count returns the result of a Count call. +func (c ValCountResult) Count() int64 { return c.Cnt } + +// Value returns the result of a Min, Max or Sum call. +func (c ValCountResult) Value() int64 { return c.Val } + +// Changed returns whether the corresponding Set or Clear call changed the value of a bit. +func (ValCountResult) Changed() bool { return false } + +// GroupCounts returns the result of a GroupBy call. +func (ValCountResult) GroupCounts() []GroupCount { return nil } + +// RowIdentifiers returns the result of a Rows call. +func (ValCountResult) RowIdentifiers() RowIdentifiersResult { return RowIdentifiersResult{} } + +// IntResult is returned from Count call. +type IntResult int64 + +// Type is the type of this result. +func (IntResult) Type() uint32 { return QueryResultTypeUint64 } + +// Row returns a RowResult. +func (IntResult) Row() RowResult { return RowResult{} } + +// CountItems returns a CountResultItem slice. +func (IntResult) CountItems() []CountResultItem { return nil } + +// CountItem returns a CountResultItem +func (IntResult) CountItem() CountResultItem { return CountResultItem{} } + +// Count returns the result of a Count call. +func (i IntResult) Count() int64 { return int64(i) } + +// Value returns the result of a Min, Max or Sum call. +func (IntResult) Value() int64 { return 0 } + +// Changed returns whether the corresponding Set or Clear call changed the value of a bit. +func (IntResult) Changed() bool { return false } + +// GroupCounts returns the result of a GroupBy call. +func (IntResult) GroupCounts() []GroupCount { return nil } + +// RowIdentifiers returns the result of a Rows call. +func (IntResult) RowIdentifiers() RowIdentifiersResult { return RowIdentifiersResult{} } + +// BoolResult is returned from Set and Clear calls. +type BoolResult bool + +// Type is the type of this result. +func (BoolResult) Type() uint32 { return QueryResultTypeBool } + +// Row returns a RowResult. +func (BoolResult) Row() RowResult { return RowResult{} } + +// CountItems returns a CountResultItem slice. +func (BoolResult) CountItems() []CountResultItem { return nil } + +// CountItem returns a CountResultItem +func (BoolResult) CountItem() CountResultItem { return CountResultItem{} } + +// Count returns the result of a Count call. +func (BoolResult) Count() int64 { return 0 } + +// Value returns the result of a Min, Max or Sum call. +func (BoolResult) Value() int64 { return 0 } + +// Changed returns whether the corresponding Set or Clear call changed the value of a bit. +func (b BoolResult) Changed() bool { return bool(b) } + +// GroupCounts returns the result of a GroupBy call. +func (BoolResult) GroupCounts() []GroupCount { return nil } + +// RowIdentifiers returns the result of a Rows call. +func (BoolResult) RowIdentifiers() RowIdentifiersResult { return RowIdentifiersResult{} } + +// NilResult is returned from calls which don't return a value, such as SetRowAttrs. +type NilResult struct{} + +// Type is the type of this result. +func (NilResult) Type() uint32 { return QueryResultTypeNil } + +// Row returns a RowResult. +func (NilResult) Row() RowResult { return RowResult{} } + +// CountItems returns a CountResultItem slice. +func (NilResult) CountItems() []CountResultItem { return nil } + +// CountItem returns a CountResultItem +func (NilResult) CountItem() CountResultItem { return CountResultItem{} } + +// Count returns the result of a Count call. +func (NilResult) Count() int64 { return 0 } + +// Value returns the result of a Min, Max or Sum call. +func (NilResult) Value() int64 { return 0 } + +// Changed returns whether the corresponding Set or Clear call changed the value of a bit. +func (NilResult) Changed() bool { return false } + +// GroupCounts returns the result of a GroupBy call. +func (NilResult) GroupCounts() []GroupCount { return nil } + +// RowIdentifiers returns the result of a Rows call. +func (NilResult) RowIdentifiers() RowIdentifiersResult { return RowIdentifiersResult{} } + +// FieldRow represents a Group in a GroupBy call result. +type FieldRow struct { + FieldName string `json:"field"` + RowID uint64 `json:"rowID"` + RowKey string `json:"rowKey"` + Value *int64 `json:"value,omitempty"` +} + +// GroupCount contains groups and their count in a GroupBy call result. +type GroupCount struct { + Groups []FieldRow `json:"groups"` + Count int64 `json:"count"` + Agg int64 `json:"agg"` +} + +// GroupCountResult is returned from GroupBy call. +type GroupCountResult []GroupCount + +// Type is the type of this result. +func (GroupCountResult) Type() uint32 { return QueryResultTypeGroupCounts } + +// Row returns a RowResult. +func (GroupCountResult) Row() RowResult { return RowResult{} } + +// CountItems returns a CountResultItem slice. +func (GroupCountResult) CountItems() []CountResultItem { return nil } + +// CountItem returns a CountResultItem +func (GroupCountResult) CountItem() CountResultItem { return CountResultItem{} } + +// Count returns the result of a Count call. +func (GroupCountResult) Count() int64 { return 0 } + +// Value returns the result of a Min, Max or Sum call. +func (GroupCountResult) Value() int64 { return 0 } + +// Changed returns whether the corresponding Set or Clear call changed the value of a bit. +func (GroupCountResult) Changed() bool { return false } + +// GroupCounts returns the result of a GroupBy call. +func (r GroupCountResult) GroupCounts() []GroupCount { return r } + +// RowIdentifiers returns the result of a Rows call. +func (GroupCountResult) RowIdentifiers() RowIdentifiersResult { return RowIdentifiersResult{} } + +// RowIdentifiersResult is returned from a Rows call. +type RowIdentifiersResult struct { + IDs []uint64 `json:"ids"` + Keys []string `json:"keys,omitempty"` +} + +// Type is the type of this result. +func (RowIdentifiersResult) Type() uint32 { return QueryResultTypeRowIdentifiers } + +// Row returns a RowResult. +func (RowIdentifiersResult) Row() RowResult { return RowResult{} } + +// CountItems returns a CountResultItem slice. +func (RowIdentifiersResult) CountItems() []CountResultItem { return nil } + +// CountItem returns a CountResultItem +func (RowIdentifiersResult) CountItem() CountResultItem { return CountResultItem{} } + +// Count returns the result of a Count call. +func (RowIdentifiersResult) Count() int64 { return 0 } + +// Value returns the result of a Min, Max or Sum call. +func (RowIdentifiersResult) Value() int64 { return 0 } + +// Changed returns whether the corresponding Set or Clear call changed the value of a bit. +func (RowIdentifiersResult) Changed() bool { return false } + +// GroupCounts returns the result of a GroupBy call. +func (RowIdentifiersResult) GroupCounts() []GroupCount { return nil } + +// RowIdentifiers returns the result of a Rows call. +func (r RowIdentifiersResult) RowIdentifiers() RowIdentifiersResult { return r } + +func groupCountsFromInternal(items *internal.GroupCounts) GroupCountResult { + result := make([]GroupCount, 0, len(items.Groups)) + for _, g := range items.Groups { + groups := make([]FieldRow, 0, len(g.Group)) + for _, f := range g.Group { + fr := FieldRow{ + FieldName: f.Field, + RowID: f.RowID, + RowKey: f.RowKey, + } + if f.Value != nil { + fr.Value = &f.Value.Value + } + groups = append(groups, fr) + } + result = append(result, GroupCount{ + Groups: groups, + Count: int64(g.Count), + Agg: int64(g.Agg), + }) + } + return GroupCountResult(result) +} + +const ( + stringType = 1 + intType = 2 + boolType = 3 + floatType = 4 +) + +func convertInternalAttrsToMap(attrs []*internal.Attr) (attrsMap map[string]interface{}, err error) { + attrsMap = make(map[string]interface{}, len(attrs)) + for _, attr := range attrs { + switch attr.Type { + case stringType: + attrsMap[attr.Key] = attr.StringValue + case intType: + attrsMap[attr.Key] = attr.IntValue + case boolType: + attrsMap[attr.Key] = attr.BoolValue + case floatType: + attrsMap[attr.Key] = attr.FloatValue + default: + return nil, errors.New("Unknown attribute type") + } + } + + return attrsMap, nil +} + +// ColumnItem represents data about a column. +// Column data is only returned if QueryOptions.Columns was set to true. +type ColumnItem struct { + ID uint64 `json:"id,omitempty"` + Key string `json:"key,omitempty"` + Attributes map[string]interface{} `json:"attributes,omitempty"` +} + +func newColumnItemFromInternal(column *internal.ColumnAttrSet) (ColumnItem, error) { + attrs, err := convertInternalAttrsToMap(column.Attrs) + if err != nil { + return ColumnItem{}, err + } + return ColumnItem{ + ID: column.ID, + Key: column.Key, + Attributes: attrs, + }, nil +} diff --git a/client/response_test.go b/client/response_test.go new file mode 100644 index 000000000..7c228305e --- /dev/null +++ b/client/response_test.go @@ -0,0 +1,363 @@ +// Copyright 2017 Pilosa Corp. +// +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions +// are met: +// +// 1. Redistributions of source code must retain the above copyright +// notice, this list of conditions and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright +// notice, this list of conditions and the following disclaimer in the +// documentation and/or other materials provided with the distribution. +// +// 3. Neither the name of the copyright holder nor the names of its +// contributors may be used to endorse or promote products derived +// from this software without specific prior written permission. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND +// CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, +// INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF +// MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR +// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, +// BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH +// DAMAGE. + +package client + +import ( + "encoding/json" + "fmt" + "log" + "reflect" + "testing" + + "github.com/pilosa/pilosa/v2/internal" +) + +func TestNewRowResultFromInternal(t *testing.T) { + targetAttrs := map[string]interface{}{ + "name": "some string", + "age": int64(95), + "registered": true, + "height": 1.83, + } + targetColumns := []uint64{5, 10} + attrs := []*internal.Attr{ + {Key: "name", StringValue: "some string", Type: 1}, + {Key: "age", IntValue: 95, Type: 2}, + {Key: "registered", BoolValue: true, Type: 3}, + {Key: "height", FloatValue: 1.83, Type: 4}, + } + row := &internal.Row{ + Attrs: attrs, + Columns: []uint64{5, 10}, + } + result, err := newRowResultFromInternal(row) + if err != nil { + t.Fatalf("Failed with error: %s", err) + } + // assertMapEquals(t, targetAttrs, result.Attributes) + if !reflect.DeepEqual(targetAttrs, result.Attributes) { + t.Fatal() + } + if !reflect.DeepEqual(targetColumns, result.Columns) { + t.Fatal() + } +} + +func TestNewQueryResponseFromInternal(t *testing.T) { + targetAttrs := map[string]interface{}{ + "name": "some string", + "age": int64(95), + "registered": true, + "height": 1.83, + } + targetColumns := []uint64{5, 10} + targetCountItems := []CountResultItem{ + {ID: 10, Count: 100}, + } + attrs := []*internal.Attr{ + {Key: "name", StringValue: "some string", Type: 1}, + {Key: "age", IntValue: 95, Type: 2}, + {Key: "registered", BoolValue: true, Type: 3}, + {Key: "height", FloatValue: 1.83, Type: 4}, + } + row := &internal.Row{ + Attrs: attrs, + Columns: []uint64{5, 10}, + } + pairs := []*internal.Pair{ + {ID: 10, Count: 100}, + } + response := &internal.QueryResponse{ + Results: []*internal.QueryResult{ + {Type: QueryResultTypeRow, Row: row}, + {Type: QueryResultTypePairs, Pairs: pairs}, + }, + Err: "", + } + qr, err := newQueryResponseFromInternal(response) + if err != nil { + t.Fatalf("Failed with error: %s", err) + } + if qr.ErrorMessage != "" { + t.Fatalf("ErrorMessage should be empty") + } + if !qr.Success { + t.Fatalf("IsSuccess should be true") + } + + results := qr.Results() + if len(results) != 2 { + t.Fatalf("Number of results should be 2") + } + if results[0] != qr.Result() { + t.Fatalf("Result() should return the first result") + } + if !reflect.DeepEqual(targetAttrs, results[0].Row().Attributes) { + t.Fatalf("The row result should contain the attributes") + } + if !reflect.DeepEqual(targetColumns, results[0].Row().Columns) { + t.Fatalf("The row result should contain the columns") + } + if !reflect.DeepEqual(targetCountItems, results[1].CountItems()) { + t.Fatalf("The response should include count items") + } +} + +func TestNewQueryResponseWithErrorFromInternal(t *testing.T) { + response := &internal.QueryResponse{ + Err: "some error", + } + qr, err := newQueryResponseFromInternal(response) + if err != nil { + t.Fatalf("Failed with error: %s", err) + } + if qr.ErrorMessage != "some error" { + t.Fatalf("The response should include the error message") + } + if qr.Success { + t.Fatalf("IsSuccess should be false") + } + if qr.Result() != nil { + t.Fatalf("If there are no results, Result should return nil") + } +} + +func TestNewQueryResponseFromInternalFailure(t *testing.T) { + attrs := []*internal.Attr{ + {Key: "name", StringValue: "some string", Type: 99}, + } + row := &internal.Row{ + Attrs: attrs, + } + response := &internal.QueryResponse{ + Results: []*internal.QueryResult{{Type: QueryResultTypeRow, Row: row}}, + } + qr, err := newQueryResponseFromInternal(response) + if qr != nil && err == nil { + t.Fatalf("Should have failed") + } + response = &internal.QueryResponse{ + ColumnAttrSets: []*internal.ColumnAttrSet{{ID: 1, Attrs: attrs}}, + } + qr, err = newQueryResponseFromInternal(response) + if qr != nil && err == nil { + t.Fatalf("Should have failed") + } +} + +func TestCountResultItemToString(t *testing.T) { + tests := []struct { + item *CountResultItem + expected string + }{ + {item: &CountResultItem{ID: 100, Count: 50}, expected: "100:50"}, + {item: &CountResultItem{Key: "blah", Count: 50}, expected: "blah:50"}, + {item: &CountResultItem{Key: "blah", ID: 22, Count: 50}, expected: "blah:50"}, + {item: &CountResultItem{Key: "blah", ID: 22}, expected: "blah:0"}, + {item: &CountResultItem{}, expected: "0:0"}, + } + + for i, tst := range tests { + t.Run(fmt.Sprintf("%d: ", i), func(t *testing.T) { + if tst.expected != tst.item.String() { + t.Fatalf("%s != %s", tst.expected, tst.item.String()) + } + }) + } +} + +func TestMarshalResults(t *testing.T) { + attrs := []*internal.Attr{ + {Key: "name", StringValue: "some string", Type: 1}, + {Key: "age", IntValue: 95, Type: 2}, + {Key: "registered", BoolValue: true, Type: 3}, + {Key: "height", FloatValue: 1.83, Type: 4}, + } + row := &internal.Row{ + Attrs: attrs, + Columns: []uint64{5, 10}, + } + pairs := []*internal.Pair{ + {ID: 10, Count: 100}, + } + pbufResults := []*internal.QueryResult{ + {Type: QueryResultTypeRow, Row: row}, + {Type: QueryResultTypePairs, Pairs: pairs}, + } + resultJSONStrings := make([]string, len(pbufResults)) + for i, pr := range pbufResults { + r, err := newQueryResultFromInternal(pr) + if err != nil { + t.Fatal(err) + } + b, err := json.Marshal(r) + if err != nil { + t.Fatal(err) + } + resultJSONStrings[i] = string(b) + } + targetJSON := []string{ + `{"attrs":{"age":95,"height":1.83,"name":"some string","registered":true},"columns":[5,10],"keys":[]}`, + `[{"id":10,"count":100}]`, + } + for i := range targetJSON { + if sortedString(targetJSON[i]) != sortedString(resultJSONStrings[i]) { + t.Fatalf("%v != %v ", targetJSON[i], resultJSONStrings[i]) + } + } + +} + +func TestUnknownQueryResultType(t *testing.T) { + result := &internal.QueryResult{ + Type: 999, + } + _, err := newQueryResultFromInternal(result) + if err != ErrUnknownType { + t.Fatalf("Should have failed with ErrUnknownType") + } +} + +func TestTopNResult(t *testing.T) { + result := TopNResult{ + CountResultItem{ID: 100, Count: 10}, + } + expectResult(t, result, QueryResultTypePairsField, RowResult{}, []CountResultItem{{100, "", 10}}, 0, 0, false, nil, RowIdentifiersResult{}) +} + +func TestRowResult(t *testing.T) { + result := RowResult{ + Columns: []uint64{1, 2, 3}, + } + targetBmp := RowResult{ + Columns: []uint64{1, 2, 3}, + } + expectResult(t, result, QueryResultTypeRow, targetBmp, nil, 0, 0, false, nil, RowIdentifiersResult{}) +} + +func TestRowResultNilColumns(t *testing.T) { + result := RowResult{ + Columns: nil, + } + _, err := result.MarshalJSON() + if err != nil { + t.Fatal(err) + } +} + +func TestSumCountResult(t *testing.T) { + result := ValCountResult{ + Val: 100, + Cnt: 50, + } + expectResult(t, result, QueryResultTypeValCount, RowResult{}, nil, 100, 50, false, nil, RowIdentifiersResult{}) +} + +func TestIntResult(t *testing.T) { + result := IntResult(11) + expectResult(t, result, QueryResultTypeUint64, RowResult{}, nil, 0, 11, false, nil, RowIdentifiersResult{}) +} + +func TestBoolResult(t *testing.T) { + result := BoolResult(true) + expectResult(t, result, QueryResultTypeBool, RowResult{}, nil, 0, 0, true, nil, RowIdentifiersResult{}) +} + +func TestNilResult(t *testing.T) { + result := NilResult{} + expectResult(t, result, QueryResultTypeNil, RowResult{}, nil, 0, 0, false, nil, RowIdentifiersResult{}) +} + +func TestGroupCountResult(t *testing.T) { + result := GroupCountResult{ + {Groups: []FieldRow{{FieldName: "f1", RowID: 1}}, Count: 2}, + {Groups: []FieldRow{{FieldName: "f1", RowID: 2}}, Count: 1}, + } + expectResult(t, result, QueryResultTypeGroupCounts, RowResult{}, nil, 0, 0, false, []GroupCount{ + {Groups: []FieldRow{{FieldName: "f1", RowID: 1}}, Count: 2}, + {Groups: []FieldRow{{FieldName: "f1", RowID: 2}}, Count: 1}, + }, RowIdentifiersResult{}) +} + +func TestGroupCountWithValueResult(t *testing.T) { + var a, b int64 = -1, 1 + + result := GroupCountResult{ + {Groups: []FieldRow{{FieldName: "f1", Value: &a}}, Count: 1}, + {Groups: []FieldRow{{FieldName: "f1", Value: &b}}, Count: 1}, + } + + var aa, bb int64 = -1, 1 + expectResult(t, result, QueryResultTypeGroupCounts, RowResult{}, nil, 0, 0, false, []GroupCount{ + {Groups: []FieldRow{{FieldName: "f1", Value: &aa}}, Count: 1}, + {Groups: []FieldRow{{FieldName: "f1", Value: &bb}}, Count: 1}, + }, RowIdentifiersResult{}) +} + +func TestRowIdentifiersResult(t *testing.T) { + result := RowIdentifiersResult{ + IDs: []uint64{1, 2, 3, 4}, + } + expectResult(t, result, QueryResultTypeRowIdentifiers, RowResult{}, nil, 0, 0, false, nil, RowIdentifiersResult{ + IDs: []uint64{1, 2, 3, 4}, + }) +} + +func expectResult(t *testing.T, r QueryResult, resultType uint32, bmp RowResult, + countItems []CountResultItem, sum int64, count int64, changed bool, + groupCounts []GroupCount, rowIdentifiers RowIdentifiersResult) { + if resultType != r.Type() { + log.Fatalf("Result type: %d != %d", resultType, r.Type()) + } + if !reflect.DeepEqual(bmp, r.Row()) { + log.Fatalf("Row: %v != %v", bmp, r.Row()) + } + if !reflect.DeepEqual(countItems, r.CountItems()) { + log.Fatalf("Count items: %v != %v", countItems, r.CountItems()) + } + if count != r.Count() { + log.Fatalf("Count: %d != %d", count, r.Count()) + } + if sum != r.Value() { + log.Fatalf("Sum: %d != %d", sum, r.Value()) + } + if changed != r.Changed() { + log.Fatalf("Changed: %v != %v", changed, r.Changed()) + } + if !reflect.DeepEqual(groupCounts, r.GroupCounts()) { + log.Fatalf("Group counts: %v != %v", groupCounts, r.GroupCounts()) + } + if !reflect.DeepEqual(rowIdentifiers, r.RowIdentifiers()) { + log.Fatalf("Row identifiers: %v != %v", rowIdentifiers, r.RowIdentifiers()) + } +} diff --git a/client/shardnodes.go b/client/shardnodes.go new file mode 100644 index 000000000..836385c44 --- /dev/null +++ b/client/shardnodes.go @@ -0,0 +1,49 @@ +package client + +import ( + "sync" + + pnet "github.com/pilosa/pilosa/v2/net" +) + +type shardNodes struct { + data map[string]map[uint64][]*pnet.URI + mu *sync.RWMutex +} + +func newShardNodes() shardNodes { + return shardNodes{ + data: make(map[string]map[uint64][]*pnet.URI), + mu: &sync.RWMutex{}, + } +} + +func (s shardNodes) Get(index string, shard uint64) ([]*pnet.URI, bool) { + s.mu.RLock() + defer s.mu.RUnlock() + if idx, ok := s.data[index]; ok { + if uris, ok := idx[shard]; ok { + return uris, true + } + } + return nil, false +} + +func (s shardNodes) Put(index string, shard uint64, uris []*pnet.URI) { + s.mu.Lock() + defer s.mu.Unlock() + idx, ok := s.data[index] + if !ok { + idx = make(map[uint64][]*pnet.URI) + } + idx[shard] = uris + s.data[index] = idx +} + +func (s shardNodes) Invalidate() { + s.mu.Lock() + defer s.mu.Unlock() + for k := range s.data { + delete(s.data, k) + } +} diff --git a/client/tracer.go b/client/tracer.go new file mode 100644 index 000000000..660bc0098 --- /dev/null +++ b/client/tracer.go @@ -0,0 +1,104 @@ +// Copyright 2017 Pilosa Corp. +// +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions +// are met: +// +// 1. Redistributions of source code must retain the above copyright +// notice, this list of conditions and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright +// notice, this list of conditions and the following disclaimer in the +// documentation and/or other materials provided with the distribution. +// +// 3. Neither the name of the copyright holder nor the names of its +// contributors may be used to endorse or promote products derived +// from this software without specific prior written permission. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND +// CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, +// INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF +// MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR +// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, +// BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH +// DAMAGE. + +package client + +import ( + opentracing "github.com/opentracing/opentracing-go" + "github.com/opentracing/opentracing-go/log" +) + +type NoopTracer struct{} + +type NoopSpan struct{} + +func (s NoopSpan) Finish() { + // pass +} +func (s NoopSpan) FinishWithOptions(opts opentracing.FinishOptions) { + // pass +} + +func (s NoopSpan) Context() opentracing.SpanContext { + return nil +} +func (s NoopSpan) SetOperationName(operationName string) opentracing.Span { + return s +} + +func (s NoopSpan) SetTag(key string, value interface{}) opentracing.Span { + return s +} + +func (s NoopSpan) LogFields(fields ...log.Field) { + // pass +} + +func (s NoopSpan) LogKV(alternatingKeyValues ...interface{}) { + // pass +} + +func (s NoopSpan) SetBaggageItem(restrictedKey, value string) opentracing.Span { + return s +} + +func (s NoopSpan) BaggageItem(restrictedKey string) string { + return "" +} + +func (s NoopSpan) Tracer() opentracing.Tracer { + return nil +} + +func (s NoopSpan) LogEvent(event string) { + // pass +} + +func (s NoopSpan) LogEventWithPayload(event string, payload interface{}) { + // pass +} + +func (s NoopSpan) Log(data opentracing.LogData) { + // pass +} + +func (t NoopTracer) StartSpan(operationName string, opts ...opentracing.StartSpanOption) opentracing.Span { + return NoopSpan{} +} + +func (t NoopTracer) Inject(sm opentracing.SpanContext, format interface{}, carrier interface{}) error { + return nil +} + +func (t NoopTracer) Extract(format interface{}, carrier interface{}) (opentracing.SpanContext, error) { + return nil, nil +} diff --git a/client/validate.go b/client/validate.go new file mode 100644 index 000000000..7e865c91a --- /dev/null +++ b/client/validate.go @@ -0,0 +1,69 @@ +// Copyright 2017 Pilosa Corp. +// +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions +// are met: +// +// 1. Redistributions of source code must retain the above copyright +// notice, this list of conditions and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright +// notice, this list of conditions and the following disclaimer in the +// documentation and/or other materials provided with the distribution. +// +// 3. Neither the name of the copyright holder nor the names of its +// contributors may be used to endorse or promote products derived +// from this software without specific prior written permission. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND +// CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, +// INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF +// MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR +// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, +// BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH +// DAMAGE. + +package client + +import ( + "regexp" +) + +const ( + maxLabel = 64 + maxKey = 64 +) + +var labelRegex = regexp.MustCompile("^[a-zA-Z][a-zA-Z0-9_-]*$") +var keyRegex = regexp.MustCompile("^[A-Za-z0-9_{}+/=.~%:-]*$") + +// ValidLabel returns true if the given label is valid, otherwise false. +func ValidLabel(label string) bool { + return len(label) <= maxLabel && labelRegex.Match([]byte(label)) +} + +// ValidKey returns true if the given key is valid, otherwise false. +func ValidKey(key string) bool { + return len(key) <= maxKey && keyRegex.Match([]byte(key)) +} + +func validateLabel(label string) error { + if ValidLabel(label) { + return nil + } + return ErrInvalidLabel +} + +func validateKey(key string) error { + if ValidKey(key) { + return nil + } + return ErrInvalidKey +} diff --git a/client/validate_test.go b/client/validate_test.go new file mode 100644 index 000000000..ce07eb7ed --- /dev/null +++ b/client/validate_test.go @@ -0,0 +1,88 @@ +// Copyright 2017 Pilosa Corp. +// +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions +// are met: +// +// 1. Redistributions of source code must retain the above copyright +// notice, this list of conditions and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright +// notice, this list of conditions and the following disclaimer in the +// documentation and/or other materials provided with the distribution. +// +// 3. Neither the name of the copyright holder nor the names of its +// contributors may be used to endorse or promote products derived +// from this software without specific prior written permission. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND +// CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, +// INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF +// MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR +// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, +// BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH +// DAMAGE. + +package client + +import "testing" + +func TestValidateLabel(t *testing.T) { + labels := []string{ + "a", "ab", "ab1", "d_e", "A", "Bc", "B1", "aB", "b-c", + "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + } + for _, label := range labels { + if validateLabel(label) != nil { + t.Fatalf("Should be valid label: %s", label) + } + } +} + +func TestValidateLabelInvalid(t *testing.T) { + labels := []string{ + "", "1", "_", "-", "'", "^", "/", "\\", "*", "a:b", "valid?no", "yüce", + "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa1", + } + for _, label := range labels { + if validateLabel(label) == nil { + t.Fatalf("Should be invalid label: %s", label) + } + } +} + +func TestValidateKey(t *testing.T) { + keys := []string{ + "", "1", "ab", "ab1", "b-c", "d_e", "pilosa.com", + "bbf8d41c-7dba-40c4-94dc-94677b43bcf3", // UUID + "{bbf8d41c-7dba-40c4-94dc-94677b43bcf3}", // Windows GUID + "https%3A//www.pilosa.com/about/%23contact", // escaped URL + "aHR0cHM6Ly93d3cucGlsb3NhLmNvbS9hYm91dC8jY29udGFjdA==", // base64 + "urn:isbn:1234567", + "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + } + for _, key := range keys { + if validateKey(key) != nil { + t.Fatalf("Should be valid key: %s", key) + } + } +} + +func TestValidateKeyInvalid(t *testing.T) { + keys := []string{ + "\"", "'", "slice\\dice", "valid?no", "yüce", "*xyz", "with space", "